fix: potential resource leak
All checks were successful
Docker Build and Push / build-and-push (push) Successful in 4m17s

This commit is contained in:
2025-12-17 21:38:00 +07:00
parent 2725975d82
commit 6451304ed7
7 changed files with 63 additions and 237 deletions

View File

@ -4,7 +4,6 @@ import (
"bytes"
"encoding/binary"
"errors"
"fmt"
"io"
"log"
"net"
@ -16,12 +15,11 @@ import (
)
type Forwarder struct {
Listener net.Listener
TunnelType types.TunnelType
ForwardedPort uint16
SlugManager slug.Manager
Lifecycle Lifecycle
ActiveForwarder []chan struct{}
Listener net.Listener
TunnelType types.TunnelType
ForwardedPort uint16
SlugManager slug.Manager
Lifecycle Lifecycle
}
type Lifecycle interface {
@ -41,27 +39,6 @@ type ForwardingController interface {
SetLifecycle(lifecycle Lifecycle)
CreateForwardedTCPIPPayload(origin net.Addr) []byte
WriteBadGatewayResponse(dst io.Writer)
AddActiveForwarder(drop chan struct{})
DropAllForwarder() int
GetForwarderCount() int
}
func (f *Forwarder) AddActiveForwarder(drop chan struct{}) {
f.ActiveForwarder = append(f.ActiveForwarder, drop)
}
func (f *Forwarder) DropAllForwarder() int {
total := 0
for _, d := range f.ActiveForwarder {
close(d)
total += 1
}
f.ActiveForwarder = nil
return total
}
func (f *Forwarder) GetForwarderCount() int {
return len(f.ActiveForwarder)
}
func (f *Forwarder) SetLifecycle(lifecycle Lifecycle) {
@ -82,7 +59,10 @@ func (f *Forwarder) AcceptTCPConnections() {
channel, reqs, err := f.Lifecycle.GetConnection().OpenChannel("forwarded-tcpip", payload)
if err != nil {
log.Printf("Failed to open forwarded-tcpip channel: %v", err)
return
if closeErr := conn.Close(); closeErr != nil {
log.Printf("Failed to close connection: %v", closeErr)
}
continue
}
go func() {
@ -99,8 +79,7 @@ func (f *Forwarder) AcceptTCPConnections() {
}
func (f *Forwarder) HandleConnection(dst io.ReadWriter, src ssh.Channel, remoteAddr net.Addr) {
drop := make(chan struct{})
defer func(src ssh.Channel) {
defer func() {
_, err := io.Copy(io.Discard, src)
if err != nil {
log.Printf("Failed to discard connection: %v", err)
@ -108,34 +87,38 @@ func (f *Forwarder) HandleConnection(dst io.ReadWriter, src ssh.Channel, remoteA
err = src.Close()
if err != nil && !errors.Is(err, io.EOF) {
log.Printf("Error closing connection: %v", err)
log.Printf("Error closing source channel: %v", err)
}
}(src)
if closer, ok := dst.(io.Closer); ok {
err = closer.Close()
if err != nil && !errors.Is(err, io.EOF) {
log.Printf("Error closing destination connection: %v", err)
}
}
}()
log.Printf("Handling new forwarded connection from %s", remoteAddr)
done := make(chan struct{}, 2)
go func() {
_, err := io.Copy(src, dst)
if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) {
log.Printf("Error copying from conn.Reader to channel: %v", err)
}
done <- struct{}{}
}()
go func() {
select {
case <-drop:
fmt.Println("Closinggggg")
return
_, err := io.Copy(dst, src)
if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) {
log.Printf("Error copying from channel to conn.Writer: %v", err)
}
done <- struct{}{}
}()
f.AddActiveForwarder(drop)
_, err := io.Copy(dst, src)
if err != nil && !errors.Is(err, io.EOF) {
log.Printf("Error copying from channel to conn.Writer: %v", err)
}
return
<-done
}
func (f *Forwarder) SetType(tunnelType types.TunnelType) {

View File

@ -211,6 +211,9 @@ func (s *SSHSession) HandleTCPForward(req *ssh.Request, addr string, portToBind
listener, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", portToBind))
if err != nil {
s.Interaction.SendMessage(fmt.Sprintf("Port %d is already in use or restricted. Please choose a different port.\r\n", portToBind))
if setErr := portUtil.Manager.SetPortStatus(portToBind, false); setErr != nil {
log.Printf("Failed to reset port status: %v", setErr)
}
err = req.Reply(false, nil)
if err != nil {
log.Println("Failed to reply to request:", err)
@ -227,6 +230,9 @@ func (s *SSHSession) HandleTCPForward(req *ssh.Request, addr string, portToBind
err = binary.Write(buf, binary.BigEndian, uint32(portToBind))
if err != nil {
log.Println("Failed to write port to buffer:", err)
if setErr := portUtil.Manager.SetPortStatus(portToBind, false); setErr != nil {
log.Printf("Failed to reset port status: %v", setErr)
}
err = listener.Close()
if err != nil {
log.Printf("Failed to close listener: %s", err)
@ -239,6 +245,9 @@ func (s *SSHSession) HandleTCPForward(req *ssh.Request, addr string, portToBind
err = req.Reply(true, buf.Bytes())
if err != nil {
log.Println("Failed to reply to request:", err)
if setErr := portUtil.Manager.SetPortStatus(portToBind, false); setErr != nil {
log.Printf("Failed to reset port status: %v", setErr)
}
err = listener.Close()
if err != nil {
log.Printf("Failed to close listener: %s", err)

View File

@ -39,8 +39,6 @@ type Forwarder interface {
Close() error
GetTunnelType() types.TunnelType
GetForwardedPort() uint16
DropAllForwarder() int
GetForwarderCount() int
}
type Interaction struct {
@ -118,8 +116,6 @@ func (i *Interaction) handleInteractiveMode(char byte) {
switch i.InteractionType {
case types.Slug:
i.HandleSlugEditMode(char)
case types.Drop:
i.HandleDropMode(char)
}
}
@ -271,8 +267,17 @@ func (i *Interaction) returnToMainScreen() {
}
func (i *Interaction) HandleSlugCancel() {
i.SendMessage(clearScreen)
i.SendMessage("\r\n\r\n⚠ SUBDOMAIN EDIT CANCELLED ⚠️\r\n\r\n")
i.SendMessage("Press any key to continue...\r\n")
i.InteractiveMode = false
i.showMessageAndWait("\r\n\r\n⚠ SUBDOMAIN EDIT CANCELLED ⚠️\r\n\r\n")
i.InteractionType = ""
i.WaitForKeyPress()
i.SendMessage(clearScreen)
i.ShowWelcomeMessage()
i.ShowForwardingMessage()
}
func (i *Interaction) HandleSlugUpdateError() {
@ -295,7 +300,6 @@ func (i *Interaction) HandleCommand(command string) {
"/help": i.handleHelpCommand,
"/clear": i.handleClearCommand,
"/slug": i.handleSlugCommand,
"/drop": i.handleDropCommand,
}
if handler, exists := handlers[command]; exists {
@ -315,7 +319,7 @@ func (i *Interaction) handleByeCommand() {
}
func (i *Interaction) handleHelpCommand() {
i.SendMessage("\r\nAvailable commands: /bye, /help, /clear, /slug, /drop\r\n")
i.SendMessage("\r\nAvailable commands: /bye, /help, /clear, /slug\r\n")
}
func (i *Interaction) handleClearCommand() {
@ -340,13 +344,6 @@ func (i *Interaction) handleSlugCommand() {
i.SendMessage("➤ " + i.EditSlug + "." + domain)
}
func (i *Interaction) handleDropCommand() {
i.InteractiveMode = true
i.InteractionType = types.Drop
i.SendMessage(clearScreen)
i.ShowDropMessage()
}
func (i *Interaction) ShowForwardingMessage() {
domain := utils.Getenv("domain")
@ -361,64 +358,6 @@ func (i *Interaction) ShowForwardingMessage() {
}
}
func (i *Interaction) HandleDropMode(char byte) {
switch {
case char == enterChar || char == 'y' || char == 'Y':
i.executeDropAll()
case char == escapeChar || char == 'n' || char == 'N' || char == ctrlC:
i.cancelDrop()
}
}
func (i *Interaction) executeDropAll() {
count := i.Forwarder.DropAllForwarder()
message := fmt.Sprintf("Dropped %d forwarders\r\n", count)
i.showMessageAndWait(message)
}
func (i *Interaction) cancelDrop() {
i.showMessageAndWait("Dropping canceled.\r\n")
}
func (i *Interaction) showMessageAndWait(message string) {
i.SendMessage(clearScreen)
i.SendMessage(message)
i.SendMessage("Press any key to continue...\r\n")
i.InteractiveMode = false
i.InteractionType = ""
i.WaitForKeyPress()
i.SendMessage(clearScreen)
i.ShowWelcomeMessage()
i.ShowForwardingMessage()
}
func (i *Interaction) ShowDropMessage() {
confirmText := fmt.Sprintf(" ║ Drop ALL %d active connections?", i.Forwarder.GetForwarderCount())
boxWidth := calculateBoxWidth(confirmText)
box := buildDropConfirmationBox(boxWidth, confirmText)
i.SendMessage("\r\n" + box + "\r\n\r\n")
}
func buildDropConfirmationBox(boxWidth int, confirmText string) string {
topBorder := " ╔" + strings.Repeat("═", boxWidth-4) + "╗\r\n"
title := centerText("DROP CONFIRMATION", boxWidth-4)
header := " ║" + title + "║\r\n"
midBorder := " ╠" + strings.Repeat("═", boxWidth-4) + "╣\r\n"
emptyLine := " ║" + strings.Repeat(" ", boxWidth-4) + "║\r\n"
confirmLine := confirmText + strings.Repeat(" ", boxWidth-len(confirmText)+1) + "║\r\n"
controlText := " ║ [Enter/Y] Confirm [N/Esc] Cancel"
controlLine := controlText + strings.Repeat(" ", boxWidth-len(controlText)+1) + "║\r\n"
bottomBorder := " ╚" + strings.Repeat("═", boxWidth-4) + "╝\r\n"
return topBorder + header + midBorder + emptyLine + confirmLine + emptyLine + controlLine + emptyLine + bottomBorder
}
func (i *Interaction) ShowWelcomeMessage() {
asciiArt := []string{
` _______ _ _____ _ `,
@ -436,7 +375,6 @@ func (i *Interaction) ShowWelcomeMessage() {
` - '/help' : Show this help message`,
` - '/clear' : Clear the current line`,
` - '/slug' : Set custom subdomain`,
` - '/drop' : Drop all active forwarders`,
}
for _, line := range asciiArt {
@ -483,6 +421,10 @@ func (i *Interaction) WaitForKeyPress() {
if err == nil {
break
}
if err != nil {
log.Printf("Error reading keypress: %v", err)
break
}
}
}