tstest/natlab/vnet: add port mapping that might not work yet

Change-Id: Iaf274d250398973790873534b236d5cbb34fbe0e
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>

tstest/integration/nat/nat_test.go

@@ -58,6 +58,13 @@ func hard(c *vnet.Config) *vnet.Node {
 		fmt.Sprintf("10.0.%d.1/24", n), vnet.HardNAT))
 }
 
+func hardPMP(c *vnet.Config) *vnet.Node {
+	n := c.NumNodes() + 1
+	return c.AddNode(c.AddNetwork(
+		fmt.Sprintf("2.%d.%d.%d", n, n, n), // public IP
+		fmt.Sprintf("10.7.%d.1/24", n), vnet.HardNAT, vnet.NATPMP))
+}
+
 func (nt *natTest) runTest(node1, node2 addNodeFunc) {
 	t := nt.tb
 
@@ -229,3 +236,8 @@ func TestEasyHard(t *testing.T) {
 	nt := newNatTest(t)
 	nt.runTest(easy, hard)
 }
+
+func TestEasyHardPMP(t *testing.T) {
+	nt := newNatTest(t)
+	nt.runTest(easy, hardPMP)
+}

tstest/natlab/vnet/nat.go

@@ -33,7 +33,11 @@ type IPPool interface {
 	// and if so, its IP address.
 	SoleLANIP() (_ netip.Addr, ok bool)
 
-	// TODO: port availability stuff for interacting with portmapping
+	// IsPublicPortUsed reports whether the provided WAN IP+port is in use by
+	// anything. (In particular, the NAT-PMP/etc port mappers might have taken
+	// a port.) Implementations should check this before allocating a port,
+	// and then they should report IsPublicPortUsed themselves for that port.
+	IsPublicPortUsed(netip.AddrPort) bool
 }
 
 // newTableFunc is a constructor for a NAT table.
@@ -86,6 +90,10 @@ type NATTable interface {
 	// address of a machine on the local network address, usually a private
 	// LAN IP.
 	PickIncomingDst(src, dst netip.AddrPort, at time.Time) (lanDst netip.AddrPort)
+
+	// IsPublicPortUsed reports whether the provided WAN IP+port is in use by
+	// anything. The port mapper uses this to avoid grabbing an in-use port.
+	IsPublicPortUsed(netip.AddrPort) bool
 }
 
 // oneToOneNAT is a 1:1 NAT, like a typical EC2 VM.
@@ -112,6 +120,10 @@ func (n *oneToOneNAT) PickIncomingDst(src, dst netip.AddrPort, at time.Time) (la
 	return netip.AddrPortFrom(n.lanIP, dst.Port())
 }
 
+func (n *oneToOneNAT) IsPublicPortUsed(netip.AddrPort) bool {
+	return true // all ports are owned by the 1:1 NAT
+}
+
 type srcDstTuple struct {
 	src netip.AddrPort
 	dst netip.AddrPort
@@ -136,6 +148,7 @@ type lanAddrAndTime struct {
 // This is shown as "MappingVariesByDestIP: true" by netcheck, and what
 // Tailscale calls "Hard NAT".
 type hardNAT struct {
+	pool  IPPool
 	wanIP netip.Addr
 
 	out map[srcDstTuple]portMappingAndTime
@@ -144,10 +157,22 @@ type hardNAT struct {
 
 func init() {
 	registerNATType(HardNAT, func(p IPPool) (NATTable, error) {
-		return &hardNAT{wanIP: p.WANIP()}, nil
+		return &hardNAT{pool: p, wanIP: p.WANIP()}, nil
 	})
 }
 
+func (n *hardNAT) IsPublicPortUsed(ap netip.AddrPort) bool {
+	if ap.Addr() != n.wanIP {
+		return false
+	}
+	for k := range n.in {
+		if k.wanPort == ap.Port() {
+			return true
+		}
+	}
+	return false
+}
+
 func (n *hardNAT) PickOutgoingSrc(src, dst netip.AddrPort, at time.Time) (wanSrc netip.AddrPort) {
 	ko := srcDstTuple{src, dst}
 	if pm, ok := n.out[ko]; ok {
@@ -165,6 +190,10 @@ func (n *hardNAT) PickOutgoingSrc(src, dst netip.AddrPort, at time.Time) (wanSrc
 	// by tests and doesn't care about performance, this is good enough.
 	for {
 		port := rand.N(uint16(32<<10)) + 32<<10 // pick some "ephemeral" port
+		if n.pool.IsPublicPortUsed(netip.AddrPortFrom(n.wanIP, port)) {
+			continue
+		}
+
 		ki := hardKeyIn{wanPort: port, src: dst}
 		if _, ok := n.in[ki]; ok {
 			// Port already in use.
@@ -197,6 +226,7 @@ func (n *hardNAT) PickIncomingDst(src, dst netip.AddrPort, at time.Time) (lanDst
 // Unlike Linux, this implementation is capped at 32k entries and doesn't resort
 // to other allocation strategies when all 32k WAN ports are taken.
 type easyNAT struct {
+	pool    IPPool
 	wanIP   netip.Addr
 	out     map[netip.AddrPort]portMappingAndTime
 	in      map[uint16]lanAddrAndTime
@@ -205,10 +235,18 @@ type easyNAT struct {
 
 func init() {
 	registerNATType(EasyNAT, func(p IPPool) (NATTable, error) {
-		return &easyNAT{wanIP: p.WANIP()}, nil
+		return &easyNAT{pool: p, wanIP: p.WANIP()}, nil
 	})
 }
 
+func (n *easyNAT) IsPublicPortUsed(ap netip.AddrPort) bool {
+	if ap.Addr() != n.wanIP {
+		return false
+	}
+	_, ok := n.in[ap.Port()]
+	return ok
+}
+
 func (n *easyNAT) PickOutgoingSrc(src, dst netip.AddrPort, at time.Time) (wanSrc netip.AddrPort) {
 	mak.Set(&n.lastOut, srcDstTuple{src, dst}, at)
 	if pm, ok := n.out[src]; ok {
@@ -224,6 +262,9 @@ func (n *easyNAT) PickOutgoingSrc(src, dst netip.AddrPort, at time.Time) (wanSrc
 		port := 32<<10 + (start+off)%(32<<10)
 		if _, ok := n.in[port]; !ok {
 			wanAddr := netip.AddrPortFrom(n.wanIP, port)
+			if n.pool.IsPublicPortUsed(wanAddr) {
+				continue
+			}
 
 			// Found a free port.
 			mak.Set(&n.out, src, portMappingAndTime{port: port, at: at})

tstest/natlab/vnet/vnet.go

@@ -23,6 +23,7 @@
 	"fmt"
 	"io"
 	"log"
+	"math/rand/v2"
 	"net"
 	"net/http"
 	"net/http/httptest"
@@ -394,6 +395,11 @@ func (m MAC) String() string {
 	return fmt.Sprintf("%02x:%02x:%02x:%02x:%02x:%02x", m[0], m[1], m[2], m[3], m[4], m[5])
 }
 
+type portMapping struct {
+	dst    netip.AddrPort // LAN IP:port
+	expiry time.Time
+}
+
 type network struct {
 	s         *Server
 	mac       MAC
@@ -408,6 +414,7 @@ type network struct {
 	natStyle syncs.AtomicValue[NAT]
 	natMu    sync.Mutex // held while using + changing natTable
 	natTable NATTable
+	portMap  map[netip.AddrPort]portMapping // WAN ip:port -> LAN ip:port
 
 	// writeFunc is a map of MAC -> func to write to that MAC.
 	// It contains entries for connected nodes only.
@@ -1201,7 +1208,56 @@ func (n *network) doNATOut(src, dst netip.AddrPort) (newSrc netip.AddrPort) {
 func (n *network) doNATIn(src, dst netip.AddrPort) (newDst netip.AddrPort) {
 	n.natMu.Lock()
 	defer n.natMu.Unlock()
-	return n.natTable.PickIncomingDst(src, dst, time.Now())
+
+	now := time.Now()
+
+	// First see if there's a port mapping, before doing NAT.
+	if lanAP, ok := n.portMap[dst]; ok {
+		if now.Before(lanAP.expiry) {
+			return lanAP.dst
+		}
+		delete(n.portMap, dst)
+		return netip.AddrPort{}
+	}
+
+	return n.natTable.PickIncomingDst(src, dst, now)
+}
+
+// IsPublicPortUsed reports whether the given public port is currently in use.
+//
+// n.natMu must be held by the caller. (It's only called by nat implementations
+// which are always called with natMu held))
+func (n *network) IsPublicPortUsed(ap netip.AddrPort) bool {
+	_, ok := n.portMap[ap]
+	return ok
+}
+
+func (n *network) doPortMap(src netip.Addr, dstLANPort, wantExtPort uint16, sec int) (gotPort uint16, ok bool) {
+	n.natMu.Lock()
+	defer n.natMu.Unlock()
+
+	wanAP := netip.AddrPortFrom(n.wanIP, wantExtPort)
+
+	if sec == 0 {
+		lanAP, ok := n.portMap[wanAP]
+		if ok && lanAP.dst.Addr() == src {
+			delete(n.portMap, wanAP)
+		}
+		return 0, false
+	}
+
+	for try := 0; try < 20_000; try++ {
+		if !n.natTable.IsPublicPortUsed(wanAP) {
+			mak.Set(&n.portMap, wanAP, portMapping{
+				dst:    netip.AddrPortFrom(src, dstLANPort),
+				expiry: time.Now().Add(time.Duration(sec) * time.Second),
+			})
+			return wanAP.Port(), true
+		}
+		wantExtPort = rand.N(uint16(32<<10)) + 32<<10
+		wanAP = netip.AddrPortFrom(n.wanIP, wantExtPort)
+	}
+	return 0, false
 }
 
 func (n *network) createARPResponse(pkt gopacket.Packet) ([]byte, error) {
@@ -1274,8 +1330,42 @@ func (n *network) handleNATPMPRequest(req UDPPacket) {
 		return
 	}
 
+	// Map UDP request
+	if len(req.Payload) == 12 && req.Payload[0] == 0 && req.Payload[1] == 1 {
+		// https://www.rfc-editor.org/rfc/rfc6886#section-3.3
+		// "00 01 00 00 ed 40 00 00 00 00 1c 20" =>
+		//   00 ver
+		//   01 op=map UDP
+		//   00 00 reserved  (0 in request; in response, this is the result code)
+		//   ed 40 internal port 60736
+		//   00 00 suggested external port
+		//   00 00 1c 20 suggested lifetime in seconds (7200 sec = 2 hours)
+		internalPort := binary.BigEndian.Uint16(req.Payload[4:6])
+		wantExtPort := binary.BigEndian.Uint16(req.Payload[6:8])
+		lifetimeSec := binary.BigEndian.Uint32(req.Payload[8:12])
+		gotPort, ok := n.doPortMap(req.Src.Addr(), internalPort, wantExtPort, int(lifetimeSec))
+		if !ok {
+			log.Printf("NAT-PMP map request for %v:%d failed", req.Src.Addr(), internalPort)
+			return
+		}
+		res := make([]byte, 0, 12)
+		res = append(res,
+			0,     // version 0 (NAT-PMP)
+			1+128, // response to op 1
+			0, 0,  // result code success
+		)
+		res = binary.BigEndian.AppendUint16(res, internalPort)
+		res = binary.BigEndian.AppendUint16(res, gotPort)
+		res = binary.BigEndian.AppendUint32(res, lifetimeSec)
+		n.WriteUDPPacketNoNAT(UDPPacket{
+			Src:     req.Dst,
+			Dst:     req.Src,
+			Payload: res,
+		})
+		return
+	}
 	log.Printf("TODO: handle NAT-PMP packet % 02x", req.Payload)
-	// TODO: handle NAT-PMP packet 00 01 00 00 ed 40 00 00 00 00 1c 20
+
 }
 
 // UDPPacket is a UDP packet.