wiresocks: add deadline to connection proxying

Signed-off-by: Mark Pashmfouroush <mark@markpash.me>

Author: markpash

go.mod

@@ -6,6 +6,8 @@ replace github.com/Psiphon-Labs/psiphon-tunnel-core => github.com/bepass-org/psi
 
 require (
 	github.com/Psiphon-Labs/psiphon-tunnel-core v2.0.28+incompatible
+	github.com/adrg/xdg v0.4.0
+	github.com/carlmjohnson/versioninfo v0.22.5
 	github.com/fatih/color v1.16.0
 	github.com/flynn/noise v1.1.0
 	github.com/frankban/quicktest v1.14.6
@@ -16,6 +18,7 @@ require (
 	github.com/quic-go/quic-go v0.40.1
 	github.com/refraction-networking/utls v1.3.3
 	github.com/rodaine/table v1.1.1
+	github.com/things-go/go-socks5 v0.0.5
 	golang.org/x/crypto v0.21.0
 	golang.org/x/net v0.22.0
 	golang.org/x/sys v0.18.0
@@ -31,11 +34,9 @@ require (
 	github.com/Psiphon-Labs/goptlib v0.0.0-20200406165125-c0e32a7a3464 // indirect
 	github.com/Psiphon-Labs/psiphon-tls v0.0.0-20240305020009-09f917290799 // indirect
 	github.com/Psiphon-Labs/quic-go v0.0.0-20240305203241-7c4a760d03cc // indirect
-	github.com/adrg/xdg v0.4.0 // indirect
 	github.com/andybalholm/brotli v1.0.5 // indirect
 	github.com/armon/go-proxyproto v0.0.0-20180202201750-5b7edb60ff5f // indirect
 	github.com/bifurcation/mint v0.0.0-20180306135233-198357931e61 // indirect
-	github.com/carlmjohnson/versioninfo v0.22.5 // indirect
 	github.com/cheekybits/genny v0.0.0-20170328200008-9127e812e1e9 // indirect
 	github.com/cognusion/go-cache-lru v0.0.0-20170419142635-f73e2280ecea // indirect
 	github.com/dchest/siphash v1.2.3 // indirect

go.sum

@@ -204,6 +204,8 @@ github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcU
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/syndtr/gocapability v0.0.0-20200815063812-42c35b437635 h1:kdXcSzyDtseVEc4yCz2qF8ZrQvIDBJLl4S1c3GCXmoI=
 github.com/syndtr/gocapability v0.0.0-20200815063812-42c35b437635/go.mod h1:hkRG7XYTFWNJGYcbNJQlaLq0fg1yr4J4t/NcTQtrfww=
+github.com/things-go/go-socks5 v0.0.5 h1:qvKaGcBkfDrUL33SchHN93srAmYGzb4CxSM2DPYufe8=
+github.com/things-go/go-socks5 v0.0.5/go.mod h1:mtzInf8v5xmsBpHZVbIw2YQYhc4K0jRwzfsH64Uh0IQ=
 github.com/wader/filtertransport v0.0.0-20200316221534-bdd9e61eee78 h1:9sreu9e9KOihf2Y0NbpyfWhd1XFDcL4GTkPYL4IvMrg=
 github.com/wader/filtertransport v0.0.0-20200316221534-bdd9e61eee78/go.mod h1:HazXTRLhXFyq80TQp7PUXi6BKE6mS+ydEdzEqNBKopQ=
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=

wiresocks/proxy.go

@@ -2,23 +2,27 @@ package wiresocks
 
 import (
 	"context"
+	"errors"
 	"io"
 	"log/slog"
 	"net"
 	"net/netip"
+	"time"
 
 	"github.com/bepass-org/warp-plus/proxy/pkg/mixed"
 	"github.com/bepass-org/warp-plus/proxy/pkg/statute"
 	"github.com/bepass-org/warp-plus/wireguard/device"
 	"github.com/bepass-org/warp-plus/wireguard/tun/netstack"
+	"github.com/things-go/go-socks5/bufferpool"
 )
 
 // VirtualTun stores a reference to netstack network and DNS configuration
 type VirtualTun struct {
-	Tnet      *netstack.Net
-	Logger    *slog.Logger
-	Dev       *device.Device
-	Ctx       context.Context
+	Tnet   *netstack.Net
+	Logger *slog.Logger
+	Dev    *device.Device
+	Ctx    context.Context
+	pool   bufferpool.BufPool
 }
 
 // StartProxy spawns a socks5 server.
@@ -60,12 +64,18 @@ func (vt *VirtualTun) generalHandler(req *statute.ProxyRequest) error {
 	done := make(chan error, 1)
 	// Copy data from req.Conn to conn
 	go func() {
-		_, err := io.Copy(conn, req.Conn)
+		req.Conn.SetReadDeadline(time.Now().Add(15 * time.Second))
+		buf1 := vt.pool.Get()
+		defer vt.pool.Put(buf1)
+		_, err := copyConnTimeout(conn, req.Conn, buf1[:cap(buf1)], 15*time.Second)
 		done <- err
 	}()
 	// Copy data from conn to req.Conn
 	go func() {
-		_, err := io.Copy(req.Conn, conn)
+		conn.SetReadDeadline(time.Now().Add(15 * time.Second))
+		buf2 := vt.pool.Get()
+		defer vt.pool.Put(buf2)
+		_, err := copyConnTimeout(req.Conn, conn, buf2[:cap(buf2)], 15*time.Second)
 		done <- err
 	}()
 	// Wait for one of the copy operations to finish
@@ -75,10 +85,7 @@ func (vt *VirtualTun) generalHandler(req *statute.ProxyRequest) error {
 	}
 
 	// Close connections and wait for the other copy operation to finish
-	conn.Close()
-	req.Conn.Close()
 	<-done
-
 	return nil
 }
 
@@ -89,3 +96,44 @@ func (vt *VirtualTun) Stop() {
 		}
 	}
 }
+
+var errInvalidWrite = errors.New("invalid write result")
+
+func copyConnTimeout(dst net.Conn, src net.Conn, buf []byte, timeout time.Duration) (written int64, err error) {
+	if buf != nil && len(buf) == 0 {
+		panic("empty buffer in CopyBuffer")
+	}
+
+	for {
+		if err := src.SetReadDeadline(time.Now().Add(timeout)); err != nil {
+			return 0, err
+		}
+
+		nr, er := src.Read(buf)
+		if nr > 0 {
+			nw, ew := dst.Write(buf[0:nr])
+			if nw < 0 || nr < nw {
+				nw = 0
+				if ew == nil {
+					ew = errInvalidWrite
+				}
+			}
+			written += int64(nw)
+			if ew != nil {
+				err = ew
+				break
+			}
+			if nr != nw {
+				err = io.ErrShortWrite
+				break
+			}
+		}
+		if er != nil {
+			if er != io.EOF {
+				err = er
+			}
+			break
+		}
+	}
+	return written, err
+}

wiresocks/wiresocks.go

@@ -9,6 +9,7 @@ import (
 	"github.com/bepass-org/warp-plus/wireguard/conn"
 	"github.com/bepass-org/warp-plus/wireguard/device"
 	"github.com/bepass-org/warp-plus/wireguard/tun/netstack"
+	"github.com/things-go/go-socks5/bufferpool"
 )
 
 // StartWireguard creates a tun interface on netstack given a configuration
@@ -46,9 +47,10 @@ func StartWireguard(ctx context.Context, l *slog.Logger, conf *Configuration) (*
 	}
 
 	return &VirtualTun{
-		Tnet:      tnet,
-		Logger:    l.With("subsystem", "vtun"),
-		Dev:       dev,
-		Ctx:       ctx,
+		Tnet:   tnet,
+		Logger: l.With("subsystem", "vtun"),
+		Dev:    dev,
+		Ctx:    ctx,
+		pool:   bufferpool.NewPool(256 * 1024),
 	}, nil
 }