From 38aa2687acc93b04be56a85283f22b00e5f627b8 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Sun, 28 Jul 2024 11:43:01 -0700
Subject: [PATCH] make hard NAT work

Change-Id: Id7b4195639c7fbff17c5b52e3780a1fdbcefbd31
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 natlab/natlabd/nat.go     | 16 +++++++++-------
 natlab/natlabd/natlabd.go |  2 +-
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/natlab/natlabd/nat.go b/natlab/natlabd/nat.go
index 8bdb8b2b1..ebe7d8c6c 100644
--- a/natlab/natlabd/nat.go
+++ b/natlab/natlabd/nat.go
@@ -4,6 +4,8 @@ import (
 	"math/rand/v2"
 	"net/netip"
 	"time"
+
+	"tailscale.com/util/mak"
 )
 
 // NATTable is what a NAT implementation is expected to do.
@@ -101,16 +103,16 @@ func (n *hardNAT) PickOutgoingSrc(src, dst netip.AddrPort, at time.Time) (wanSrc
 	// Instead of proper data structures that would be efficient, we instead
 	// just loop a bunch and look for a free port. This project is only used
 	// by tests and doesn't care about performance, this is good enough.
-	port := src.Port() // start with the port the client used? (TODO: does FreeBSD do this? make knob?)
 	for {
+		port := rand.N(uint16(32<<10)) + 32<<10 // pick some "ephemeral" port
 		ki := hardKeyIn{wanPort: port, src: dst}
 		if _, ok := n.in[ki]; ok {
-			// Collision. Try again.
-			port = rand.N(uint16(32<<10)) + 32<<10 // pick some "ephemeral" port
+			// Port already in use.
 			continue
 		}
-		n.in[ki] = lanAddrAndTime{lanAddr: src, at: at}
-		n.out[ko] = portMappingAndTime{port: port, at: at}
+		mak.Set(&n.in, ki, lanAddrAndTime{lanAddr: src, at: at})
+		mak.Set(&n.out, ko, portMappingAndTime{port: port, at: at})
+		return netip.AddrPortFrom(n.wanIP, port)
 	}
 }
 
@@ -160,8 +162,8 @@ func (n *easyNAT) PickOutgoingSrc(src, dst netip.AddrPort, at time.Time) (wanSrc
 		port := 32<<10 + (start + off)
 		if _, ok := n.in[port]; !ok {
 			// Found a free port.
-			n.out[src] = portMappingAndTime{port: port, at: at}
-			n.in[port] = lanAddrAndTime{lanAddr: src, at: at}
+			mak.Set(&n.out, src, portMappingAndTime{port: port, at: at})
+			mak.Set(&n.in, port, lanAddrAndTime{lanAddr: src, at: at})
 			return netip.AddrPortFrom(n.wanIP, port)
 		}
 	}
diff --git a/natlab/natlabd/natlabd.go b/natlab/natlabd/natlabd.go
index f7debe970..aa1815554 100644
--- a/natlab/natlabd/natlabd.go
+++ b/natlab/natlabd/natlabd.go
@@ -81,7 +81,7 @@ func main() {
 	}
 
 	s.nodes[node1.mac] = node1
-	net1.SetNATTable(&oneToOneNAT{wanIP: net1.wanIP, lanIP: node1.lanIP})
+	net1.SetNATTable(&hardNAT{wanIP: net1.wanIP})
 	net2 := &network{
 		s:     s,
 		mac:   MAC{0x52, 0x54, 0x00, 0x01, 0x01, 0x2},