tstest/integration: add ping test w/ masquerades

Updates tailscale/corp#8020

Co-authored-by: Melanie Warrick <[email protected]>
Signed-off-by: Maisem Ali <[email protected]>

tstest/integration/integration_test.go

@@ -39,6 +39,7 @@ import (
 	"tailscale.com/tailcfg"
 	"tailscale.com/tstest"
 	"tailscale.com/tstest/integration/testcontrol"
+	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
 )
 
@@ -503,6 +504,110 @@ func TestOneNodeUpWindowsStyle(t *testing.T) {
 	d1.MustCleanShutdown(t)
 }
 
+// TestNATPing creates two nodes, n1 and n2, sets up masquerades for both and
+// tries to do bi-directional pings between them.
+func TestNATPing(t *testing.T) {
+	t.Parallel()
+	env := newTestEnv(t)
+	registerNode := func() (*testNode, key.NodePublic) {
+		n := newTestNode(t, env)
+		n.StartDaemon()
+		n.AwaitListening()
+		n.MustUp()
+		n.AwaitRunning()
+		k := n.MustStatus().Self.PublicKey
+		return n, k
+	}
+	n1, k1 := registerNode()
+	n2, k2 := registerNode()
+
+	n1IP := n1.AwaitIP()
+	n2IP := n2.AwaitIP()
+
+	n1ExternalIP := netip.MustParseAddr("100.64.1.1")
+	n2ExternalIP := netip.MustParseAddr("100.64.2.1")
+
+	tests := []struct {
+		name       string
+		pairs      []testcontrol.MasqueradePair
+		n1SeesN2IP netip.Addr
+		n2SeesN1IP netip.Addr
+	}{
+		{
+			name:       "no_nat",
+			n1SeesN2IP: n2IP,
+			n2SeesN1IP: n1IP,
+		},
+		{
+			name: "n1_has_external_ip",
+			pairs: []testcontrol.MasqueradePair{
+				{
+					Node:              k1,
+					Peer:              k2,
+					NodeMasqueradesAs: n1ExternalIP,
+				},
+			},
+			n1SeesN2IP: n2IP,
+			n2SeesN1IP: n1ExternalIP,
+		},
+		{
+			name: "n2_has_external_ip",
+			pairs: []testcontrol.MasqueradePair{
+				{
+					Node:              k2,
+					Peer:              k1,
+					NodeMasqueradesAs: n2ExternalIP,
+				},
+			},
+			n1SeesN2IP: n2ExternalIP,
+			n2SeesN1IP: n1IP,
+		},
+		{
+			name: "both_have_external_ips",
+			pairs: []testcontrol.MasqueradePair{
+				{
+					Node:              k1,
+					Peer:              k2,
+					NodeMasqueradesAs: n1ExternalIP,
+				},
+				{
+					Node:              k2,
+					Peer:              k1,
+					NodeMasqueradesAs: n2ExternalIP,
+				},
+			},
+			n1SeesN2IP: n2ExternalIP,
+			n2SeesN1IP: n1ExternalIP,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			env.Control.SetMasqueradeAddresses(tc.pairs)
+
+			s1 := n1.MustStatus()
+			n2AsN1Peer := s1.Peer[k2]
+			if got := n2AsN1Peer.TailscaleIPs[0]; got != tc.n1SeesN2IP {
+				t.Fatalf("n1 sees n2 as %v; want %v", got, tc.n1SeesN2IP)
+			}
+
+			s2 := n2.MustStatus()
+			n1AsN2Peer := s2.Peer[k1]
+			if got := n1AsN2Peer.TailscaleIPs[0]; got != tc.n2SeesN1IP {
+				t.Fatalf("n2 sees n1 as %v; want %v", got, tc.n2SeesN1IP)
+			}
+
+			if err := n1.Tailscale("ping", tc.n1SeesN2IP.String()).Run(); err != nil {
+				t.Fatal(err)
+			}
+
+			if err := n2.Tailscale("ping", tc.n2SeesN1IP.String()).Run(); err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+}
+
 func TestLogoutRemovesAllPeers(t *testing.T) {
 	t.Parallel()
 	env := newTestEnv(t)

tstest/integration/testcontrol/testcontrol.go

@@ -60,6 +60,12 @@ type Server struct {
 	pubKey     key.MachinePublic
 	privKey    key.ControlPrivate // not strictly needed vs. MachinePrivate, but handy to test type interactions.
 
+	// masquerades is the set of masquerades that should be applied to
+	// MapResponses sent to clients. It is keyed by the requesting nodes
+	// public key, and then the peer node's public key. The value is the
+	// masquerade address to use for that peer.
+	masquerades map[key.NodePublic]map[key.NodePublic]netip.Addr // node => peer => SelfNodeV4MasqAddrForThisPeer IP
+
 	noisePubKey  key.MachinePublic
 	noisePrivKey key.ControlPrivate // not strictly needed vs. MachinePrivate, but handy to test type interactions.
 
@@ -288,6 +294,48 @@ func (s *Server) serveMachine(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+// MasqueradePair is a pair of nodes and the IP address that the
+// Node masquerades as for the Peer.
+//
+// Setting this will have future MapResponses for Node to have
+// Peer.SelfNodeV4MasqAddrForThisPeer set to NodeMasqueradesAs.
+// MapResponses for the Peer will now see Node.Addresses as
+// NodeMasqueradesAs.
+type MasqueradePair struct {
+	Node              key.NodePublic
+	Peer              key.NodePublic
+	NodeMasqueradesAs netip.Addr
+}
+
+// SetMasqueradeAddresses sets the masquerade addresses for the server.
+// See MasqueradePair for more details.
+func (s *Server) SetMasqueradeAddresses(pairs []MasqueradePair) {
+	m := make(map[key.NodePublic]map[key.NodePublic]netip.Addr)
+	for _, p := range pairs {
+		if m[p.Node] == nil {
+			m[p.Node] = make(map[key.NodePublic]netip.Addr)
+		}
+		m[p.Node][p.Peer] = p.NodeMasqueradesAs
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.masquerades = m
+	s.updateLocked("SetMasqueradeAddresses", s.nodeIDsLocked(0))
+}
+
+// nodeIDsLocked returns the node IDs of all nodes in the server, except
+// for the node with the given ID.
+func (s *Server) nodeIDsLocked(except tailcfg.NodeID) []tailcfg.NodeID {
+	var ids []tailcfg.NodeID
+	for _, n := range s.nodes {
+		if n.ID == except {
+			continue
+		}
+		ids = append(ids, n.ID)
+	}
+	return ids
+}
+
 // Node returns the node for nodeKey. It's always nil or cloned memory.
 func (s *Server) Node(nodeKey key.NodePublic) *tailcfg.Node {
 	s.mu.Lock()
@@ -588,12 +636,7 @@ func (s *Server) UpdateNode(n *tailcfg.Node) (peersToUpdate []tailcfg.NodeID) {
 		panic("zero nodekey")
 	}
 	s.nodes[n.Key] = n.Clone()
-	for _, n2 := range s.nodes {
-		if n.ID != n2.ID {
-			peersToUpdate = append(peersToUpdate, n2.ID)
-		}
-	}
-	return peersToUpdate
+	return s.nodeIDsLocked(n.ID)
 }
 
 func (s *Server) incrInServeMap(delta int) {
@@ -791,11 +834,28 @@ func (s *Server) MapResponse(req *tailcfg.MapRequest) (res *tailcfg.MapResponse,
 		DNSConfig:   dns,
 		ControlTime: &t,
 	}
+
+	s.mu.Lock()
+	nodeMasqs := s.masquerades[node.Key]
+	s.mu.Unlock()
 	for _, p := range s.AllNodes() {
-		if p.StableID != node.StableID {
-			res.Peers = append(res.Peers, p)
+		if p.StableID == node.StableID {
+			continue
 		}
+		if masqIP := nodeMasqs[p.Key]; masqIP.IsValid() {
+			p.SelfNodeV4MasqAddrForThisPeer = masqIP
+		}
+
+		s.mu.Lock()
+		peerAddress := s.masquerades[p.Key][node.Key]
+		s.mu.Unlock()
+		if peerAddress.IsValid() {
+			p.Addresses[0] = netip.PrefixFrom(peerAddress, peerAddress.BitLen())
+			p.AllowedIPs[0] = netip.PrefixFrom(peerAddress, peerAddress.BitLen())
+		}
+		res.Peers = append(res.Peers, p)
 	}
+
 	sort.Slice(res.Peers, func(i, j int) bool {
 		return res.Peers[i].ID < res.Peers[j].ID
 	})