switch socks library for better context and UDP support

Author: lanrat

cmd/main.go

@@ -102,10 +102,12 @@ func main() {
 	totalNetworks := 1 << totalNetworksBits
 	l.Printf("Running with subnet size /%d and /%d prefix resulting in %d egress networks and %d options per network", *subnetBits, cidrBits, totalNetworks, hostsPerNetwork)
 
+	ctx := context.Background()
+
 	// test mode
 	if *runTest {
 		// test requests
-		err := test(context.Background(), parsedNetwork, *subnetBits)
+		err := test(ctx, parsedNetwork, *subnetBits)
 		if err != nil {
 			l.Fatal(err)
 		}
@@ -121,7 +123,7 @@ func main() {
 
 	// run subnet proxy server
 	l.Printf("Starting subnet egress proxy %s\n", *listenAddr)
-	err = runRandomSubnetProxy(*listenAddr, parsedNetwork, *subnetBits)
+	err = runRandomSubnetProxy(ctx, *listenAddr, parsedNetwork, *subnetBits)
 	if err != nil {
 		l.Fatal(err)
 	}

cmd/socks.go

@@ -6,9 +6,8 @@ import (
 	"net/netip"
 	"strconv"
 
-	"github.com/haxii/socks5"
+	"github.com/lanrat/go-socks5"
 	"github.com/lanrat/stargate"
-	"golang.org/x/sync/errgroup"
 )
 
 // runRandomSubnetProxy starts a SOCKS5 proxy server listening on listenAddr that distributes
@@ -17,7 +16,7 @@ import (
 // This is memory efficient for large IPv6 ranges as it doesn't pre-generate all addresses.
 // The function cycles through all available subnets before repeating.
 // Supports both TCP and UDP protocols simultaneously.
-func runRandomSubnetProxy(listenAddr string, parsedNetwork netip.Prefix, cidrSize uint) error {
+func runRandomSubnetProxy(ctx context.Context, listenAddr string, parsedNetwork netip.Prefix, cidrSize uint) error {
 	ipItr, err := stargate.NewRandomIPIterator(parsedNetwork, cidrSize)
 	if err != nil {
 		return err
@@ -34,10 +33,15 @@ func runRandomSubnetProxy(listenAddr string, parsedNetwork netip.Prefix, cidrSiz
 		return err
 	}
 
+	udpDial := func(ctx context.Context, network string, udpClientSrcAddr, targetUDPAddr *net.UDPAddr) (net.Conn, error) {
+		return ipItr.Dial(ctx, network, targetUDPAddr.String())
+	}
+
 	conf := &socks5.Config{
 		Logger:   l,
 		Resolver: NewDNSResolver(getCIDRNetwork(parsedNetwork)),
 		Dial:     ipItr.Dial,
+		DialUDP:  udpDial,
 		BindIP:   net.ParseIP(host),
 		BindPort: port,
 	}
@@ -46,23 +50,9 @@ func runRandomSubnetProxy(listenAddr string, parsedNetwork netip.Prefix, cidrSiz
 		return err
 	}
 
-	// Use errgroup to manage both TCP and UDP listeners
-	var g errgroup.Group
-
-	// Start TCP listener
-	g.Go(func() error {
-		l.Printf("Starting TCP SOCKS5 proxy on %s", listenAddr)
-		return server.ListenAndServe("tcp", listenAddr)
-	})
-
-	// Start UDP listener
-	g.Go(func() error {
-		l.Printf("Starting UDP SOCKS5 proxy on %s", listenAddr)
-		return server.ListenAndServe("udp", listenAddr)
-	})
-
-	// Wait for both listeners, return first error
-	return g.Wait()
+	// Start TCP SOCKS5 proxy (UDP is handled internally by the server when BindPort is set)
+	l.Printf("Starting SOCKS5 proxy on %s", listenAddr)
+	return server.ListenAndServe(ctx, listenAddr)
 }
 
 // DNSResolver implements socks5.NameResolver using the system DNS resolver.

go.mod

@@ -1,7 +1,7 @@
 module github.com/lanrat/stargate
 
 require (
-	github.com/haxii/socks5 v1.0.0
+	github.com/lanrat/go-socks5 v0.1.1
 	golang.org/x/sync v0.16.0
 )
 

go.sum

@@ -1,4 +1,4 @@
-github.com/haxii/socks5 v1.0.0 h1:78BIzd4lHibdRNOKdMwKCnnsgYLW9SeotqU+nMhWSSo=
-github.com/haxii/socks5 v1.0.0/go.mod h1:6O9Ba2yrLlvuSe/L1e84eZI8cPw6H+q1Ilr4hjgm4uY=
+github.com/lanrat/go-socks5 v0.1.1 h1:e9w0XT/p6Gts5EutFXypULpmM3Ul1mz1Vg3Q4mjkN+Q=
+github.com/lanrat/go-socks5 v0.1.1/go.mod h1:EOAqUXY0ToNGWAW4/lZ/jwlrCISCW3tSIZ7kcO579KI=
 golang.org/x/sync v0.16.0 h1:ycBJEhp9p4vXvUZNszeOq0kGTPghopOL8q0fq3vstxw=
 golang.org/x/sync v0.16.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=