wgengine/magicsock: delete Conn.UpdatePeers, derive peer state internally

[This commit is pulled out of a branch that ultimately removes the
wgcfg.Config.Peers field and removes all O(n peers) processing when
handling deltas]

magicsock.Conn.UpdatePeers existed so wgengine.Reconfig could tell
magicsock the set of WireGuard peers from cfg.Peers, used only to
garbage collect the derpRoute and peerLastDerp maps and to ReSTUN when
the first peers appear. magicsock already learns the full peer list
directly from LocalBackend via SetNetworkMap, UpsertPeer, and
RemovePeer, so do that bookkeeping there and delete the API and its
cfg.Peers use.

Updates #12542

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

wgengine/magicsock/magicsock.go

@@ -327,12 +327,6 @@ type Conn struct {
 	// for a period of time before withdrawing them.
 	endpointTracker endpointTracker
 
-	// peerSet is the set of peers that are currently configured in
-	// WireGuard. These are not used to filter inbound or outbound
-	// traffic at all, but only to track what state can be cleaned up
-	// in other maps below that are keyed by peer public key.
-	peerSet set.Set[key.NodePublic]
-
 	// peerMap tracks the networkmap Node entity for each peer
 	// by node key, node ID, and discovery key.
 	peerMap peerMap
@@ -2801,31 +2795,6 @@ func (c *Conn) SetPrivateKey(privateKey key.NodePrivate) error {
 	return nil
 }
 
-// UpdatePeers is called when the set of WireGuard peers changes. It
-// then removes any state for old peers.
-//
-// The caller passes ownership of newPeers map to UpdatePeers.
-func (c *Conn) UpdatePeers(newPeers set.Set[key.NodePublic]) {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-
-	oldPeers := c.peerSet
-	c.peerSet = newPeers
-
-	// Clean up any key.NodePublic-keyed maps for peers that no longer
-	// exist.
-	for peer := range oldPeers {
-		if !newPeers.Contains(peer) {
-			delete(c.derpRoute, peer)
-			delete(c.peerLastDerp, peer)
-		}
-	}
-
-	if len(oldPeers) == 0 && len(newPeers) > 0 {
-		go c.ReSTUN("non-zero-peers")
-	}
-}
-
 // debugRingBufferSize returns a maximum size for our set of endpoint ring
 // buffers by assuming that a single large update is ~500 bytes, and that we
 // want to not use more than 1MiB of memory on phones / 4MiB on other devices.
@@ -3168,7 +3137,21 @@ func (c *Conn) updateNodes(self tailcfg.NodeView, peers []tailcfg.NodeView) (pee
 		// Duplicate NodeIDs in the input shouldn't happen, but log if so.
 		c.logf("[unexpected] magicsock.updateNodes: %d peers input but %d unique IDs", len(peers), len(newPeers))
 	}
+	// Clean up node-key-keyed state for peers that are gone or have
+	// rotated their node key.
+	for id, prev := range c.peersByID {
+		if n, ok := newPeers[id]; !ok || n.Key() != prev.Key() {
+			delete(c.derpRoute, prev.Key())
+			delete(c.peerLastDerp, prev.Key())
+		}
+	}
+	hadPeers := len(c.peersByID) > 0
 	c.peersByID = newPeers
+	// A key-less conn can't STUN; the ReSTUN("set-private-key") on
+	// first key covers peers that arrived before the key.
+	if !hadPeers && len(newPeers) > 0 && !c.privateKey.IsZero() {
+		go c.ReSTUN("non-zero-peers")
+	}
 
 	// If the upsert pass left stale endpoints in peerMap (peers removed
 	// relative to before), clean them up.
@@ -3322,7 +3305,16 @@ func (c *Conn) UpsertPeer(n tailcfg.NodeView) {
 		return
 	}
 	flags := c.debugFlagsLocked()
+	if prev, ok := c.peersByID[n.ID()]; ok && prev.Key() != n.Key() {
+		delete(c.derpRoute, prev.Key())
+		delete(c.peerLastDerp, prev.Key())
+	}
+	hadPeers := len(c.peersByID) > 0
 	mak.Set(&c.peersByID, n.ID(), n)
+	// See the matching trigger in updateNodes for why the key check.
+	if !hadPeers && !c.privateKey.IsZero() {
+		go c.ReSTUN("non-zero-peers")
+	}
 	c.upsertPeerLocked(n, flags, debugRingBufferSize(len(c.peersByID)))
 
 	var relayUpsert candidatePeerRelay
@@ -3360,6 +3352,8 @@ func (c *Conn) RemovePeer(nid tailcfg.NodeID) {
 		return
 	}
 	delete(c.peersByID, nid)
+	delete(c.derpRoute, prev.Key())
+	delete(c.peerLastDerp, prev.Key())
 	if ep, ok := c.peerMap.endpointForNodeID(nid); ok {
 		c.peerMap.deleteEndpoint(ep)
 	}
@@ -3629,7 +3623,7 @@ func (c *Conn) shouldDoPeriodicReSTUNLocked() bool {
 	if c.networkDown() || c.homeless {
 		return false
 	}
-	if len(c.peerSet) == 0 || c.privateKey.IsZero() {
+	if len(c.peersByID) == 0 || c.privateKey.IsZero() {
 		// If no peers, not worth doing.
 		// Also don't if there's no key (not running).
 		return false

wgengine/magicsock/magicsock_test.go

@@ -21,6 +21,7 @@
 	"os"
 	"reflect"
 	"runtime"
+	"slices"
 	"strconv"
 	"strings"
 	"sync"
@@ -69,6 +70,7 @@
 	"tailscale.com/util/clientmetric"
 	"tailscale.com/util/eventbus"
 	"tailscale.com/util/eventbus/eventbustest"
+	"tailscale.com/util/mak"
 	"tailscale.com/util/must"
 	"tailscale.com/util/racebuild"
 	"tailscale.com/util/set"
@@ -365,11 +367,6 @@ func meshStacks(logf logger.Logf, mutateNetmap func(idx int, nm *netmap.NetworkM
 		for i, m := range ms {
 			nm := buildNetmapLocked(i)
 			m.conn.SetNetworkMap(nm.SelfNode, nm.Peers)
-			peerSet := make(set.Set[key.NodePublic], len(nm.Peers))
-			for _, peer := range nm.Peers {
-				peerSet.Add(peer.Key())
-			}
-			m.conn.UpdatePeers(peerSet)
 			wg, err := nmcfg.WGCfg(ms[i].privateKey, nm, logf, 0, "")
 			if err != nil {
 				// We're too far from the *testing.T to be graceful,
@@ -2328,6 +2325,116 @@ func TestSetNetworkMapWithNoPeers(t *testing.T) {
 	}
 }
 
+// Tests that the full-set peer update path (SetNetworkMap) and the
+// incremental per-peer paths (UpsertPeer, RemovePeer) maintain the
+// node-key-keyed DERP state (derpRoute and peerLastDerp) identically
+// when peers are removed or rotate their node keys. LocalBackend uses
+// either path depending on whether it received a full netmap or a
+// delta, so a change to the cleanup in one path must be made to both.
+func TestPeerDERPStateCleanup(t *testing.T) {
+	peerA := &tailcfg.Node{ID: 1, Key: randNodeKey(), DiscoKey: randDiscoKey(), Endpoints: eps("192.168.1.2:1")}
+	peerB := &tailcfg.Node{ID: 2, Key: randNodeKey(), DiscoKey: randDiscoKey(), Endpoints: eps("192.168.1.2:2")}
+	peerARotated := peerA.Clone()
+	peerARotated.Key = randNodeKey()
+
+	steps := []struct {
+		name  string
+		peers []*tailcfg.Node // the desired peer set after the step
+	}{
+		{"initial", []*tailcfg.Node{peerA, peerB}},
+		{"rotate-A-key", []*tailcfg.Node{peerARotated, peerB}},
+		{"remove-B", []*tailcfg.Node{peerARotated}},
+		{"remove-all", nil},
+	}
+
+	// seedDERPState populates the node-key-keyed DERP maps for each
+	// current peer, as receiving DERP traffic from them would.
+	seedDERPState := func(c *Conn) {
+		c.mu.Lock()
+		defer c.mu.Unlock()
+		for _, n := range c.peersByID {
+			mak.Set(&c.derpRoute, n.Key(), derpRoute{regionID: 1})
+			mak.Set(&c.peerLastDerp, n.Key(), 1)
+		}
+	}
+
+	// derpKeys returns a canonical dump of the node-key-keyed DERP
+	// state that peer removals and key rotations must clean up.
+	derpKeys := func(c *Conn) []string {
+		c.mu.Lock()
+		defer c.mu.Unlock()
+		var got []string
+		for k := range c.derpRoute {
+			got = append(got, "derpRoute "+k.String())
+		}
+		for k := range c.peerLastDerp {
+			got = append(got, "peerLastDerp "+k.String())
+		}
+		slices.Sort(got)
+		return got
+	}
+
+	// stateString returns a canonical dump of all the per-peer state
+	// that the two update paths must maintain identically.
+	stateString := func(c *Conn) string {
+		c.mu.Lock()
+		lines := []string{fmt.Sprintf("peerMap %d", c.peerMap.nodeCount())}
+		for id, n := range c.peersByID {
+			lines = append(lines, fmt.Sprintf("peer %d %s", id, n.Key()))
+		}
+		c.mu.Unlock()
+		lines = append(lines, derpKeys(c)...)
+		slices.Sort(lines)
+		return strings.Join(lines, "\n")
+	}
+
+	connFull := newTestConn(t) // updated via SetNetworkMap
+	defer connFull.Close()
+	connInc := newTestConn(t) // updated via UpsertPeer/RemovePeer
+	defer connInc.Close()
+
+	prevKeyByID := map[tailcfg.NodeID]key.NodePublic{}
+	for _, step := range steps {
+		connFull.SetNetworkMap(tailcfg.NodeView{}, nodeViews(step.peers))
+
+		curIDs := set.Set[tailcfg.NodeID]{}
+		for _, n := range step.peers {
+			curIDs.Add(n.ID)
+			connInc.UpsertPeer(n.View())
+		}
+		for id := range prevKeyByID {
+			if !curIDs.Contains(id) {
+				connInc.RemovePeer(id)
+			}
+		}
+
+		// The DERP state seeded before this step must survive only for
+		// peers still present with an unchanged node key.
+		var want []string
+		for _, n := range step.peers {
+			if k, ok := prevKeyByID[n.ID]; ok && k == n.Key {
+				want = append(want, "derpRoute "+n.Key.String(), "peerLastDerp "+n.Key.String())
+			}
+		}
+		slices.Sort(want)
+		if got := derpKeys(connFull); !slices.Equal(got, want) {
+			t.Errorf("step %q: SetNetworkMap path DERP state = %q; want %q", step.name, got, want)
+		}
+
+		// Both update paths must have produced identical state.
+		if full, inc := stateString(connFull), stateString(connInc); full != inc {
+			t.Errorf("step %q: state mismatch between update paths\nSetNetworkMap path:\n%s\nUpsertPeer/RemovePeer path:\n%s", step.name, full, inc)
+		}
+
+		clear(prevKeyByID)
+		for _, n := range step.peers {
+			prevKeyByID[n.ID] = n.Key
+		}
+		seedDERPState(connFull)
+		seedDERPState(connInc)
+	}
+}
+
 // newWireguard starts up a new wireguard-go device attached to a test tun, and
 // returns the device, tun and endpoint port. To add peers call device.IpcSet
 // with UAPI instructions.

wgengine/userspace.go

@@ -970,7 +970,6 @@ func (e *userspaceEngine) Reconfig(cfg *wgcfg.Config, routerCfg *router.Config,
 
 	e.lastCfgFull = *cfg.Clone()
 
-	e.magicConn.UpdatePeers(peerSet)
 	e.magicConn.SetPreferredPort(listenPort)
 	e.magicConn.UpdatePMTUD()