UDP expose connection tracking (#104)

- Add connection tracking to the UDP expose functionality. 
- Fix a bug in the TCP proxy used by tunnel and expose that would sometimes close one end of the proxy too soon
- Added server messages when expose connections are made, similar to tunnel connection messages
- Server UDP tunnel connection messages now only trigger on first connections (when tracking starts), not every packet
- Bump dockerfile golang version to 1.24 to match current code requirements
- docker-compose file for demo environment now also mounts the src folder to /wiretap-live to make development testing much faster; code can be edited on the host and is immediately available for compilation in all containers.

Author: Aptimex

docker-compose.yml

@@ -51,6 +51,8 @@ services:
       - net.ipv6.conf.all.disable_ipv6=0
     environment:
     - DISPLAY=host.docker.internal:0
+    volumes: 
+      - ./src:/wiretap-live/
   server: 
     depends_on: 
       - client 
@@ -71,6 +73,8 @@ services:
       - net.ipv6.conf.all.disable_ipv6=0
     environment:
     - DISPLAY=host.docker.internal:0
+    volumes: 
+      - ./src:/wiretap-live/
   target:
     depends_on: 
       - client 
@@ -87,6 +91,8 @@ services:
       - SYS_MODULE
     sysctls:
       - net.ipv6.conf.all.disable_ipv6=0
+    volumes: 
+    - ./src:/wiretap-live/
   target2:
     depends_on: 
       - client 
@@ -100,6 +106,8 @@ services:
       - SYS_MODULE
     sysctls:
       - net.ipv6.conf.all.disable_ipv6=0
+    volumes: 
+      - ./src:/wiretap-live/
 
 
 networks:

src/cmd/serve.go

@@ -23,10 +23,10 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/network/ipv4"
 	"gvisor.dev/gvisor/pkg/tcpip/network/ipv6"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
+	gicmp "gvisor.dev/gvisor/pkg/tcpip/transport/icmp"
 	"gvisor.dev/gvisor/pkg/tcpip/transport/raw"
 	gtcp "gvisor.dev/gvisor/pkg/tcpip/transport/tcp"
 	gudp "gvisor.dev/gvisor/pkg/tcpip/transport/udp"
-	gicmp "gvisor.dev/gvisor/pkg/tcpip/transport/icmp"
 
 	"wiretap/peer"
 	"wiretap/transport/api"

src/transport/api/api.go

@@ -533,7 +533,7 @@ func handleExpose(tnet *netstack.Net, exposeMap *map[ExposeTuple]ExposeConn, exp
 				}
 
 				// Bind successful, expose port.
-				go transport.ForwardUdpPort(
+				go transport.ForwardUdpPortWithTracking(
 					tnet.Stack(),
 					c,
 					localAddr,

src/transport/transport.go

@@ -11,6 +11,7 @@ import (
 	"net/netip"
 	"strconv"
 	"sync"
+	"time"
 
 	"github.com/armon/go-socks5"
 	"github.com/google/gopacket"
@@ -39,6 +40,7 @@ type ConnCounts struct {
 }
 
 var connCounts ConnCounts
+const UDP_TIMEOUT = 60 * time.Second
 
 func init() {
 	connCounts.counts = make(map[netip.Addr]int)
@@ -139,24 +141,30 @@ func SendPacket(s *stack.Stack, packet []byte, addr *tcpip.FullAddress, netProto
 }
 
 func Proxy(src net.Conn, dst net.Conn) {
+	defer src.Close()
+	defer dst.Close()
 	var wg sync.WaitGroup
 
+	//log.Printf("Proxying between %v <-> %v and %v <-> %v\n", src.LocalAddr(), src.RemoteAddr(), dst.LocalAddr(), dst.RemoteAddr())
+
 	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		_, err := io.Copy(src, dst)
 		if err != nil {
 			log.Printf("error copying between connections: %v\n", err)
 		}
-		src.Close()
-		wg.Done()
 	}()
 
 	// Copy from peer to new connection.
-	_, nerr := io.Copy(dst, src)
-	if nerr != nil {
-		log.Printf("error copying between connections: %v\n", nerr)
-	}
-	dst.Close()
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		_, nerr := io.Copy(dst, src)
+		if nerr != nil {
+			log.Printf("error copying between connections: %v\n", nerr)
+		}
+	}()
 
 	// Wait for both copies to finish.
 	wg.Wait()
@@ -172,9 +180,12 @@ func ForwardTcpPort(s *stack.Stack, l net.Listener, localAddr tcpip.FullAddress,
 			return
 		}
 
+		log.Printf("(client [%v]:%v) <- Expose: TCP <- %v", remoteAddr.Addr, remoteAddr.Port, conn.RemoteAddr().String())
+
 		// Proxy between conns.
 		go func() {
 			var nc net.Conn
+			// since the localAddr has a port value of 0 it should auto-select a random available source port
 			nc, err = gonet.DialTCPWithBind(
 				ctx,
 				s,
@@ -195,6 +206,7 @@ func ForwardTcpPort(s *stack.Stack, l net.Listener, localAddr tcpip.FullAddress,
 
 // ForwardUdpPort proxies UDP datagrams by forwarding datagrams to a peer, and then returns responses to the last remote address to talk to this endpoint.
 // No connection tracking is in place at this time.
+// DEPRECATED. Use ForwardUdpPortWithTracking instead.
 func ForwardUdpPort(s *stack.Stack, conn *net.UDPConn, localAddr tcpip.FullAddress, remoteAddr tcpip.FullAddress, np tcpip.NetworkProtocolNumber) {
 	var wg sync.WaitGroup
 	var clientAddr *netip.AddrPort
@@ -248,6 +260,7 @@ func ForwardUdpPort(s *stack.Stack, conn *net.UDPConn, localAddr tcpip.FullAddre
 		}
 		lock.Lock()
 		if clientAddr == nil {
+			log.Println("UDP client addr null, cannot forward response")
 			lock.Unlock()
 			continue
 		}
@@ -263,6 +276,172 @@ func ForwardUdpPort(s *stack.Stack, conn *net.UDPConn, localAddr tcpip.FullAddre
 	wg.Wait()
 }
 
+// ForwardUdpPortWithTracking proxies UDP datagrams by forwarding datagrams to and from a peer. All connections are tracked so that the client responses can be sent to the correct sender
+//
+// "conn" is the listening socket on the "real" network (the port being forwarded). 
+// "LocalAddr" should be the API listener for this server inside Wiretap's network (src), but port 0 (so a random ephemeral port is assigned).
+// "remoteAddr" is the Client's IP(v6) address and port in Wiretap's network (dst)
+func ForwardUdpPortWithTracking(s *stack.Stack, conn *net.UDPConn, localAddr tcpip.FullAddress, remoteAddr tcpip.FullAddress, np tcpip.NetworkProtocolNumber) {
+	var wg sync.WaitGroup
+	const bufSize = 65535
+	type trackedConn struct {
+		udpConn    *gonet.UDPConn
+		lastActive time.Time
+	}
+	var connTrack = make(map[netip.AddrPort]*trackedConn)
+
+	var ctLock sync.Mutex
+	var newConn = make(chan netip.AddrPort)
+
+	// Sanity check that we can create a connection to the target client port
+	_, err := gonet.DialUDP(s, &localAddr, &remoteAddr, np)
+	if err != nil {
+		log.Println("failed to proxy UDP conn:", err)
+		conn.Close()
+		return
+	}
+
+	
+	// Watchdog: periodically check for idle connections
+	var stopWatchdog = make(chan bool)
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		ticker := time.NewTicker(30 * time.Second)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ticker.C:
+				expire := time.Now().Add(-1 * UDP_TIMEOUT)
+				ctLock.Lock()
+				for addr, clientConn := range connTrack {
+					if clientConn.lastActive.Before(expire) {
+						log.Printf("Closing idle UDP forward: (client %v) <- Expose: UDP <- %v", clientConn.udpConn.RemoteAddr().String(), addr.String())
+						clientConn.udpConn.Close()
+						delete(connTrack, addr)
+					}
+				}
+				ctLock.Unlock()
+			case <-stopWatchdog:
+				return
+			}
+		}
+	}()
+
+	// Process incoming packets (new or existing “connections”)
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		buf := make([]byte, bufSize)
+
+		for {
+			n, addr, err := conn.ReadFromUDPAddrPort(buf)
+			if err != nil {
+				// This "listener" gets closed when the forward is removed via API
+				log.Println("conn closed:", err)
+				close(newConn) //signal the other goroutines to shut down
+				stopWatchdog <-true
+
+				ctLock.Lock()
+				for _, t := range connTrack {
+					t.udpConn.Close()
+				}
+				ctLock.Unlock()
+				
+				return
+			}
+
+			ctLock.Lock()
+			clientConn, exists := connTrack[addr]
+			ctLock.Unlock()
+			if !exists {
+				// If this connection is not already being tracked, set it up
+				udpConn, err := gonet.DialUDP(s, &localAddr, &remoteAddr, np)
+				if err != nil {
+					log.Println("failed to establish new proxy conn:", err)
+					continue
+				}
+				clientConn = &trackedConn{udpConn: udpConn, lastActive: time.Now()}
+				ctLock.Lock()
+				connTrack[addr] = clientConn
+				ctLock.Unlock()
+
+				// This blocks until it gets picked up by the response handler routine below
+				newConn <- addr
+
+			} else {
+				clientConn.lastActive = time.Now()
+			}
+
+			// Forward payload to the client
+			// We currently have no way to capture Dest Unreachable ICMP here, so the first Write() will never fail due to that. 
+			// But the network stack will remember them, so the second Write() will fail and trigger the cleanup. 
+			_, err = clientConn.udpConn.Write(buf[:n])
+			if err != nil {
+				log.Println("failed to forward UDP packet to client:", err)
+				clientConn.udpConn.Close()
+				continue
+			}
+		}
+	}()
+
+	// Process new connections – set up routines that handle return (outgoing) packets
+	for {
+		ncAddr, ok := <-newConn
+		if !ok {
+			break
+		}
+
+		ctLock.Lock()
+		clientConn, exists := connTrack[ncAddr]
+		ctLock.Unlock()
+		if !exists {
+			log.Printf("new UDP forward connection %v marked for processing but not found\n", ncAddr)
+			continue
+		}
+
+		clientRemoteAddr := clientConn.udpConn.RemoteAddr().String()
+		log.Printf("(client %v) <- Expose: UDP <- %v", clientRemoteAddr, ncAddr.String())
+
+		// Spawn routine to forward responses back to the original sender
+		wg.Add(1)
+		go func(targetAddr netip.AddrPort) {
+			defer wg.Done()
+			buf := make([]byte, bufSize)
+
+			for {
+				n, err := clientConn.udpConn.Read(buf)
+				if err != nil {
+					log.Printf("response conn closed: (client %v) <- Expose: UDP <- %v : %s", clientRemoteAddr, targetAddr.String(), err)
+					clientConn.udpConn.Close()
+
+					ctLock.Lock()
+					delete(connTrack, targetAddr)
+					ctLock.Unlock()
+					break
+				}
+
+				ctLock.Lock()
+				if t, ok := connTrack[targetAddr]; ok {
+					t.lastActive = time.Now()
+				}
+				ctLock.Unlock()
+
+				_, err = conn.WriteToUDPAddrPort(buf[:n], targetAddr)
+				if err != nil {
+					log.Println("failed to send client UDP response:", err)
+					continue
+				}
+			}
+		}(ncAddr)
+	}
+
+
+	wg.Wait()
+	log.Printf("All routines for UDP forward %v successfully shut down\n", conn.LocalAddr().String())
+}
+
 // ForwardTcpPort proxies TCP connections by accepting connections and proxying them back to the client.
 func ForwardDynamic(s *stack.Stack, l *net.Listener, localAddr tcpip.FullAddress, remoteAddr tcpip.FullAddress, np tcpip.NetworkProtocolNumber) {
 	dialer := func(ctx context.Context, network, addr string) (net.Conn, error) {

src/transport/udp/udp.go

@@ -57,8 +57,6 @@ type Config struct {
 // TODO: Clean this up. Can't use UDPForwarder because it doesn't offer a way to return false, which is required to send Unreachables.
 func Handler(c Config) func(stack.TransportEndpointID, *stack.PacketBuffer) bool {
 	return func(teid stack.TransportEndpointID, pkb *stack.PacketBuffer) bool {
-		log.Printf("(client %s) - Transport: UDP -> %s", net.JoinHostPort(teid.RemoteAddress.String(), fmt.Sprint(teid.RemotePort)), net.JoinHostPort(teid.LocalAddress.String(), fmt.Sprint(teid.LocalPort)))
-
 		packetClone := pkb.Clone()
 		go func() {
 			newPacket(packetClone, c.Tnet.Stack())
@@ -129,7 +127,7 @@ func getDataFromPacket(packet *stack.PacketBuffer) []byte {
 	return transHeader.Payload()
 }
 
-// NewPacket handles every new packet and sending it to the proper UDP dialer.
+// NewPacket handles every new packet and sends it to the proper UDP dialer.
 func newPacket(packet *stack.PacketBuffer, s *stack.Stack) {
 	netHeader := packet.Network()
 	transHeader := header.UDP(netHeader.Payload())
@@ -160,6 +158,7 @@ func newPacket(packet *stack.PacketBuffer, s *stack.Stack) {
 	pktChan = make(chan *stack.PacketBuffer, 1)
 	connMapWrite(conn, pktChan)
 
+	log.Printf("(client %v) - Transport: UDP -> %v", source, dest)
 	go handleConn(conn, port, s)
 
 	pktChan <- packet.Clone()
@@ -202,7 +201,7 @@ func handleConn(conn udpConn, port int, s *stack.Stack) {
 		sourceMapDecrement(conn.Source)
 	}()
 
-	err = newConn.SetDeadline(time.Now().Add(30 * time.Second))
+	err = newConn.SetDeadline(time.Now().Add(transport.UDP_TIMEOUT))
 	if err != nil {
 		log.Println("failed to set deadline", err)
 	}
@@ -230,7 +229,7 @@ func handleConn(conn udpConn, port int, s *stack.Stack) {
 			}
 
 			// Reset timer, we got a packet.
-			err = newConn.SetDeadline(time.Now().Add(30 * time.Second))
+			err = newConn.SetDeadline(time.Now().Add(transport.UDP_TIMEOUT))
 			if err != nil {
 				log.Println("failed to set deadline:", err)
 			}
@@ -261,7 +260,7 @@ func handleConn(conn udpConn, port int, s *stack.Stack) {
 		}
 
 		// Reset timer, we got a packet.
-		err = newConn.SetDeadline(time.Now().Add(30 * time.Second))
+		err = newConn.SetDeadline(time.Now().Add(transport.UDP_TIMEOUT))
 		if err != nil {
 			log.Println("failed to set deadline:", err)
 		}

wiretap.Dockerfile

@@ -1,4 +1,4 @@
-FROM golang:1.23
+FROM golang:1.24
 
 ARG http_proxy
 ARG https_proxy