From 06ac8c179974240d3c5cf8d64f9c44bedb8bf4cd Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 9 Jan 2026 17:41:34 +0000
Subject: [PATCH] Use b.UDP as source in return packets with proper routing

- Use b.UDP (actual response origin) as source address when available
- Build route from srcAddr (b.UDP or w.dest) to w.src (original client)
- Set UDP header with srcAddr.Port as source, w.src.Port as destination
- Validates address family matches w.src instead of w.dest
- Properly implements NAT with custom source addresses

Co-authored-by: RPRX <63339210+RPRX@users.noreply.github.com>
---
 proxy/tun/handler.go | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/proxy/tun/handler.go b/proxy/tun/handler.go
index 5002ed69..85e04b61 100644
--- a/proxy/tun/handler.go
+++ b/proxy/tun/handler.go
@@ -142,9 +142,15 @@ type udpWriter struct {
 
 func (w *udpWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
 	for _, b := range mb {
-		// Validate return packet address family matches expected destination
-		if b.UDP != nil && b.UDP.Address.Family() != w.dest.Address.Family() {
-			errors.LogWarning(context.Background(), "UDP return packet address family mismatch: expected ", w.dest.Address.Family(), ", got ", b.UDP.Address.Family())
+		// Use b.UDP as source if available, otherwise use w.dest
+		srcAddr := w.dest
+		if b.UDP != nil {
+			srcAddr = *b.UDP
+		}
+
+		// Validate address family matches
+		if srcAddr.Address.Family() != w.src.Address.Family() {
+			errors.LogWarning(context.Background(), "UDP return packet address family mismatch: expected ", w.src.Address.Family(), ", got ", srcAddr.Address.Family())
 			b.Release()
 			continue
 		}
@@ -154,10 +160,11 @@ func (w *udpWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
 			netProto = header.IPv6ProtocolNumber
 		}
 
+		// Build route from actual response source to original client
 		route, err := w.stack.FindRoute(
 			defaultNIC,
+			tcpip.AddrFromSlice(srcAddr.Address.IP()),
 			tcpip.AddrFromSlice(w.src.Address.IP()),
-			tcpip.AddrFromSlice(w.dest.Address.IP()),
 			netProto,
 			false,
 		)
@@ -172,8 +179,8 @@ func (w *udpWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
 		})
 		udp := header.UDP(pkt.TransportHeader().Push(header.UDPMinimumSize))
 		udp.Encode(&header.UDPFields{
-			SrcPort: uint16(w.src.Port),
-			DstPort: uint16(w.dest.Port),
+			SrcPort: uint16(srcAddr.Port),
+			DstPort: uint16(w.src.Port),
 			Length:  uint16(pkt.Size()),
 		})
 		xsum := route.PseudoHeaderChecksum(header.UDPProtocolNumber, uint16(pkt.Size()))