From e237e6ff795c99271b6e5a8321d95a0ec66058e6 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Thu, 8 Aug 2024 10:20:11 -0700
Subject: [PATCH] portmap fixes

Change-Id: Ia847580ba523acacadcb5fa8f87ccea98dc7ce41
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 cmd/vnet/vnet-main.go              |  3 ++-
 tstest/integration/nat/nat_test.go | 12 ++++++++---
 tstest/natlab/vnet/vnet.go         | 34 +++++++++++++++++++++++++-----
 3 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/cmd/vnet/vnet-main.go b/cmd/vnet/vnet-main.go
index acf85db0e..5fb83bce4 100644
--- a/cmd/vnet/vnet-main.go
+++ b/cmd/vnet/vnet-main.go
@@ -14,6 +14,7 @@
 	"time"
 
 	"tailscale.com/tstest/natlab/vnet"
+	"tailscale.com/types/logger"
 )
 
 var (
@@ -78,7 +79,7 @@ func main() {
 				log.Printf("NodeStatus: %v", err)
 				return
 			}
-			log.Printf("NodeStatus: %q", st)
+			log.Printf("NodeStatus: %v", logger.AsJSON(st))
 		}
 		for {
 			time.Sleep(5 * time.Second)
diff --git a/tstest/integration/nat/nat_test.go b/tstest/integration/nat/nat_test.go
index 20ca4c66e..7e99bf216 100644
--- a/tstest/integration/nat/nat_test.go
+++ b/tstest/integration/nat/nat_test.go
@@ -227,7 +227,7 @@ func ping(ctx context.Context, c *vnet.NodeAgentClient, target netip.Addr) (*ipn
 	n := 0
 	var res *ipnstate.PingResult
 	anyPong := false
-	for {
+	for n < 10 {
 		n++
 		pr, err := c.PingWithOpts(ctx, target, tailcfg.PingDisco, tailscale.PingOpts{})
 		if err != nil {
@@ -242,9 +242,16 @@ func ping(ctx context.Context, c *vnet.NodeAgentClient, target netip.Addr) (*ipn
 		if pr.DERPRegionID == 0 {
 			return pr, nil
 		}
-		time.Sleep(time.Second)
+		select {
+		case <-ctx.Done():
+		case <-time.After(time.Second):
+		}
 		res = pr
 	}
+	if res == nil {
+		return nil, errors.New("no ping response")
+	}
+	return res, nil
 }
 
 func up(ctx context.Context, c *vnet.NodeAgentClient) error {
@@ -270,7 +277,6 @@ func TestEasyEasy(t *testing.T) {
 }
 
 func TestEasyHard(t *testing.T) {
-	t.Skip()
 	nt := newNatTest(t)
 	nt.runTest(easy, hard)
 }
diff --git a/tstest/natlab/vnet/vnet.go b/tstest/natlab/vnet/vnet.go
index 98e791b97..2dea18735 100644
--- a/tstest/natlab/vnet/vnet.go
+++ b/tstest/natlab/vnet/vnet.go
@@ -1214,11 +1214,16 @@ func (n *network) doNATIn(src, dst netip.AddrPort) (newDst netip.AddrPort) {
 	// First see if there's a port mapping, before doing NAT.
 	if lanAP, ok := n.portMap[dst]; ok {
 		if now.Before(lanAP.expiry) {
+			log.Printf("XXX NAT: doNatIn: port mapping %v=>%v", dst, lanAP.dst)
 			return lanAP.dst
 		}
+		log.Printf("XXX NAT: doNatIn: port mapping EXPIRED for %v=>%v", dst, lanAP.dst)
 		delete(n.portMap, dst)
 		return netip.AddrPort{}
 	}
+	if len(n.portMap) > 0 {
+		log.Printf("XXX NAT: doNatIn: no port mapping for %v; have %v", dst, n.portMap)
+	}
 
 	return n.natTable.PickIncomingDst(src, dst, now)
 }
@@ -1236,7 +1241,12 @@ func (n *network) doPortMap(src netip.Addr, dstLANPort, wantExtPort uint16, sec
 	n.natMu.Lock()
 	defer n.natMu.Unlock()
 
+	if !n.portmap {
+		return 0, false
+	}
+
 	wanAP := netip.AddrPortFrom(n.wanIP, wantExtPort)
+	dst := netip.AddrPortFrom(src, dstLANPort)
 
 	if sec == 0 {
 		lanAP, ok := n.portMap[wanAP]
@@ -1246,12 +1256,24 @@ func (n *network) doPortMap(src netip.Addr, dstLANPort, wantExtPort uint16, sec
 		return 0, false
 	}
 
+	// See if they already have a mapping and extend expiry if so.
+	for k, v := range n.portMap {
+		if v.dst == dst {
+			n.portMap[k] = portMapping{
+				dst:    dst,
+				expiry: time.Now().Add(time.Duration(sec) * time.Second),
+			}
+			return k.Port(), true
+		}
+	}
+
 	for try := 0; try < 20_000; try++ {
-		if !n.natTable.IsPublicPortUsed(wanAP) {
+		if wanAP.Port() > 0 && !n.natTable.IsPublicPortUsed(wanAP) {
 			mak.Set(&n.portMap, wanAP, portMapping{
-				dst:    netip.AddrPortFrom(src, dstLANPort),
+				dst:    dst,
 				expiry: time.Now().Add(time.Duration(sec) * time.Second),
 			})
+			log.Printf("XXX allocated NAT mapping from %v to %v", wanAP, dst)
 			return wanAP.Port(), true
 		}
 		wantExtPort = rand.N(uint16(32<<10)) + 32<<10
@@ -1310,6 +1332,9 @@ func (n *network) createARPResponse(pkt gopacket.Packet) ([]byte, error) {
 }
 
 func (n *network) handleNATPMPRequest(req UDPPacket) {
+	if !n.portmap {
+		return
+	}
 	if string(req.Payload) == "\x00\x00" {
 		// https://www.rfc-editor.org/rfc/rfc6886#section-3.2
 
@@ -1348,12 +1373,13 @@ func (n *network) handleNATPMPRequest(req UDPPacket) {
 			log.Printf("NAT-PMP map request for %v:%d failed", req.Src.Addr(), internalPort)
 			return
 		}
-		res := make([]byte, 0, 12)
+		res := make([]byte, 0, 16)
 		res = append(res,
 			0,     // version 0 (NAT-PMP)
 			1+128, // response to op 1
 			0, 0,  // result code success
 		)
+		res = binary.BigEndian.AppendUint32(res, uint32(time.Now().Unix()))
 		res = binary.BigEndian.AppendUint16(res, internalPort)
 		res = binary.BigEndian.AppendUint16(res, gotPort)
 		res = binary.BigEndian.AppendUint32(res, lifetimeSec)
@@ -1439,11 +1465,9 @@ func (s *Server) takeAgentConnOne(n *node) (_ *agentConn, ok bool) {
 	for ac := range s.agentConns {
 		if ac.node == n {
 			s.agentConns.Delete(ac)
-			log.Printf("XXX takeAgentConnOne HIT for %v", n.mac)
 			return ac, true
 		}
 	}
-	log.Printf("XXX takeAgentConnOne MISS for %v", n.mac)
 	return nil, false
 }