From 740d55fde481c0ce6c64b2de5c16c8180f9caf4c Mon Sep 17 00:00:00 2001
From: eyedeekay <hankhill19580@gmail.com>
Date: Fri, 7 Feb 2025 12:38:10 -0500
Subject: [PATCH] Fix UDP handler in SOCKS

---
 lib/socks/client/handle.go | 30 ++++++++++++++++++++++--------
 lib/socks/client/socks.go  |  1 +
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/lib/socks/client/handle.go b/lib/socks/client/handle.go
index 6dd2dbc..e504c52 100644
--- a/lib/socks/client/handle.go
+++ b/lib/socks/client/handle.go
@@ -1,15 +1,19 @@
 package socks
 
 import (
+	"context"
+	"fmt"
 	"net"
 
 	"github.com/go-i2p/go-forward/config"
 	"github.com/go-i2p/go-forward/packet"
 	"github.com/go-i2p/go-forward/stream"
+	"github.com/go-i2p/i2pkeys"
 	"github.com/txthinking/socks5"
 )
 
 var socksHandler socks5.Handler = &SOCKS{}
+var forwardConfig = config.DefaultConfig()
 
 // TCPHandle implements socks5.Handler.
 func (s *SOCKS) TCPHandle(_ *socks5.Server, conn *net.TCPConn, req *socks5.Request) error {
@@ -20,28 +24,38 @@ func (s *SOCKS) TCPHandle(_ *socks5.Server, conn *net.TCPConn, req *socks5.Reque
 	}
 	defer i2pConn.Close()
 
-	// Forward data between TCP and I2P connections
-	err = stream.Forward(nil, conn, i2pConn, config.DefaultConfig())
+	ctx := context.Background()
+	err = stream.Forward(ctx, conn, i2pConn, forwardConfig)
 	return err
 }
 
 // UDPHandle implements socks5.Handler.
-func (s *SOCKS) UDPHandle(_ *socks5.Server, addr *net.UDPAddr, d *socks5.Datagram) error {
+func (s *SOCKS) UDPHandle(_ *socks5.Server, addr *net.UDPAddr, data *socks5.Datagram) error {
 	// Connect to destination through I2P
-	i2pConn, err := s.Garlic.DialRemote("udp", d.Address())
+	i2pConn, err := s.Garlic.DialRemote("udp", data.Address())
 	if err != nil {
 		return err
 	}
 	defer i2pConn.Close()
 
-	// Create UDP connection back to client
+	remoteAddr, err := i2pkeys.Lookup(data.Address())
+	if err != nil {
+		return fmt.Errorf("failed to lookup remote address: %w", err)
+	}
+
+	// Send initial datagram
+	if _, err := i2pConn.WriteTo(data.Data, remoteAddr); err != nil {
+		return fmt.Errorf("failed to write initial datagram: %w", err)
+	}
+
+	// Set up UDP connection for replies
 	conn, err := net.DialUDP("udp", nil, addr)
 	if err != nil {
 		return err
 	}
 	defer conn.Close()
 
-	// Forward packets between UDP and I2P connections
-	err = packet.Forward(nil, conn, i2pConn, config.DefaultConfig())
-	return err
+	// Forward subsequent packets
+	ctx := context.Background()
+	return packet.Forward(ctx, conn, i2pConn, forwardConfig)
 }
diff --git a/lib/socks/client/socks.go b/lib/socks/client/socks.go
index f6c8473..623561c 100644
--- a/lib/socks/client/socks.go
+++ b/lib/socks/client/socks.go
@@ -130,6 +130,7 @@ func (s *SOCKS) Start() error {
 	s.Server = server
 	s.I2PTunnelStatus = i2ptunnel.I2PTunnelStatusStarting
 	s.Server.Handle = s
+	s.I2PTunnelStatus = i2ptunnel.I2PTunnelStatusRunning
 	s.Server.ListenAndServe(s)
 
 	return nil