Merge pull request 'fix: potential resource leak' (#38) from staging into main

Reviewed-on: #38

server/http.go

@@ -298,13 +298,22 @@ func forwardRequest(cw *CustomWriter, initialRequest *RequestHeaderFactory, sshS
 	channel, reqs, err := sshSession.Lifecycle.GetConnection().OpenChannel("forwarded-tcpip", payload)
 	if err != nil {
 		log.Printf("Failed to open forwarded-tcpip channel: %v", err)
+		if closer, ok := cw.writer.(io.Closer); ok {
+			if closeErr := closer.Close(); closeErr != nil {
+				log.Printf("Failed to close connection: %v", closeErr)
+			}
+		}
 		return
 	}
 
 	go func() {
+		defer func() {
+			if r := recover(); r != nil {
+				log.Printf("Panic in request handler goroutine: %v", r)
+			}
+		}()
 		for req := range reqs {
-			err := req.Reply(false, nil)
-			if err != nil {
+			if err := req.Reply(false, nil); err != nil {
 				log.Printf("Failed to reply to request: %v", err)
 				return
 			}

server/middleware.go

@@ -44,119 +44,3 @@ func (ff *ForwardedFor) HandleRequest(header *RequestHeaderFactory) error {
 	header.Set("X-Forwarded-For", host)
 	return nil
 }
-
-//TODO: Implement caching atau enggak
-//const maxCacheSize = 50 * 1024 * 1024
-//
-//type DiskCacheMiddleware struct {
-//	dir       string
-//	mu        sync.Mutex
-//	file      *os.File
-//	path      string
-//	cacheable bool
-//}
-//
-//func NewDiskCacheMiddleware() *DiskCacheMiddleware {
-//	return &DiskCacheMiddleware{dir: "cache"}
-//}
-//
-//func (c *DiskCacheMiddleware) ensureDir() error {
-//	return os.MkdirAll(c.dir, 0755)
-//}
-//
-//func (c *DiskCacheMiddleware) cacheKey(method, path string) string {
-//	return fmt.Sprintf("%s_%s.cache", method, base64.URLEncoding.EncodeToString([]byte(path)))
-//}
-//
-//func (c *DiskCacheMiddleware) filePath(method, path string) string {
-//	return filepath.Join(c.dir, c.cacheKey(method, path))
-//}
-//
-//func fileExists(path string) bool {
-//	_, err := os.Stat(path)
-//	if err == nil {
-//		return true
-//	}
-//	if os.IsNotExist(err) {
-//		return false
-//	}
-//	return false
-//}
-//
-//func canCacheRequest(header *RequestHeaderFactory) bool {
-//	if header.Method != "GET" {
-//		return false
-//	}
-//
-//	if cacheControl := header.Get("Cache-Control"); cacheControl != "" {
-//		if strings.Contains(cacheControl, "no-store") || strings.Contains(cacheControl, "private") || strings.Contains(cacheControl, "no-cache") || strings.Contains(cacheControl, "max-age=0") {
-//			return false
-//		}
-//	}
-//
-//	if header.Get("Authorization") != "" {
-//		return false
-//	}
-//
-//	if header.Get("Cookie") != "" {
-//		return false
-//	}
-//
-//	return true
-//}
-//
-//func (c *DiskCacheMiddleware) HandleRequest(header *RequestHeaderFactory) error {
-//	if !canCacheRequest(header) {
-//		c.cacheable = false
-//		return nil
-//	}
-//
-//	c.cacheable = true
-//	_ = c.ensureDir()
-//	path := c.filePath(header.Method, header.Path)
-//
-//	if fileExists(path + ".finish") {
-//		c.file = nil
-//		return nil
-//	}
-//
-//	if c.file != nil {
-//		err := c.file.Close()
-//		if err != nil {
-//			return err
-//		}
-//		err = os.Rename(c.path, c.path+".finish")
-//		if err != nil {
-//			return err
-//		}
-//	}
-//
-//	c.path = path
-//	f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
-//	if err != nil {
-//		return err
-//	}
-//
-//	c.file = f
-//
-//	return nil
-//}
-//
-//func (c *DiskCacheMiddleware) HandleResponse(header *ResponseHeaderFactory, body []byte) error {
-//	if !c.cacheable {
-//		return nil
-//	}
-//
-//	if c.file == nil {
-//		header.Set("X-Cache", "HIT")
-//		return nil
-//	}
-//
-//	_, err := c.file.Write(body)
-//	if err != nil {
-//		return err
-//	}
-//
-//	header.Set("X-Cache", "MISS")
-//	return nil
-//}

session/forwarder/forwarder.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"encoding/binary"
 	"errors"
-	"fmt"
 	"io"
 	"log"
 	"net"
@@ -21,7 +20,6 @@ type Forwarder struct {
 	ForwardedPort uint16
 	SlugManager   slug.Manager
 	Lifecycle     Lifecycle
-	ActiveForwarder []chan struct{}
 }
 
 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
-		}
-	}()
-
-	f.AddActiveForwarder(drop)
-
 		_, err := io.Copy(dst, src)
-
-	if err != nil && !errors.Is(err, io.EOF) {
+		if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) {
 			log.Printf("Error copying from channel to conn.Writer: %v", err)
 		}
-	return
+		done <- struct{}{}
+	}()
+
+	<-done
 }
 
 func (f *Forwarder) SetType(tunnelType types.TunnelType) {

session/handler.go

@@ -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)

session/interaction/interaction.go

@@ -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
+		}
 	}
 }
 

types/types.go

@@ -19,7 +19,6 @@ type InteractionType string
 
 const (
 	Slug InteractionType = "SLUG"
-	Drop InteractionType = "DROP"
 )
 
 var BadGatewayResponse = []byte("HTTP/1.1 502 Bad Gateway\r\n" +