Improve nat performance

device/nat.go

@@ -134,6 +134,14 @@ func IPv4SourceEquals(packet []byte, addr [4]byte) bool {
 // else, so this errs on the long side.
 var natEntryTTL = 5 * time.Minute
 
+// natRefreshInterval bounds how often a busy flow's entry timestamp actually
+// gets rewritten. A saturating connection (e.g. iperf) calls FixOutboundSource
+// or FixInboundDest on every single packet - refreshing on every one of them
+// would mean a map write (and, for new entries, a full-table prune) at line
+// rate instead of at most once per interval. natEntryTTL is minutes, so
+// resolution at this granularity costs nothing.
+const natRefreshInterval = time.Second
+
 type natKey struct {
 	proto uint8
 	port  uint16
@@ -182,11 +190,24 @@ func (n *ExitNodeNAT) FixOutboundSource(packet []byte, correctSrc [4]byte) {
 	}
 
 	key := natKey{proto, srcPort}
+	now := time.Now()
 
 	n.mu.Lock()
-	_, existed := n.seen[key]
-	n.seen[key] = time.Now()
-	n.prune()
+	t, existed := n.seen[key]
+	if existed && now.Sub(t) < natRefreshInterval {
+		// Already recorded recently enough - skip the write entirely. This is
+		// the common case for a busy flow: every packet gets here, but only
+		// one per interval needs to touch the map.
+		n.mu.Unlock()
+		return
+	}
+	n.seen[key] = now
+	if !existed {
+		// Only prune when the table is actually growing (a new connection),
+		// not on every packet - this is an O(map size) scan and the map only
+		// ever gains entries here.
+		n.prune()
+	}
 	n.mu.Unlock()
 
 	if !existed {
@@ -206,16 +227,17 @@ func (n *ExitNodeNAT) FixInboundDest(packet []byte, wrongDst [4]byte) {
 	}
 
 	key := natKey{proto, dstPort}
+	now := time.Now()
 
 	n.mu.Lock()
 	t, tracked := n.seen[key]
-	expired := tracked && time.Since(t) > natEntryTTL
+	expired := tracked && now.Sub(t) > natEntryTTL
 	if tracked {
 		if expired {
 			delete(n.seen, key)
 			tracked = false
-		} else {
-			n.seen[key] = time.Now()
+		} else if now.Sub(t) >= natRefreshInterval {
+			n.seen[key] = now
 		}
 	}
 	n.mu.Unlock()

device/nat_test.go

@@ -335,3 +335,49 @@ func TestExitNodeNATEntryExpires(t *testing.T) {
 		t.Errorf("expired entry was still translated: got %x, want unchanged %x", reply, original)
 	}
 }
+
+// BenchmarkExitNodeNATSteadyStateOutbound simulates a single busy flow (e.g.
+// an iperf upload) hammering FixOutboundSource, as happens for real since the
+// OS keeps stamping every packet of an affected socket with the wrong source
+// for the connection's whole lifetime, not just its first packet. Before the
+// refresh-throttling/prune-on-insert-only fix, every call here paid for a
+// map write plus a full-table prune; steady state should now be a single
+// lock/lookup/compare with no write and no allocation.
+func BenchmarkExitNodeNATSteadyStateOutbound(b *testing.B) {
+	wrongSrc := [4]byte{100, 89, 128, 9}
+	correctSrc := [4]byte{100, 89, 128, 4}
+	serverIP := [4]byte{100, 89, 128, 1}
+
+	nat := NewExitNodeNAT()
+	packet := exitNodeNATTestPacket(wrongSrc, serverIP, 52746, 80)
+	nat.FixOutboundSource(packet, correctSrc) // prime the entry
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		// FixOutboundSource rewrites in place, so re-derive a wrong-source
+		// packet each iteration rather than measuring the already-correct
+		// (no-op) fast path.
+		packet := exitNodeNATTestPacket(wrongSrc, serverIP, 52746, 80)
+		nat.FixOutboundSource(packet, correctSrc)
+	}
+}
+
+// BenchmarkExitNodeNATSteadyStateInbound is BenchmarkExitNodeNATSteadyStateOutbound's
+// counterpart for the download direction / ACK stream.
+func BenchmarkExitNodeNATSteadyStateInbound(b *testing.B) {
+	wrongSrc := [4]byte{100, 89, 128, 9}
+	correctSrc := [4]byte{100, 89, 128, 4}
+	serverIP := [4]byte{100, 89, 128, 1}
+
+	nat := NewExitNodeNAT()
+	outbound := exitNodeNATTestPacket(wrongSrc, serverIP, 52746, 80)
+	nat.FixOutboundSource(outbound, correctSrc) // establish the tracked port
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		reply := exitNodeNATTestPacket(serverIP, correctSrc, 80, 52746)
+		nat.FixInboundDest(reply, wrongSrc)
+	}
+}