From 7bb1be2d20cc05a9d92de2c093786986f799bdb0 Mon Sep 17 00:00:00 2001
From: Christian Speckner <christian.speckner@mayflower.de>
Date: Fri, 31 Jan 2025 17:09:16 +0100
Subject: [PATCH] Make sure that closing one direction closes the other, too.
 (#159)

* Make sure that closing one direction closes the other, too.

* Pacify linter.
---
 go.mod     |  1 -
 go.sum     |  2 --
 http.go    | 26 +++++++++++++-------------
 routine.go | 51 ++++++++++-----------------------------------------
 4 files changed, 23 insertions(+), 57 deletions(-)

diff --git a/go.mod b/go.mod
index 9d1b1dd..ec664d3 100644
--- a/go.mod
+++ b/go.mod
@@ -9,7 +9,6 @@ require (
 	github.com/akamensky/argparse v1.4.0
 	github.com/go-ini/ini v1.67.0
 	github.com/landlock-lsm/go-landlock v0.0.0-20240216195629-efb66220540a
-	github.com/sourcegraph/conc v0.3.0
 	github.com/things-go/go-socks5 v0.0.5
 	golang.org/x/net v0.23.0
 	golang.zx2c4.com/wireguard v0.0.0-20231211153847-12269c276173
diff --git a/go.sum b/go.sum
index 23a91dc..4c70481 100644
--- a/go.sum
+++ b/go.sum
@@ -12,8 +12,6 @@ github.com/landlock-lsm/go-landlock v0.0.0-20240216195629-efb66220540a h1:dz+a1M
 github.com/landlock-lsm/go-landlock v0.0.0-20240216195629-efb66220540a/go.mod h1:1NY/VPO8xm3hXw3f+M65z+PJDLUaZA5cu7OfanxoUzY=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/sourcegraph/conc v0.3.0 h1:OQTbbt6P72L20UqAkXXuLOj79LfEanQ+YQFNpLA9ySo=
-github.com/sourcegraph/conc v0.3.0/go.mod h1:Sdozi7LEKbFPqYX2/J+iBAM6HpqSLTASQIKqDmF7Mt0=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/things-go/go-socks5 v0.0.5 h1:qvKaGcBkfDrUL33SchHN93srAmYGzb4CxSM2DPYufe8=
diff --git a/http.go b/http.go
index ebaa822..88a7ef4 100644
--- a/http.go
+++ b/http.go
@@ -10,8 +10,6 @@ import (
 	"net"
 	"net/http"
 	"strings"
-
-	"github.com/sourcegraph/conc"
 )
 
 const proxyAuthHeaderKey = "Proxy-Authorization"
@@ -32,7 +30,7 @@ func (s *HTTPServer) authenticate(req *http.Request) (int, error) {
 
 	auth := req.Header.Get(proxyAuthHeaderKey)
 	if auth == "" {
-		return http.StatusProxyAuthRequired, fmt.Errorf(http.StatusText(http.StatusProxyAuthRequired))
+		return http.StatusProxyAuthRequired, fmt.Errorf("%s", http.StatusText(http.StatusProxyAuthRequired))
 	}
 
 	enc := strings.TrimPrefix(auth, "Basic ")
@@ -131,17 +129,19 @@ func (s *HTTPServer) serve(conn net.Conn) {
 		log.Println("dial proxy failed: peer nil")
 		return
 	}
+
 	go func() {
-		wg := conc.NewWaitGroup()
-		wg.Go(func() {
-			_, err = io.Copy(conn, peer)
-			_ = conn.Close()
-		})
-		wg.Go(func() {
-			_, err = io.Copy(peer, conn)
-			_ = peer.Close()
-		})
-		wg.Wait()
+		defer conn.Close()
+		defer peer.Close()
+
+		_, _ = io.Copy(conn, peer)
+	}()
+
+	go func() {
+		defer conn.Close()
+		defer peer.Close()
+
+		_, _ = io.Copy(peer, conn)
 	}()
 }
 
diff --git a/routine.go b/routine.go
index 8c65456..edfc793 100644
--- a/routine.go
+++ b/routine.go
@@ -24,7 +24,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/sourcegraph/conc"
 	"github.com/things-go/go-socks5"
 	"github.com/things-go/go-socks5/bufferpool"
 
@@ -190,6 +189,9 @@ func (c CredentialValidator) Valid(username, password string) bool {
 
 // connForward copy data from `from` to `to`
 func connForward(from io.ReadWriteCloser, to io.ReadWriteCloser) {
+	defer from.Close()
+	defer to.Close()
+
 	_, err := io.Copy(to, from)
 	if err != nil {
 		errorLogger.Printf("Cannot forward traffic: %s\n", err.Error())
@@ -212,20 +214,8 @@ func tcpClientForward(vt *VirtualTun, raddr *addressPort, conn net.Conn) {
 		return
 	}
 
-	go func() {
-		wg := conc.NewWaitGroup()
-		wg.Go(func() {
-			connForward(sconn, conn)
-		})
-		wg.Go(func() {
-			connForward(conn, sconn)
-		})
-		wg.Wait()
-		_ = sconn.Close()
-		_ = conn.Close()
-		sconn = nil
-		conn = nil
-	}()
+	go connForward(sconn, conn)
+	go connForward(conn, sconn)
 }
 
 // STDIOTcpForward starts a new connection via wireguard and forward traffic from `conn`
@@ -250,18 +240,8 @@ func STDIOTcpForward(vt *VirtualTun, raddr *addressPort) {
 		return
 	}
 
-	go func() {
-		wg := conc.NewWaitGroup()
-		wg.Go(func() {
-			connForward(os.Stdin, sconn)
-		})
-		wg.Go(func() {
-			connForward(sconn, stdout)
-		})
-		wg.Wait()
-		_ = sconn.Close()
-		sconn = nil
-	}()
+	go connForward(os.Stdin, sconn)
+	go connForward(sconn, stdout)
 }
 
 // SpawnRoutine spawns a local TCP server which acts as a proxy to the specified target
@@ -311,20 +291,9 @@ func tcpServerForward(vt *VirtualTun, raddr *addressPort, conn net.Conn) {
 		return
 	}
 
-	go func() {
-		gr := conc.NewWaitGroup()
-		gr.Go(func() {
-			connForward(sconn, conn)
-		})
-		gr.Go(func() {
-			connForward(conn, sconn)
-		})
-		gr.Wait()
-		_ = sconn.Close()
-		_ = conn.Close()
-		sconn = nil
-		conn = nil
-	}()
+	go connForward(sconn, conn)
+	go connForward(conn, sconn)
+
 }
 
 // SpawnRoutine spawns a TCP server on wireguard which acts as a proxy to the specified target