Localhost working - is this the best way to do it?

netstack2/handlers.go

@@ -62,22 +62,24 @@ const (
 
 // TCPHandler handles TCP connections from netstack
 type TCPHandler struct {
-	stack *stack.Stack
+	stack        *stack.Stack
+	proxyHandler *ProxyHandler
 }
 
 // UDPHandler handles UDP connections from netstack
 type UDPHandler struct {
-	stack *stack.Stack
+	stack        *stack.Stack
+	proxyHandler *ProxyHandler
 }
 
 // NewTCPHandler creates a new TCP handler
-func NewTCPHandler(s *stack.Stack) *TCPHandler {
+func NewTCPHandler(s *stack.Stack, ph *ProxyHandler) *TCPHandler {
-	return &TCPHandler{stack: s}
+	return &TCPHandler{stack: s, proxyHandler: ph}
 }
 
 // NewUDPHandler creates a new UDP handler
-func NewUDPHandler(s *stack.Stack) *UDPHandler {
+func NewUDPHandler(s *stack.Stack, ph *ProxyHandler) *UDPHandler {
-	return &UDPHandler{stack: s}
+	return &UDPHandler{stack: s, proxyHandler: ph}
 }
 
 // InstallTCPHandler installs the TCP forwarder on the stack
@@ -125,7 +127,16 @@ func (h *TCPHandler) handleTCPConn(netstackConn *gonet.TCPConn, id stack.Transpo
 
 	logger.Info("TCP Forwarder: Handling connection %s:%d -> %s:%d", srcIP, srcPort, dstIP, dstPort)
 
-	targetAddr := fmt.Sprintf("%s:%d", dstIP, dstPort)
+	// Check if there's a destination rewrite for this connection (e.g., localhost targets)
+	actualDstIP := dstIP
+	if h.proxyHandler != nil {
+		if rewrittenAddr, ok := h.proxyHandler.LookupDestinationRewrite(srcIP, dstIP, dstPort, uint8(tcp.ProtocolNumber)); ok {
+			actualDstIP = rewrittenAddr.String()
+			logger.Info("TCP Forwarder: Using rewritten destination %s (original: %s)", actualDstIP, dstIP)
+		}
+	}
+
+	targetAddr := fmt.Sprintf("%s:%d", actualDstIP, dstPort)
 
 	// Create context with timeout for connection establishment
 	ctx, cancel := context.WithTimeout(context.Background(), tcpConnectTimeout)
@@ -238,7 +249,16 @@ func (h *UDPHandler) handleUDPConn(netstackConn *gonet.UDPConn, id stack.Transpo
 
 	logger.Info("UDP Forwarder: Handling connection %s:%d -> %s:%d", srcIP, srcPort, dstIP, dstPort)
 
-	targetAddr := fmt.Sprintf("%s:%d", dstIP, dstPort)
+	// Check if there's a destination rewrite for this connection (e.g., localhost targets)
+	actualDstIP := dstIP
+	if h.proxyHandler != nil {
+		if rewrittenAddr, ok := h.proxyHandler.LookupDestinationRewrite(srcIP, dstIP, dstPort, uint8(udp.ProtocolNumber)); ok {
+			actualDstIP = rewrittenAddr.String()
+			logger.Info("UDP Forwarder: Using rewritten destination %s (original: %s)", actualDstIP, dstIP)
+		}
+	}
+
+	targetAddr := fmt.Sprintf("%s:%d", actualDstIP, dstPort)
 
 	// Resolve target address
 	remoteUDPAddr, err := net.ResolveUDPAddr("udp", targetAddr)

netstack2/proxy.go

@@ -145,6 +145,14 @@ type connKey struct {
 	proto   uint8
 }
 
+// destKey identifies a destination for handler lookups (without source port since it may change)
+type destKey struct {
+	srcIP   string
+	dstIP   string
+	dstPort uint16
+	proto   uint8
+}
+
 // natState tracks NAT translation state for reverse translation
 type natState struct {
 	originalDst netip.Addr // Original destination before DNAT
@@ -160,6 +168,7 @@ type ProxyHandler struct {
 	udpHandler        *UDPHandler
 	subnetLookup      *SubnetLookup
 	natTable          map[connKey]*natState
+	destRewriteTable  map[destKey]netip.Addr // Maps original dest to rewritten dest for handler lookups
 	natMu             sync.RWMutex
 	enabled           bool
 }
@@ -178,10 +187,11 @@ func NewProxyHandler(options ProxyHandlerOptions) (*ProxyHandler, error) {
 	}
 
 	handler := &ProxyHandler{
-		enabled:      true,
+		enabled:          true,
-		subnetLookup: NewSubnetLookup(),
+		subnetLookup:     NewSubnetLookup(),
-		natTable:     make(map[connKey]*natState),
+		natTable:         make(map[connKey]*natState),
-		proxyEp:      channel.New(1024, uint32(options.MTU), ""),
+		destRewriteTable: make(map[destKey]netip.Addr),
+		proxyEp:          channel.New(1024, uint32(options.MTU), ""),
 		proxyStack: stack.New(stack.Options{
 			NetworkProtocols: []stack.NetworkProtocolFactory{
 				ipv4.NewProtocol,
@@ -198,7 +208,7 @@ func NewProxyHandler(options ProxyHandlerOptions) (*ProxyHandler, error) {
 
 	// Initialize TCP handler if enabled
 	if options.EnableTCP {
-		handler.tcpHandler = NewTCPHandler(handler.proxyStack)
+		handler.tcpHandler = NewTCPHandler(handler.proxyStack, handler)
 		if err := handler.tcpHandler.InstallTCPHandler(); err != nil {
 			return nil, fmt.Errorf("failed to install TCP handler: %v", err)
 		}
@@ -206,7 +216,7 @@ func NewProxyHandler(options ProxyHandlerOptions) (*ProxyHandler, error) {
 
 	// Initialize UDP handler if enabled
 	if options.EnableUDP {
-		handler.udpHandler = NewUDPHandler(handler.proxyStack)
+		handler.udpHandler = NewUDPHandler(handler.proxyStack, handler)
 		if err := handler.udpHandler.InstallUDPHandler(); err != nil {
 			return nil, fmt.Errorf("failed to install UDP handler: %v", err)
 		}
@@ -251,6 +261,27 @@ func (p *ProxyHandler) RemoveSubnetRule(sourcePrefix, destPrefix netip.Prefix) {
 	p.subnetLookup.RemoveSubnet(sourcePrefix, destPrefix)
 }
 
+// LookupDestinationRewrite looks up the rewritten destination for a connection
+// This is used by TCP/UDP handlers to find the actual target address
+func (p *ProxyHandler) LookupDestinationRewrite(srcIP, dstIP string, dstPort uint16, proto uint8) (netip.Addr, bool) {
+	if p == nil || !p.enabled {
+		return netip.Addr{}, false
+	}
+
+	key := destKey{
+		srcIP:   srcIP,
+		dstIP:   dstIP,
+		dstPort: dstPort,
+		proto:   proto,
+	}
+
+	p.natMu.RLock()
+	defer p.natMu.RUnlock()
+
+	addr, ok := p.destRewriteTable[key]
+	return addr, ok
+}
+
 // resolveRewriteAddress resolves a rewrite address which can be either:
 // - An IP address with CIDR notation (e.g., "192.168.1.1/32") - returns the IP directly
 // - A plain IP address (e.g., "192.168.1.1") - returns the IP directly
@@ -407,6 +438,14 @@ func (p *ProxyHandler) HandleIncomingPacket(packet []byte) bool {
 				proto:   uint8(protocol),
 			}
 
+			// Key for handler lookups (doesn't include srcPort for flexibility)
+			dKey := destKey{
+				srcIP:   srcAddr.String(),
+				dstIP:   dstAddr.String(),
+				dstPort: dstPort,
+				proto:   uint8(protocol),
+			}
+
 			// Check if we already have a NAT entry for this connection
 			p.natMu.RLock()
 			existingEntry, exists := p.natTable[key]
@@ -437,14 +476,23 @@ func (p *ProxyHandler) HandleIncomingPacket(packet []byte) bool {
 					originalDst: dstAddr,
 					rewrittenTo: newDst,
 				}
+				// Store destination rewrite for handler lookups
+				p.destRewriteTable[dKey] = newDst
 				p.natMu.Unlock()
 				logger.Debug("New NAT entry for connection: %s -> %s", dstAddr, newDst)
 			}
 
-			// Rewrite the packet
+			// Check if target is loopback - if so, don't rewrite packet destination
-			packet = p.rewritePacketDestination(packet, newDst)
+			// as gVisor will drop martian packets. Instead, the handlers will use
-			if packet == nil {
+			// destRewriteTable to find the actual target address.
-				return false
+			if !newDst.IsLoopback() {
+				// Rewrite the packet only for non-loopback destinations
+				packet = p.rewritePacketDestination(packet, newDst)
+				if packet == nil {
+					return false
+				}
+			} else {
+				logger.Debug("Target is loopback, not rewriting packet - handlers will use rewrite table")
 			}
 		}