From 759780508a3179e87b0cde5be3302f5569d03c85 Mon Sep 17 00:00:00 2001 From: Owen Schwartz Date: Sun, 19 Jan 2025 22:46:00 -0500 Subject: [PATCH] Resolve TCP hanging but port is in use issue --- go.mod | 1 + go.sum | 2 ++ proxy/manager.go | 87 +++++++++++++++++++++++------------------------- proxy/types.go | 21 ++++++------ 4 files changed, 56 insertions(+), 55 deletions(-) diff --git a/go.mod b/go.mod index 66a9c2a..2cc0c19 100644 --- a/go.mod +++ b/go.mod @@ -10,6 +10,7 @@ require ( github.com/google/btree v1.1.2 // indirect github.com/gorilla/websocket v1.5.3 // indirect golang.org/x/crypto v0.28.0 // indirect + golang.org/x/exp v0.0.0-20250106191152-7588d65b2ba8 // indirect golang.org/x/net v0.30.0 // indirect golang.org/x/sys v0.26.0 // indirect golang.org/x/time v0.7.0 // indirect diff --git a/go.sum b/go.sum index 571b2d6..d95ab3a 100644 --- a/go.sum +++ b/go.sum @@ -4,6 +4,8 @@ github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aN github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= golang.org/x/crypto v0.28.0 h1:GBDwsMXVQi34v5CCYUm2jkJvu4cbtru2U4TN2PSyQnw= golang.org/x/crypto v0.28.0/go.mod h1:rmgy+3RHxRZMyY0jjAJShp2zgEdOqj2AO7U0pYmeQ7U= +golang.org/x/exp v0.0.0-20250106191152-7588d65b2ba8 h1:yqrTHse8TCMW1M1ZCP+VAR/l0kKxwaAIqN/il7x4voA= +golang.org/x/exp v0.0.0-20250106191152-7588d65b2ba8/go.mod h1:tujkw807nyEEAamNbDrEGzRav+ilXA7PCRAd6xsmwiU= golang.org/x/net v0.30.0 h1:AcW1SDZMkb8IpzCdQUaIq2sP4sZ4zw+55h6ynffypl4= golang.org/x/net v0.30.0/go.mod h1:2wGyMJ5iFasEhkwi13ChkO/t1ECNC4X4eBKkVFyYFlU= golang.org/x/sys v0.26.0 h1:KHjCJyddX0LoSTb3J+vWpupP9p0oznkqVk/IfjymZbo= diff --git a/proxy/manager.go b/proxy/manager.go index 6aa3128..ae89b3e 100644 --- a/proxy/manager.go +++ b/proxy/manager.go @@ -9,6 +9,7 @@ import ( "time" "github.com/fosrl/newt/logger" + "gvisor.dev/gvisor/pkg/tcpip/adapters/gonet" "golang.zx2c4.com/wireguard/tun/netstack" ) @@ -19,13 +20,12 @@ func NewProxyManager(tnet *netstack.Net) *ProxyManager { } } -func (pm *ProxyManager) AddTarget(protocol, listen string, port int, target string) { +func (pm *ProxyManager) AddTarget(protocol, listen string, port int, target string) error { pm.Lock() defer pm.Unlock() logger.Info("Adding target: %s://%s:%d -> %s", protocol, listen, port, target) - - newTarget := ProxyTarget{ + newTarget := &ProxyTarget{ Protocol: protocol, Listen: listen, Port: port, @@ -35,6 +35,7 @@ func (pm *ProxyManager) AddTarget(protocol, listen string, port int, target stri } pm.targets = append(pm.targets, newTarget) + return nil } func (pm *ProxyManager) RemoveTarget(protocol, listen string, port int) error { @@ -54,31 +55,21 @@ func (pm *ProxyManager) RemoveTarget(protocol, listen string, port int) error { // Signal the serving goroutine to stop select { case <-target.cancel: - // Channel is already closed, no need to close it again + // Channel is already closed default: close(target.cancel) } - // Close the appropriate listener/connection based on protocol + // Close the listener/connection target.Lock() switch protocol { case "tcp": if target.listener != nil { - select { - case <-target.cancel: - // Listener was already closed by Stop() - default: - target.listener.Close() - } + target.listener.Close() } case "udp": if target.udpConn != nil { - select { - case <-target.cancel: - // Connection was already closed by Stop() - default: - target.udpConn.Close() - } + target.udpConn.Close() } } target.Unlock() @@ -86,7 +77,6 @@ func (pm *ProxyManager) RemoveTarget(protocol, listen string, port int) error { // Wait for the target to fully stop <-target.done - // Remove the target from the slice pm.targets = append(pm.targets[:i], pm.targets[i+1:]...) return nil } @@ -99,9 +89,7 @@ func (pm *ProxyManager) Start() error { pm.RLock() defer pm.RUnlock() - for i := range pm.targets { - target := &pm.targets[i] - + for _, target := range pm.targets { target.Lock() // If target is already running, skip it if target.listener != nil || target.udpConn != nil { @@ -110,7 +98,6 @@ func (pm *ProxyManager) Start() error { } // Mark the target as starting by creating a nil listener/connection - // This prevents other goroutines from trying to start it if strings.ToLower(target.Protocol) == "tcp" { target.listener = nil } else { @@ -135,10 +122,11 @@ func (pm *ProxyManager) Stop() error { defer pm.Unlock() var wg sync.WaitGroup - for i := range pm.targets { - target := &pm.targets[i] + for _, target := range pm.targets { wg.Add(1) - go func(t *ProxyTarget) { + // Create a new variable in the loop to avoid closure issues + t := target // Take a local copy + go func() { defer wg.Done() close(t.cancel) t.Lock() @@ -151,7 +139,7 @@ func (pm *ProxyManager) Stop() error { t.Unlock() // Wait for the target to fully stop <-t.done - }(target) + }() } wg.Wait() return nil @@ -220,32 +208,41 @@ func (pm *ProxyManager) handleTCPConnection(clientConn net.Conn, target string, } defer serverConn.Close() - var wg sync.WaitGroup - wg.Add(2) + // Create error channels for both copy operations + errc1 := make(chan error, 1) + errc2 := make(chan error, 1) - // Client -> Server + // Copy from client to server go func() { - defer wg.Done() - select { - case <-done: - return - default: - io.Copy(serverConn, clientConn) - } + _, err := io.Copy(serverConn, clientConn) + errc1 <- err }() - // Server -> Client + // Copy from server to client go func() { - defer wg.Done() - select { - case <-done: - return - default: - io.Copy(clientConn, serverConn) - } + _, err := io.Copy(clientConn, serverConn) + errc2 <- err }() - wg.Wait() + // Wait for either copy to finish or done signal + select { + case <-done: + // Gracefully close connections without type assertions + if closer, ok := clientConn.(interface{ CloseRead() error }); ok { + closer.CloseRead() + } + if closer, ok := serverConn.(*gonet.TCPConn); ok { + closer.CloseRead() + } + case err := <-errc1: + if err != nil { + logger.Info("Error copying client->server: %v", err) + } + case err := <-errc2: + if err != nil { + logger.Info("Error copying server->client: %v", err) + } + } } func (pm *ProxyManager) serveUDP(target *ProxyTarget) { diff --git a/proxy/types.go b/proxy/types.go index 10bd046..f189596 100644 --- a/proxy/types.go +++ b/proxy/types.go @@ -9,19 +9,20 @@ import ( ) type ProxyTarget struct { - Protocol string - Listen string - Port int - Target string - cancel chan struct{} // Channel to signal shutdown - done chan struct{} // Channel to signal completion - listener net.Listener // For TCP - udpConn net.PacketConn // For UDP - sync.Mutex // Protect access to connection + Protocol string + Listen string + Port int + Target string + cancel chan struct{} // Channel to signal shutdown + done chan struct{} // Channel to signal completion + listener net.Listener // For TCP + udpConn net.PacketConn // For UDP + sync.Mutex // Protect access to connection + activeConns sync.Map } type ProxyManager struct { - targets []ProxyTarget + targets []*ProxyTarget tnet *netstack.Net log *log.Logger sync.RWMutex // Protect access to targets slice