Perform stateful nat

device/nat.go

@@ -1,57 +1,84 @@
 package device
 
-import "encoding/binary"
+import (
+	"encoding/binary"
+	"sync"
+	"time"
+
+	"github.com/fosrl/newt/logger"
+)
+
+const (
+	ipv4SrcOffset = 12
+	ipv4DstOffset = 16
+)
 
 // FixIPv4Source rewrites an IPv4 packet's source address to correctSrc if it
-// doesn't already match, incrementally fixing up the IPv4 header checksum
-// and (for TCP/UDP) the transport checksum so the packet stays valid.
+// doesn't already match. It returns whether a rewrite happened.
+func FixIPv4Source(packet []byte, correctSrc [4]byte) bool {
+	return fixIPv4Address(packet, ipv4SrcOffset, correctSrc)
+}
+
+// FixIPv4Dest rewrites an IPv4 packet's destination address to correctDst if
+// it doesn't already match. It returns whether a rewrite happened.
+func FixIPv4Dest(packet []byte, correctDst [4]byte) bool {
+	return fixIPv4Address(packet, ipv4DstOffset, correctDst)
+}
+
+// fixIPv4Address rewrites the IPv4 address at the given header offset (source
+// or destination) to newAddr if it doesn't already match, incrementally
+// fixing up the IPv4 header checksum and (for TCP/UDP) the transport
+// checksum so the packet stays valid.
 //
-// The common case - source already correct - is a single 4-byte comparison
+// The common case - address already correct - is a single 4-byte comparison
 // and nothing else, so this is safe to call unconditionally on every packet
 // matched by a MiddleDevice rule. When a rewrite is needed, checksums are
 // updated via the RFC 1624 incremental method (add the delta of the changed
 // 16-bit words) rather than a full recompute over the packet, since only the
-// address field changed. ICMP has no pseudo-header dependency on the IP
-// addresses, so its checksum is left untouched. Non-IPv4 or malformed
-// packets are left untouched.
-func FixIPv4Source(packet []byte, correctSrc [4]byte) {
+// address field changed. The formula is agnostic to which field (source or
+// destination) changed - both are covered by the IPv4 header checksum and
+// the TCP/UDP pseudo-header checksum identically. ICMP has no pseudo-header
+// dependency on the IP addresses, so its checksum is left untouched.
+// Non-IPv4 or malformed packets are left untouched.
+func fixIPv4Address(packet []byte, offset int, newAddr [4]byte) bool {
 	if len(packet) < 20 || packet[0]>>4 != 4 {
-		return
+		return false
 	}
 
-	if packet[12] == correctSrc[0] && packet[13] == correctSrc[1] &&
-		packet[14] == correctSrc[2] && packet[15] == correctSrc[3] {
-		return
+	if packet[offset] == newAddr[0] && packet[offset+1] == newAddr[1] &&
+		packet[offset+2] == newAddr[2] && packet[offset+3] == newAddr[3] {
+		return false
 	}
 
 	ihl := int(packet[0]&0x0f) * 4
 	if ihl < 20 || len(packet) < ihl {
-		return
+		return false
 	}
 
-	oldSrc := [4]byte{packet[12], packet[13], packet[14], packet[15]}
+	old := [4]byte{packet[offset], packet[offset+1], packet[offset+2], packet[offset+3]}
 
 	ipChecksum := binary.BigEndian.Uint16(packet[10:12])
-	binary.BigEndian.PutUint16(packet[10:12], checksumAdjust(ipChecksum, oldSrc[:], correctSrc[:]))
+	binary.BigEndian.PutUint16(packet[10:12], checksumAdjust(ipChecksum, old[:], newAddr[:]))
 
 	switch packet[9] {
 	case 6: // TCP
 		if len(packet) >= ihl+20 {
 			off := ihl + 16
-			old := binary.BigEndian.Uint16(packet[off : off+2])
-			binary.BigEndian.PutUint16(packet[off:off+2], checksumAdjust(old, oldSrc[:], correctSrc[:]))
+			c := binary.BigEndian.Uint16(packet[off : off+2])
+			binary.BigEndian.PutUint16(packet[off:off+2], checksumAdjust(c, old[:], newAddr[:]))
 		}
 	case 17: // UDP
 		if len(packet) >= ihl+8 {
 			off := ihl + 6
-			old := binary.BigEndian.Uint16(packet[off : off+2])
-			if old != 0 { // zero means checksum not used - must stay zero
-				binary.BigEndian.PutUint16(packet[off:off+2], checksumAdjust(old, oldSrc[:], correctSrc[:]))
+			c := binary.BigEndian.Uint16(packet[off : off+2])
+			if c != 0 { // zero means checksum not used - must stay zero
+				binary.BigEndian.PutUint16(packet[off:off+2], checksumAdjust(c, old[:], newAddr[:]))
 			}
 		}
 	}
 
-	copy(packet[12:16], correctSrc[:])
+	copy(packet[offset:offset+4], newAddr[:])
+	return true
 }
 
 // checksumAdjust incrementally updates a ones-complement checksum after some
@@ -73,3 +100,145 @@ func checksumAdjust(checksum uint16, old, new []byte) uint16 {
 
 	return ^uint16(sum)
 }
+
+// ipv4L4Ports extracts the TCP/UDP source and destination ports from an IPv4
+// packet. ok is false for anything else (non-IPv4, non-TCP/UDP, malformed).
+func ipv4L4Ports(packet []byte) (proto uint8, srcPort, dstPort uint16, ok bool) {
+	if len(packet) < 20 || packet[0]>>4 != 4 {
+		return 0, 0, 0, false
+	}
+	proto = packet[9]
+	if proto != 6 && proto != 17 {
+		return 0, 0, 0, false
+	}
+	ihl := int(packet[0]&0x0f) * 4
+	if ihl < 20 || len(packet) < ihl+4 {
+		return 0, 0, 0, false
+	}
+	srcPort = binary.BigEndian.Uint16(packet[ihl : ihl+2])
+	dstPort = binary.BigEndian.Uint16(packet[ihl+2 : ihl+4])
+	return proto, srcPort, dstPort, true
+}
+
+// IPv4SourceEquals reports whether packet's IPv4 source address equals addr.
+func IPv4SourceEquals(packet []byte, addr [4]byte) bool {
+	return len(packet) >= 16 && packet[0]>>4 == 4 &&
+		packet[12] == addr[0] && packet[13] == addr[1] && packet[14] == addr[2] && packet[15] == addr[3]
+}
+
+// natEntryTTL bounds how long an ExitNodeNAT entry is honored without being
+// refreshed by further traffic on the same port. It's a var rather than a
+// const so tests can shrink it. Chosen generously relative to typical
+// request/response traffic - the only cost of expiring too early is the
+// original bug reappearing for that one flow, not corruption of anything
+// else, so this errs on the long side.
+var natEntryTTL = 5 * time.Minute
+
+type natKey struct {
+	proto uint8
+	port  uint16
+}
+
+// ExitNodeNAT tracks which local (protocol, port) pairs had their outbound
+// source address corrected by FixOutboundSource, so FixInboundDest can
+// translate the destination of the matching inbound reply back to the
+// address the local OS socket actually expects.
+//
+// This statefulness exists because rewriting the outbound packet's source
+// only changes what goes out on the wire - it does not change the local
+// kernel's own record of the connection's local address, which was already
+// selected and cached (in the socket's own connection state) at connect()/
+// send() time, before this packet ever reached this interception point.
+// Without also translating the reply's destination back, the OS can't match
+// the exit node's response to the socket waiting for it, and the request
+// hangs even though the corrected outbound packet reached the server fine.
+//
+// Entries are keyed by local port only (not the full flow), refreshed on
+// every match, and expire after natEntryTTL of inactivity - both so a later,
+// unrelated connection that happens to reuse the same ephemeral port isn't
+// wrongly treated as needing translation (e.g. one that was never affected
+// because it bound explicitly to the correct address), and so the table
+// doesn't grow unbounded over a long-lived tunnel.
+type ExitNodeNAT struct {
+	mu   sync.Mutex
+	seen map[natKey]time.Time
+}
+
+func NewExitNodeNAT() *ExitNodeNAT {
+	return &ExitNodeNAT{seen: make(map[natKey]time.Time)}
+}
+
+// FixOutboundSource rewrites packet's source to correctSrc (see
+// FixIPv4Source) and, if a rewrite was needed, remembers the packet's source
+// port so FixInboundDest knows to translate the reply back.
+func (n *ExitNodeNAT) FixOutboundSource(packet []byte, correctSrc [4]byte) {
+	if !FixIPv4Source(packet, correctSrc) {
+		return
+	}
+
+	proto, srcPort, _, ok := ipv4L4Ports(packet)
+	if !ok {
+		return
+	}
+
+	key := natKey{proto, srcPort}
+
+	n.mu.Lock()
+	_, existed := n.seen[key]
+	n.seen[key] = time.Now()
+	n.prune()
+	n.mu.Unlock()
+
+	if !existed {
+		logger.Debug("ExitNodeNAT: corrected outbound source for proto=%d port=%d", proto, srcPort)
+	}
+}
+
+// FixInboundDest rewrites packet's destination to wrongDst, but only if its
+// destination port matches an outbound flow FixOutboundSource actually
+// corrected - otherwise this connection was never affected by the bug (e.g.
+// a socket explicitly bound to the correct address already) and must be
+// left alone.
+func (n *ExitNodeNAT) FixInboundDest(packet []byte, wrongDst [4]byte) {
+	proto, _, dstPort, ok := ipv4L4Ports(packet)
+	if !ok {
+		return
+	}
+
+	key := natKey{proto, dstPort}
+
+	n.mu.Lock()
+	t, tracked := n.seen[key]
+	expired := tracked && time.Since(t) > natEntryTTL
+	if tracked {
+		if expired {
+			delete(n.seen, key)
+			tracked = false
+		} else {
+			n.seen[key] = time.Now()
+		}
+	}
+	n.mu.Unlock()
+
+	if expired {
+		logger.Warn("ExitNodeNAT: entry for proto=%d port=%d expired before a reply arrived on it - that flow's replies will be dropped by the OS from here on", proto, dstPort)
+	}
+
+	if !tracked {
+		return
+	}
+
+	FixIPv4Dest(packet, wrongDst)
+}
+
+// prune removes expired entries. Called with n.mu held, only from
+// FixOutboundSource so the cost is amortized over new outbound connections
+// rather than paid on every packet.
+func (n *ExitNodeNAT) prune() {
+	now := time.Now()
+	for k, t := range n.seen {
+		if now.Sub(t) > natEntryTTL {
+			delete(n.seen, k)
+		}
+	}
+}

device/nat_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"testing"
+	"time"
 )
 
 // onesComplementSum computes an RFC 1071 ones-complement checksum from
@@ -218,3 +219,119 @@ func TestFixIPv4SourceMalformedPacketNoPanic(t *testing.T) {
 	FixIPv4Source([]byte{0x45, 0x00, 0x00}, correctSrc)
 	FixIPv4Source([]byte{0x60, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, correctSrc) // IPv6 version nibble
 }
+
+func TestFixIPv4DestUDP(t *testing.T) {
+	src := [4]byte{192, 168, 1, 1}
+	wrongDst := [4]byte{10, 0, 0, 1}
+	correctDst := [4]byte{10, 0, 0, 2}
+	payload := []byte("reply")
+
+	udp := buildUDPSegment(src, wrongDst, payload)
+	ip := buildIPv4Header(src, wrongDst, 17, len(udp))
+	packet := append(ip, udp...)
+
+	if !FixIPv4Dest(packet, correctDst) {
+		t.Fatal("expected FixIPv4Dest to report a rewrite")
+	}
+	if got := [4]byte{packet[16], packet[17], packet[18], packet[19]}; got != correctDst {
+		t.Fatalf("dest = %v, want %v", got, correctDst)
+	}
+	verifyIPv4HeaderChecksum(t, packet)
+	verifyUDPChecksum(t, packet, src, correctDst)
+}
+
+// exitNodeNATTestPacket builds a minimal IPv4/UDP packet with the given
+// addresses and ports, for exercising ExitNodeNAT's port-based tracking.
+func exitNodeNATTestPacket(src, dst [4]byte, srcPort, dstPort uint16) []byte {
+	seg := make([]byte, 8)
+	binary.BigEndian.PutUint16(seg[0:2], srcPort)
+	binary.BigEndian.PutUint16(seg[2:4], dstPort)
+	binary.BigEndian.PutUint16(seg[4:6], uint16(len(seg)))
+
+	pseudo := make([]byte, 12+len(seg))
+	copy(pseudo[0:4], src[:])
+	copy(pseudo[4:8], dst[:])
+	pseudo[9] = 17
+	binary.BigEndian.PutUint16(pseudo[10:12], uint16(len(seg)))
+	copy(pseudo[12:], seg)
+	csum := onesComplementSum(pseudo)
+	if csum == 0 {
+		csum = 0xffff
+	}
+	binary.BigEndian.PutUint16(seg[6:8], csum)
+
+	ip := buildIPv4Header(src, dst, 17, len(seg))
+	return append(ip, seg...)
+}
+
+func TestExitNodeNATRoundTrip(t *testing.T) {
+	wrongSrc := [4]byte{100, 89, 128, 9}  // primary/site tunnel IP (the bug's default pick)
+	correctSrc := [4]byte{100, 89, 128, 4} // exit node's secondary tunnel IP
+	serverIP := [4]byte{100, 89, 128, 1}
+	const localPort = 52746
+
+	nat := NewExitNodeNAT()
+
+	// Outbound: kernel picked the wrong source; our fix rewrites it and should
+	// remember the local port so the reply gets translated.
+	outbound := exitNodeNATTestPacket(wrongSrc, serverIP, localPort, 80)
+	nat.FixOutboundSource(outbound, correctSrc)
+	if got := [4]byte{outbound[12], outbound[13], outbound[14], outbound[15]}; got != correctSrc {
+		t.Fatalf("outbound source = %v, want %v", got, correctSrc)
+	}
+
+	// Inbound reply: correctly addressed to correctSrc (the exit node saw the
+	// fixed source), but the OS's own connection state still expects wrongSrc.
+	reply := exitNodeNATTestPacket(serverIP, correctSrc, 80, localPort)
+	nat.FixInboundDest(reply, wrongSrc)
+	if got := [4]byte{reply[16], reply[17], reply[18], reply[19]}; got != wrongSrc {
+		t.Fatalf("reply dest = %v, want %v (translated back for the OS to match the socket)", got, wrongSrc)
+	}
+	verifyIPv4HeaderChecksum(t, reply)
+}
+
+func TestExitNodeNATUntrackedPortPassesThrough(t *testing.T) {
+	wrongSrc := [4]byte{100, 89, 128, 9}
+	correctSrc := [4]byte{100, 89, 128, 4}
+	serverIP := [4]byte{100, 89, 128, 1}
+	const localPort = 55555 // never seen by FixOutboundSource
+
+	nat := NewExitNodeNAT()
+
+	// A socket that was already, legitimately bound to correctSrc: its reply
+	// must not be touched, since translating it would misroute it away from
+	// the socket that's actually expecting it.
+	reply := exitNodeNATTestPacket(serverIP, correctSrc, 80, localPort)
+	original := append([]byte(nil), reply...)
+	nat.FixInboundDest(reply, wrongSrc)
+
+	if !bytes.Equal(reply, original) {
+		t.Errorf("untracked port was translated: got %x, want unchanged %x", reply, original)
+	}
+}
+
+func TestExitNodeNATEntryExpires(t *testing.T) {
+	origTTL := natEntryTTL
+	natEntryTTL = 10 * time.Millisecond
+	defer func() { natEntryTTL = origTTL }()
+
+	wrongSrc := [4]byte{100, 89, 128, 9}
+	correctSrc := [4]byte{100, 89, 128, 4}
+	serverIP := [4]byte{100, 89, 128, 1}
+	const localPort = 52746
+
+	nat := NewExitNodeNAT()
+
+	outbound := exitNodeNATTestPacket(wrongSrc, serverIP, localPort, 80)
+	nat.FixOutboundSource(outbound, correctSrc)
+
+	time.Sleep(50 * time.Millisecond)
+
+	reply := exitNodeNATTestPacket(serverIP, correctSrc, 80, localPort)
+	original := append([]byte(nil), reply...)
+	nat.FixInboundDest(reply, wrongSrc)
+
+	if !bytes.Equal(reply, original) {
+		t.Errorf("expired entry was still translated: got %x, want unchanged %x", reply, original)
+	}
+}

olm/connect.go

@@ -3,6 +3,7 @@ package olm
 import (
 	"encoding/json"
 	"fmt"
+	"net/netip"
 	"os"
 	"runtime"
 	"strconv"
@@ -168,6 +169,11 @@ func (o *Olm) handleConnect(msg websocket.WSMessage) {
 	if strings.Contains(interfaceIP, "/") {
 		interfaceIP = strings.Split(interfaceIP, "/")[0]
 	}
+	if addr, err := netip.ParseAddr(interfaceIP); err == nil {
+		o.primaryTunnelIP = addr
+	} else {
+		logger.Warn("Failed to parse tunnel IP %q: %v", interfaceIP, err)
+	}
 
 	// Create and start DNS proxy
 	o.dnsProxy, err = dns.NewDNSProxy(o.middleDev, o.tunnelConfig.MTU, wgData.UtilitySubnet, o.tunnelConfig.UpstreamDNS, o.tunnelConfig.TunnelDNS, interfaceIP, o.tunnelConfig.MatchDomains, o.tunnelConfig.PublicDNS)

olm/exitnode.go

@@ -132,18 +132,55 @@ persistent_keepalive_interval=%d`, util.FixKey(cfg.PublicKey), allowedIP, resolv
 	// exit node, which the exit node's WireGuard AllowedIPs filtering then
 	// silently drops. Fix it in-tunnel: intercept outbound packets addressed to
 	// the exit node and rewrite their source back to tunnelIP before WireGuard
-	// encrypts them. The fast path (source already correct) is cheap enough to
-	// also leave this on for a macOS CLI run, where the route above already
-	// gets it right.
+	// encrypts them.
+	//
+	// That alone isn't enough for anything that expects a reply (TCP, or any
+	// request/response over UDP): rewriting the outbound packet only changes
+	// what goes out on the wire - it doesn't change the OS's own connection
+	// state, which already recorded the *wrong* (primary) address as this
+	// socket's local address at connect()/send() time, before the packet ever
+	// reached this interception point. When the exit node's reply comes back
+	// correctly addressed to tunnelIP, the OS can't match it to a socket whose
+	// local address it thinks is the primary tunnel IP, and silently drops it -
+	// the connection hangs even though the corrected request reached the server
+	// fine. So also intercept inbound replies from the exit node and translate
+	// their destination back to the primary address, but only for flows we
+	// actually corrected outbound (tracked by ExitNodeNAT) - a socket that
+	// happened to already be bound to the correct address must be left alone.
+	//
+	// The fast path (address already correct) is cheap enough to also leave
+	// this on for a macOS CLI run, where the route above already gets it right.
 	if o.middleDev != nil && (runtime.GOOS == "darwin" || runtime.GOOS == "ios") {
-		if serverAddr, err := netip.ParseAddr(strings.Split(cfg.ServerIP, "/")[0]); err == nil {
-			if correctSrc, err := netip.ParseAddr(tunnelIPForRoute); err == nil && correctSrc.Is4() {
-				src := correctSrc.As4()
-				o.middleDev.AddRule(serverAddr, func(packet []byte) bool {
-					olmDevice.FixIPv4Source(packet, src)
-					return false
-				})
-			}
+		serverAddr, errS := netip.ParseAddr(strings.Split(cfg.ServerIP, "/")[0])
+		correctAddr, errC := netip.ParseAddr(tunnelIPForRoute)
+		switch {
+		case errS != nil || errC != nil || !correctAddr.Is4():
+			logger.Warn("Exit node NAT: skipping source-NAT setup, invalid address (server=%v tunnel=%v)", errS, errC)
+		case !o.primaryTunnelIP.IsValid() || !o.primaryTunnelIP.Is4():
+			logger.Warn("Exit node NAT: skipping source-NAT setup, no primary tunnel IP recorded")
+		default:
+			correctSrc := correctAddr.As4()
+			wrongSrc := o.primaryTunnelIP.As4()
+			serverSrc := serverAddr.As4()
+			nat := olmDevice.NewExitNodeNAT()
+
+			o.middleDev.AddRule(serverAddr, func(packet []byte) bool {
+				nat.FixOutboundSource(packet, correctSrc)
+				return false
+			})
+			o.middleDev.AddRule(correctAddr, func(packet []byte) bool {
+				// Only packets that actually came from the exit node's own
+				// peer should ever be translated - this rule's key (tunnelIP)
+				// is also used by the ICMP connectivity monitor's own address,
+				// so a defensive source check keeps this from ever touching
+				// unrelated traffic that happens to be addressed to tunnelIP.
+				if olmDevice.IPv4SourceEquals(packet, serverSrc) {
+					nat.FixInboundDest(packet, wrongSrc)
+				}
+				return false
+			})
+
+			logger.Debug("Exit node NAT: intercepting traffic to %s, translating source/dest between %s (primary) and %s (exit node secondary)", serverAddr, o.primaryTunnelIP, correctAddr)
 		}
 	}
 
@@ -212,6 +249,15 @@ func (o *Olm) removeExitNodePeerLocked() error {
 		if serverAddr, err := netip.ParseAddr(strings.Split(cfg.ServerIP, "/")[0]); err == nil {
 			o.middleDev.RemoveRule(serverAddr)
 		}
+		// Also removes the ICMP connectivity monitor's own rule under this same
+		// key (pm.ClearExitNode, called just above, already does this too - see
+		// RemoveRule's doc comment on it clearing every rule for a key rather
+		// than being handler-specific), so this call is normally a harmless
+		// no-op by the time it runs; kept for defensiveness/independence from
+		// that other subsystem's cleanup ordering.
+		if tunnelAddr, err := netip.ParseAddr(strings.Split(cfg.TunnelIP, "/")[0]); err == nil {
+			o.middleDev.RemoveRule(tunnelAddr)
+		}
 	}
 
 	serverIPForRoute := strings.Split(cfg.ServerIP, "/")[0] + "/32"

olm/olm.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"net/netip"
 	_ "net/http/pprof"
 	"os"
 	"os/exec"
@@ -62,6 +63,15 @@ type Olm struct {
 	// secondary address on the same interface/WireGuard device as the site peers.
 	exitNode   *ExitNodeConfig
 	exitNodeMu sync.Mutex
+
+	// primaryTunnelIP is the site tunnel's own address (wgData.TunnelIP), set once
+	// per connect in handleConnect. It's the interface's first/primary address -
+	// on macOS/iOS NetworkExtension, an unbound outbound socket's source gets
+	// stamped with this address by default even when the traffic should use an
+	// exit node's secondary address instead (see connectExitNode's NAT setup),
+	// and inbound replies need to be translated back to it for the OS to match
+	// them to the socket that's waiting.
+	primaryTunnelIP netip.Addr
 	// Power mode management
 	currentPowerMode string
 	powerModeMu      sync.Mutex