diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go index 8e4c6bba7..6b759682e 100644 --- a/wgengine/magicsock/magicsock.go +++ b/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 diff --git a/wgengine/magicsock/magicsock_test.go b/wgengine/magicsock/magicsock_test.go index 84b4e999e..69d64b43f 100644 --- a/wgengine/magicsock/magicsock_test.go +++ b/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. diff --git a/wgengine/userspace.go b/wgengine/userspace.go index 32683e3c7..7d289670a 100644 --- a/wgengine/userspace.go +++ b/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()