ipn/ipnlocal,wgengine: move disco-key change detection to nodeBackend

Engine.Reconfig previously diffed cfg.Peers disco keys against the
previous config to find restarted peers and flush their WireGuard
sessions, with a TSMP-learned-key map to suppress resets for key
changes that arrived over a working session. That was the last
per-peer state computed from wgcfg.Config.Peers inside the engine,
and it only ran on full reconfigs, so incremental netmap deltas
never got session resets at all.

Move the detection into nodeBackend, which sees every peer change:
full netmaps in SetNetMap and incremental upserts in
UpdateNetmapDelta both now report which peers changed disco keys,
with the same TSMP suppression and mismatch accounting as before.
LocalBackend acts on the result via a new Engine.ResetDevicePeer
method, which just removes the peer from the WireGuard device and
lets the peer lookup func lazily re-create it with fresh state.

LocalBackend.PatchDiscoKey now records TSMP-learned keys in
nodeBackend instead of forwarding to the engine, so the engine's
PatchDiscoKey method and tsmpLearnedDisco map are gone. The
controlclient patchDiscoKeyer interface becomes the exported
DiscoKeyUpdater so LocalBackend can compile-time assert that it
implements it, alongside its NetmapDeltaUpdater friends, replacing
the test that asserted the same of the engine.

This is one of the last steps toward removing Peers from wgcfg.Config.

Updates #12542

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

control/controlclient/auto.go

@@ -575,13 +575,13 @@ func (mrs mapRoutineState) UpdateUserProfiles(profiles map[tailcfg.UserID]tailcf
 	}
 }
 
-var _ patchDiscoKeyer = mapRoutineState{}
+var _ DiscoKeyUpdater = mapRoutineState{}
 
 func (mrs mapRoutineState) PatchDiscoKey(pub key.NodePublic, disco key.DiscoPublic) {
 	c := mrs.c
 	c.mu.Lock()
 	goodState := c.loggedIn && c.inMapPoll
-	dun, ok := c.observer.(patchDiscoKeyer)
+	dun, ok := c.observer.(DiscoKeyUpdater)
 	mapCtx := c.mapCtx
 	c.mu.Unlock()
 

control/controlclient/direct.go

@@ -286,10 +286,10 @@ type UserProfileUpdater interface {
 	UpdateUserProfiles(profiles map[tailcfg.UserID]tailcfg.UserProfileView) bool
 }
 
-// patchDiscoKeyer is an optional interface that can be implemented by an [Observer] to be
+// DiscoKeyUpdater is an optional interface that can be implemented by an [Observer] to be
 // notified about node disco keys received out-of-band from control, via
 // existing connection state.
-type patchDiscoKeyer interface {
+type DiscoKeyUpdater interface {
 	// PatchDiscoKey reports to the receiver that the specified disco key
 	// for node was obtained out-of-band from control.
 	PatchDiscoKey(key.NodePublic, key.DiscoPublic)

control/controlclient/map.go

@@ -362,7 +362,7 @@ func (ms *mapSession) handleNonKeepAliveMapResponse(ctx context.Context, resp *t
 }
 
 func (ms *mapSession) tryMarkDiscoAsLearnedFromTSMP(res *tailcfg.MapResponse) {
-	dun, ok := ms.netmapUpdater.(patchDiscoKeyer)
+	dun, ok := ms.netmapUpdater.(DiscoKeyUpdater)
 	if !ok {
 		return
 	}

control/controlclient/map_test.go

@@ -34,9 +34,7 @@ import (
 	"tailscale.com/util/eventbus/eventbustest"
 	"tailscale.com/util/mak"
 	"tailscale.com/util/must"
-	"tailscale.com/util/usermetric"
 	"tailscale.com/util/zstdframe"
-	"tailscale.com/wgengine"
 )
 
 func eps(s ...string) []netip.AddrPort {
@@ -1949,20 +1947,6 @@ func TestLearnZstdOfKeepAlive(t *testing.T) {
 	}
 }
 
-func TestPathDiscokeyerImplementations(t *testing.T) {
-	bus := eventbustest.NewBus(t)
-	ht := health.NewTracker(bus)
-	reg := new(usermetric.Registry)
-	e, err := wgengine.NewFakeUserspaceEngine(t.Logf, 0, ht, reg, bus)
-	if err != nil {
-		t.Fatal(err)
-	}
-	t.Cleanup(e.Close)
-	if _, ok := e.(patchDiscoKeyer); !ok {
-		t.Error("wgengine.userspaceEngine must implement patchDiscoKeyer")
-	}
-}
-
 func TestPeerIDAndKeyByTailscaleIP(t *testing.T) {
 	peerKey1 := key.NewNode().Public()
 	peerKey2 := key.NewNode().Public()

ipn/ipnlocal/local.go

@@ -2144,16 +2144,11 @@ func (b *LocalBackend) setControlClientStatusLocked(c controlclient.Client, st c
 	b.authReconfigLocked()
 }
 
+// PatchDiscoKey records that a peer's new disco key was learned via TSMP,
+// so the netmap update carrying the same change need not reset the peer's
+// WireGuard session. It implements [controlclient.DiscoKeyUpdater].
 func (b *LocalBackend) PatchDiscoKey(pub key.NodePublic, disco key.DiscoPublic) {
-	// PatchDiscoKey mirrors the implementation of [controlclient.patchDiscoKeyer].
-	// It is implemented here to avoid the dependency edge to controlclient, but must be kept
-	// in sync with the original implementation.
-	type patchDiscoKeyer interface {
-		PatchDiscoKey(key.NodePublic, key.DiscoPublic)
-	}
-	if e, ok := b.e.(patchDiscoKeyer); ok {
-		e.PatchDiscoKey(pub, disco)
-	}
+	b.currentNode().recordTSMPLearnedDisco(pub, disco)
 }
 
 type preferencePolicyInfo struct {
@@ -2436,6 +2431,7 @@ var (
 	_ controlclient.NetmapDeltaUpdater  = (*LocalBackend)(nil)
 	_ controlclient.PacketFilterUpdater = (*LocalBackend)(nil)
 	_ controlclient.UserProfileUpdater  = (*LocalBackend)(nil)
+	_ controlclient.DiscoKeyUpdater     = (*LocalBackend)(nil)
 )
 
 // UpdateNetmapDelta implements controlclient.NetmapDeltaUpdater.
@@ -2464,7 +2460,7 @@ func (b *LocalBackend) UpdateNetmapDelta(muts []netmap.NodeMutation) (handled bo
 	muts = b.tkaFilterDeltaMutsLocked(muts)
 	needsAuthReconfig := netmapDeltaNeedsAuthReconfig(cn, muts)
 
-	changedAllowedIPs, _ := cn.UpdateNetmapDelta(muts)
+	deltaRes, _ := cn.UpdateNetmapDelta(muts)
 	if buildfeatures.HasDrive {
 		// Drive's lazy remotes-source caches its rebuild keyed by this
 		// generation, so any delta — peer add/remove, address change,
@@ -2511,10 +2507,17 @@ func (b *LocalBackend) UpdateNetmapDelta(muts []netmap.NodeMutation) (handled bo
 	// the delta was applied above. Removed peers appear with a nil
 	// value and are removed from the device; SyncDevicePeer reads
 	// the live per-peer config, so only the keys matter here.
-	for k := range changedAllowedIPs {
+	for k := range deltaRes.ChangedAllowedIPs {
 		b.e.SyncDevicePeer(k)
 	}
 
+	// Reset the WireGuard session for peers whose disco key changed in
+	// a way that indicates a restart, flushing their dead session keys;
+	// each such peer is lazily re-created on demand with current state.
+	for k := range deltaRes.DiscoChanged {
+		b.e.ResetDevicePeer(k)
+	}
+
 	// Force a full authReconfig + SetSelfNode on any peer add or
 	// remove. The WireGuard device itself no longer depends on this:
 	// the SyncDevicePeer calls above keep its peer set delta-correct,
@@ -7289,7 +7292,7 @@ func (b *LocalBackend) setNetMapLocked(nm *netmap.NetworkMap) {
 	if nm != nil {
 		login = cmp.Or(profileFromView(nm.UserProfiles[nm.User()]).LoginName, "<missing-profile>")
 	}
-	b.currentNode().SetNetMap(nm)
+	discoChanged := b.currentNode().SetNetMap(nm)
 	if ms, ok := b.sys.MagicSock.GetOK(); ok {
 		if nm != nil {
 			if nm.Cached {
@@ -7301,6 +7304,12 @@ func (b *LocalBackend) setNetMapLocked(nm *netmap.NetworkMap) {
 			ms.SetNetworkMap(tailcfg.NodeView{}, nil)
 		}
 	}
+	// Reset the WireGuard session for peers whose disco key changed in
+	// a way that indicates a restart, flushing their dead session keys;
+	// each such peer is lazily re-created on demand with current state.
+	for _, k := range discoChanged {
+		b.e.ResetDevicePeer(k)
+	}
 	if login != b.activeLogin {
 		b.logf("active login: %v", login)
 		b.activeLogin = login
@@ -9174,6 +9183,11 @@ var (
 	metricNetmapDeltaPeerPatched  = clientmetric.NewCounter("localbackend_netmap_delta_peer_patched")
 	metricUpdatePacketFilter      = clientmetric.NewCounter("localbackend_update_packet_filter")
 	metricUpdateUserProfiles      = clientmetric.NewCounter("localbackend_update_user_profiles")
+
+	// metricTSMPLearnedKeyMismatch counts netmap updates carrying a peer
+	// disco key that doesn't match the one previously learned via TSMP
+	// for the same peer. See [nodeBackend.discoChangedLocked].
+	metricTSMPLearnedKeyMismatch = clientmetric.NewCounter("magicsock_tsmp_learned_key_mismatch")
 )
 
 func (b *LocalBackend) stateEncrypted() opt.Bool {

ipn/ipnlocal/node_backend.go

@@ -32,6 +32,7 @@ import (
 	"tailscale.com/util/dnsname"
 	"tailscale.com/util/eventbus"
 	"tailscale.com/util/mak"
+	"tailscale.com/util/set"
 	"tailscale.com/util/slicesx"
 	"tailscale.com/util/testenv"
 	"tailscale.com/wgengine/filter"
@@ -166,6 +167,14 @@ type nodeBackend struct {
 	// nil in production, where no test installs a waiter.
 	keyWaitersForTest map[key.NodePublic]chan struct{}
 
+	// tsmpLearnedDisco records, per node key, a peer disco key that was
+	// learned via TSMP (that is, over an existing WireGuard session with
+	// that peer). When a netmap update later reports the same disco key
+	// change, the peer's WireGuard session does not need to be reset,
+	// because the change demonstrably arrived over a working session.
+	// See [nodeBackend.discoChangedLocked].
+	tsmpLearnedDisco map[key.NodePublic]key.DiscoPublic
+
 	// routeMgr tracks this node's view of which IPs route to which
 	// peers and publishes lock-free snapshots for the data plane and
 	// the OS router. It is initialized once and immutable, but its
@@ -579,7 +588,7 @@ func (nb *nodeBackend) netMapWithPeers() *netmap.NetworkMap {
 	return nm
 }
 
-func (nb *nodeBackend) SetNetMap(nm *netmap.NetworkMap) {
+func (nb *nodeBackend) SetNetMap(nm *netmap.NetworkMap) (discoChanged []key.NodePublic) {
 	nb.mu.Lock()
 	defer nb.mu.Unlock()
 	nb.netMap = nm
@@ -587,7 +596,7 @@ func (nb *nodeBackend) SetNetMap(nm *netmap.NetworkMap) {
 	nb.updateNodeByKeyLocked()
 	nb.updateNodeByStableIDLocked()
 	nb.updateNodeByNameLocked()
-	nb.updatePeersLocked()
+	discoChanged = nb.updatePeersLocked()
 	nb.signalKeyWaitersForTestLocked()
 	if nm != nil {
 		nb.userProfiles = maps.Clone(nm.UserProfiles)
@@ -600,6 +609,7 @@ func (nb *nodeBackend) SetNetMap(nm *netmap.NetworkMap) {
 		nb.packetFilter = nil
 		nb.derpMapViewPub.Publish(tailcfg.DERPMapView{})
 	}
+	return discoChanged
 }
 
 // AwaitNodeKeyForTest returns a channel that is closed once a peer with the
@@ -780,10 +790,19 @@ func (nb *nodeBackend) ExtraDNSByName(hostname string) (_ netip.Addr, ok bool) {
 	return ip, ok
 }
 
-func (nb *nodeBackend) updatePeersLocked() {
+func (nb *nodeBackend) updatePeersLocked() (discoChanged []key.NodePublic) {
 	nm := nb.netMap
 	oldIDs := slices.Collect(maps.Keys(nb.peers))
 
+	// Snapshot the previous disco keys (by node key) before the peers
+	// map is repopulated, to detect restarted peers below.
+	var prevDisco map[key.NodePublic]key.DiscoPublic
+	for _, p := range nb.peers {
+		if !p.DiscoKey().IsZero() {
+			mak.Set(&prevDisco, p.Key(), p.DiscoKey())
+		}
+	}
+
 	if nm == nil {
 		nb.peers = nil
 	} else {
@@ -794,6 +813,18 @@ func (nb *nodeBackend) updatePeersLocked() {
 		})
 	}
 
+	for _, p := range nb.peers {
+		if prev, ok := prevDisco[p.Key()]; ok && nb.discoChangedLocked(p.Key(), prev, p.DiscoKey()) {
+			discoChanged = append(discoChanged, p.Key())
+		}
+	}
+	// Drop TSMP-learned disco keys for peers no longer in the netmap.
+	for k := range nb.tsmpLearnedDisco {
+		if _, ok := nb.nodeByKey[k]; !ok {
+			delete(nb.tsmpLearnedDisco, k)
+		}
+	}
+
 	// Resync the route manager to the new peer set. Upserts of
 	// unchanged peers are cheap no-ops; this is a full-netmap event,
 	// so O(n) work is expected here.
@@ -807,6 +838,52 @@ func (nb *nodeBackend) updatePeersLocked() {
 		rt.UpsertPeer(p)
 	}
 	rt.Commit()
+	return discoChanged
+}
+
+// recordTSMPLearnedDisco notes that a peer's new disco key was learned via
+// TSMP, so the netmap update carrying the same change need not reset the
+// peer's WireGuard session. See the [nodeBackend.tsmpLearnedDisco] field doc.
+func (nb *nodeBackend) recordTSMPLearnedDisco(pub key.NodePublic, disco key.DiscoPublic) {
+	nb.mu.Lock()
+	defer nb.mu.Unlock()
+	mak.Set(&nb.tsmpLearnedDisco, pub, disco)
+}
+
+// discoChangedLocked reports whether a peer's disco key change from prev to
+// cur should reset the peer's WireGuard session. A changed disco key means
+// the peer restarted, so any existing session key material is dead weight;
+// resetting lets the handshake start over immediately. The exception is a
+// key change already learned via TSMP: that arrived over a working WireGuard
+// session with the peer, so the session is demonstrably fine and is kept.
+//
+// It consumes any [nodeBackend.tsmpLearnedDisco] entry for pub.
+// nb.mu must be held.
+func (nb *nodeBackend) discoChangedLocked(pub key.NodePublic, prev, cur key.DiscoPublic) bool {
+	if prev.IsZero() || cur.IsZero() || prev == cur {
+		return false
+	}
+	if discoTSMP, ok := nb.tsmpLearnedDisco[pub]; ok {
+		delete(nb.tsmpLearnedDisco, pub)
+		if discoTSMP == cur {
+			nb.logf("nodeBackend: skipping WireGuard session reset (TSMP key): %s changed from %q to %q",
+				pub.ShortString(), prev, cur)
+			return false
+		}
+		// The new disco key does not match what we received via
+		// TSMP for this peer. This is unexpected, though possible
+		// if processing a change in a large netmap ends up taking
+		// longer than the 2 second timeout in
+		// [controlclient.mapRoutineState.UpdateNetmapDelta], or if
+		// the context is cancelled mid update. Log the event, and reset
+		// the session as it is possibly a stale entry in the map
+		// instead of a TSMP disco key update that led us here.
+		nb.logf("nodeBackend: [unexpected] using TSMP key for %s (control stale): tsmp=%q control=%q old=%q",
+			pub.ShortString(), discoTSMP, cur, prev)
+		metricTSMPLearnedKeyMismatch.Add(1)
+	}
+	nb.logf("nodeBackend: peer %s disco key changed from %q to %q", pub.ShortString(), prev, cur)
+	return true
 }
 
 // routePrefs is the subset of routing-relevant prefs (and derived
@@ -901,16 +978,30 @@ func (nb *nodeBackend) mergeUserProfiles(profiles map[tailcfg.UserID]tailcfg.Use
 	}
 }
 
+// netmapDeltaResult describes the side effects of applying netmap
+// delta mutations that the caller must propagate.
+type netmapDeltaResult struct {
+	// ChangedAllowedIPs are the peers whose allowed source prefixes
+	// changed, as described by [routemanager.Result.AllowedIPs]; the
+	// caller syncs those peers to the WireGuard device. In
+	// particular, the value for a key will be nil when that peer was
+	// removed.
+	ChangedAllowedIPs map[key.NodePublic][]netip.Prefix
+
+	// DiscoChanged is the set of peers whose disco key changed in a
+	// way that requires a WireGuard session reset (see
+	// [nodeBackend.discoChangedLocked]).
+	DiscoChanged set.Set[key.NodePublic]
+}
+
 // UpdateNetmapDelta applies the given netmap mutations to the live
-// peer state. It returns the peers whose allowed source prefixes
-// changed (as described by [routemanager.Result.AllowedIPs]) so the
-// caller can sync those peers to the WireGuard device, and reports
-// whether it handled all of the mutations.
-func (nb *nodeBackend) UpdateNetmapDelta(muts []netmap.NodeMutation) (changedAllowedIPs map[key.NodePublic][]netip.Prefix, handled bool) {
+// peer state. It returns the side effects for the caller to
+// propagate, and reports whether it handled all of the mutations.
+func (nb *nodeBackend) UpdateNetmapDelta(muts []netmap.NodeMutation) (res netmapDeltaResult, handled bool) {
 	nb.mu.Lock()
 	defer nb.mu.Unlock()
 	if nb.netMap == nil {
-		return nil, false
+		return res, false
 	}
 
 	// Locally cloned mutable nodes, to avoid calling AsStruct (clone)
@@ -925,14 +1016,23 @@ func (nb *nodeBackend) UpdateNetmapDelta(muts []netmap.NodeMutation) (changedAll
 	// get a full netmap soon and reset all our state anyway.
 	rt := nb.routeMgr.Begin()
 	defer func() {
-		res := rt.Commit()
-		changedAllowedIPs = res.AllowedIPs
+		res.ChangedAllowedIPs = rt.Commit().AllowedIPs
 	}()
 
 	for _, m := range muts {
 		switch m := m.(type) {
 		case netmap.NodeMutationUpsert:
 			nid := m.Node.ID()
+			if old, ok := nb.peers[nid]; ok {
+				if old.Key() == m.Node.Key() {
+					if nb.discoChangedLocked(m.Node.Key(), old.DiscoKey(), m.Node.DiscoKey()) {
+						res.DiscoChanged.Make()
+						res.DiscoChanged.Add(m.Node.Key())
+					}
+				} else {
+					delete(nb.tsmpLearnedDisco, old.Key())
+				}
+			}
 			mak.Set(&nb.peers, nid, m.Node)
 			for _, ipp := range m.Node.Addresses().All() {
 				if ipp.IsSingleIP() {
@@ -956,6 +1056,7 @@ func (nb *nodeBackend) UpdateNetmapDelta(muts []netmap.NodeMutation) (changedAll
 				delete(nb.nodeByKey, old.Key())
 				delete(nb.nodeByWGString, old.Key().WireGuardGoString())
 				delete(nb.nodeByStableID, old.StableID())
+				delete(nb.tsmpLearnedDisco, old.Key())
 				nb.removeNodeNameLocked(old.Name())
 				delete(nb.peers, nid)
 				rt.RemovePeer(nid)
@@ -969,7 +1070,7 @@ func (nb *nodeBackend) UpdateNetmapDelta(muts []netmap.NodeMutation) (changedAll
 			nv, ok := nb.peers[nid]
 			if !ok {
 				// TODO(bradfitz): unexpected metric?
-				return nil, false
+				return res, false
 			}
 			n = nv.AsStruct()
 			mak.Set(&mutableNodes, nv.ID(), n)
@@ -980,7 +1081,7 @@ func (nb *nodeBackend) UpdateNetmapDelta(muts []netmap.NodeMutation) (changedAll
 		nb.peers[nid] = n.View()
 	}
 	nb.signalKeyWaitersForTestLocked()
-	return nil, true
+	return res, true
 }
 
 // unlockedNodesPermitted reports whether any peer with theUnsignedPeerAPIOnly bool set true has any of its allowed IPs

ipn/ipnlocal/node_backend_test.go

@@ -345,17 +345,17 @@ func TestNodeBackendRouteManager(t *testing.T) {
 
 	// Incremental deltas: add peer 3, remove peer 1.
 	p3 := mkPeer(3, "stable3", "100.64.0.3/32")
-	changed, handled := nb.UpdateNetmapDelta([]netmap.NodeMutation{
+	deltaRes, handled := nb.UpdateNetmapDelta([]netmap.NodeMutation{
 		netmap.NodeMutationUpsert{Node: p3},
 		netmap.MakeNodeMutationRemove(1),
 	})
 	if !handled {
 		t.Fatal("UpdateNetmapDelta not handled")
 	}
-	if len(changed) != 2 || changed[p3.Key()] == nil {
+	if changed := deltaRes.ChangedAllowedIPs; len(changed) != 2 || changed[p3.Key()] == nil {
 		t.Errorf("UpdateNetmapDelta changed = %v; want entries for %v and %v", changed, p3.Key(), p1.Key())
 	}
-	if v, ok := changed[p1.Key()]; !ok || v != nil {
+	if v, ok := deltaRes.ChangedAllowedIPs[p1.Key()]; !ok || v != nil {
 		t.Errorf("UpdateNetmapDelta changed[%v] = %v, %v; want nil, true for removed peer", p1.Key(), v, ok)
 	}
 	wantPeerFor("100.64.0.3", p3)
@@ -366,3 +366,142 @@ func TestNodeBackendRouteManager(t *testing.T) {
 	wantPeerFor("100.64.0.3", tailcfg.NodeView{})
 	wantPeerFor("100.64.0.2", p2)
 }
+
+// TestNodeBackendDiscoChanged exercises the full-netmap disco change
+// detection: a peer whose disco key changes has restarted and needs its
+// WireGuard session reset, unless the new key was already learned over
+// TSMP (that is, over a working WireGuard session with the peer).
+func TestNodeBackendDiscoChanged(t *testing.T) {
+	nb := newNodeBackend(t.Context(), tstest.WhileTestRunningLogger(t), eventbus.New())
+
+	nk := key.NewNode().Public()
+	mkNetMap := func(disco key.DiscoPublic) *netmap.NetworkMap {
+		n := &tailcfg.Node{
+			ID:       1,
+			Key:      nk,
+			DiscoKey: disco,
+			HomeDERP: 1,
+		}
+		return &netmap.NetworkMap{Peers: []tailcfg.NodeView{n.View()}}
+	}
+	newDisco := func() key.DiscoPublic { return key.NewDisco().Public() }
+
+	// A brand-new peer is not a disco change.
+	d1 := newDisco()
+	if got := nb.SetNetMap(mkNetMap(d1)); len(got) != 0 {
+		t.Errorf("SetNetMap(new peer) discoChanged = %v; want none", got)
+	}
+
+	// A changed disco key requires a session reset.
+	d2 := newDisco()
+	if got := nb.SetNetMap(mkNetMap(d2)); !slices.Contains(got, nk) {
+		t.Errorf("SetNetMap(changed disco) discoChanged = %v; want %v", got, nk)
+	}
+
+	// An unchanged disco key does not.
+	if got := nb.SetNetMap(mkNetMap(d2)); len(got) != 0 {
+		t.Errorf("SetNetMap(same disco) discoChanged = %v; want none", got)
+	}
+
+	// A change already learned via TSMP is suppressed...
+	d3 := newDisco()
+	nb.recordTSMPLearnedDisco(nk, d3)
+	if got := nb.SetNetMap(mkNetMap(d3)); len(got) != 0 {
+		t.Errorf("SetNetMap(TSMP-learned disco) discoChanged = %v; want none", got)
+	}
+
+	// ...but the TSMP entry is consumed, so the next change resets again.
+	d4 := newDisco()
+	if got := nb.SetNetMap(mkNetMap(d4)); !slices.Contains(got, nk) {
+		t.Errorf("SetNetMap(after TSMP entry consumed) discoChanged = %v; want %v", got, nk)
+	}
+
+	// A TSMP-learned key that doesn't match the netmap's new key still
+	// resets the session and bumps the mismatch metric.
+	before := metricTSMPLearnedKeyMismatch.Value()
+	nb.recordTSMPLearnedDisco(nk, newDisco())
+	d5 := newDisco()
+	if got := nb.SetNetMap(mkNetMap(d5)); !slices.Contains(got, nk) {
+		t.Errorf("SetNetMap(TSMP mismatch) discoChanged = %v; want %v", got, nk)
+	}
+	if delta := metricTSMPLearnedKeyMismatch.Value() - before; delta != 1 {
+		t.Errorf("metricTSMPLearnedKeyMismatch delta = %d; want 1", delta)
+	}
+
+	// Removing the peer garbage-collects its TSMP entry: after the peer
+	// comes back, a change to the once-recorded key is a normal reset.
+	d6 := newDisco()
+	nb.recordTSMPLearnedDisco(nk, d6)
+	nb.SetNetMap(&netmap.NetworkMap{})
+	nb.SetNetMap(mkNetMap(d5))
+	if got := nb.SetNetMap(mkNetMap(d6)); !slices.Contains(got, nk) {
+		t.Errorf("SetNetMap(after TSMP entry GC) discoChanged = %v; want %v", got, nk)
+	}
+
+	// Transitions to or from a zero disco key never reset.
+	if got := nb.SetNetMap(mkNetMap(key.DiscoPublic{})); len(got) != 0 {
+		t.Errorf("SetNetMap(to zero disco) discoChanged = %v; want none", got)
+	}
+	if got := nb.SetNetMap(mkNetMap(d1)); len(got) != 0 {
+		t.Errorf("SetNetMap(from zero disco) discoChanged = %v; want none", got)
+	}
+}
+
+// TestNodeBackendDiscoChangedDelta is like TestNodeBackendDiscoChanged
+// but for the incremental path: disco changes arriving as
+// [netmap.NodeMutationUpsert] deltas.
+func TestNodeBackendDiscoChangedDelta(t *testing.T) {
+	nb := newNodeBackend(t.Context(), tstest.WhileTestRunningLogger(t), eventbus.New())
+
+	mkNode := func(k key.NodePublic, disco key.DiscoPublic) tailcfg.NodeView {
+		return (&tailcfg.Node{ID: 1, Key: k, DiscoKey: disco, HomeDERP: 1}).View()
+	}
+	newDisco := func() key.DiscoPublic { return key.NewDisco().Public() }
+	apply := func(muts ...netmap.NodeMutation) set.Set[key.NodePublic] {
+		t.Helper()
+		deltaRes, handled := nb.UpdateNetmapDelta(muts)
+		if !handled {
+			t.Fatal("UpdateNetmapDelta not handled")
+		}
+		return deltaRes.DiscoChanged
+	}
+
+	nk := key.NewNode().Public()
+	d1 := newDisco()
+	nb.SetNetMap(&netmap.NetworkMap{Peers: []tailcfg.NodeView{mkNode(nk, d1)}})
+
+	// An upserted peer with a changed disco key needs a session reset.
+	d2 := newDisco()
+	if got := apply(netmap.NodeMutationUpsert{Node: mkNode(nk, d2)}); !got.Contains(nk) {
+		t.Errorf("upsert(changed disco) discoChanged = %v; want %v", got, nk)
+	}
+
+	// An unchanged disco key does not.
+	if got := apply(netmap.NodeMutationUpsert{Node: mkNode(nk, d2)}); len(got) != 0 {
+		t.Errorf("upsert(same disco) discoChanged = %v; want none", got)
+	}
+
+	// A change already learned via TSMP is suppressed.
+	d3 := newDisco()
+	nb.recordTSMPLearnedDisco(nk, d3)
+	if got := apply(netmap.NodeMutationUpsert{Node: mkNode(nk, d3)}); len(got) != 0 {
+		t.Errorf("upsert(TSMP-learned disco) discoChanged = %v; want none", got)
+	}
+
+	// A node key rotation replaces the WireGuard peer outright, so no
+	// disco-based reset is reported.
+	nk2 := key.NewNode().Public()
+	if got := apply(netmap.NodeMutationUpsert{Node: mkNode(nk2, newDisco())}); len(got) != 0 {
+		t.Errorf("upsert(rotated node key) discoChanged = %v; want none", got)
+	}
+
+	// Removing the peer garbage-collects its TSMP entry: after the peer
+	// comes back, a change to the once-recorded key is a normal reset.
+	d4 := newDisco()
+	nb.recordTSMPLearnedDisco(nk2, d4)
+	apply(netmap.MakeNodeMutationRemove(1))
+	apply(netmap.NodeMutationUpsert{Node: mkNode(nk2, newDisco())})
+	if got := apply(netmap.NodeMutationUpsert{Node: mkNode(nk2, d4)}); !got.Contains(nk2) {
+		t.Errorf("upsert(after TSMP entry GC) discoChanged = %v; want %v", got, nk2)
+	}
+}

ipn/ipnlocal/state_test.go

@@ -2014,7 +2014,8 @@ func (e *mockEngine) SetPeerByIPPacketFunc(func(netip.Addr) (_ key.NodePublic, o
 func (e *mockEngine) SetPeerForIPFunc(func(netip.Addr) (_ wgengine.PeerForIP, ok bool))  {}
 func (e *mockEngine) SetPeerConfigFunc(func(key.NodePublic) (allowedIPs []netip.Prefix, ok bool)) {
 }
-func (e *mockEngine) SyncDevicePeer(key.NodePublic) {}
+func (e *mockEngine) SyncDevicePeer(key.NodePublic)  {}
+func (e *mockEngine) ResetDevicePeer(key.NodePublic) {}
 func (e *mockEngine) SetPeerSessionStateFunc(func(key.NodePublic, wgengine.PeerWireGuardState)) {
 }
 func (e *mockEngine) SetNetLogSource(wgengine.NetLogSource)                            {}

wgengine/userspace.go

@@ -50,7 +50,6 @@ import (
 	"tailscale.com/util/eventbus"
 	"tailscale.com/util/execqueue"
 	"tailscale.com/util/mak"
-	"tailscale.com/util/set"
 	"tailscale.com/util/singleflight"
 	"tailscale.com/util/testenv"
 	"tailscale.com/util/usermetric"
@@ -89,8 +88,7 @@ type userspaceEngine struct {
 	birdClient     BIRDClient          // or nil
 	controlKnobs   *controlknobs.Knobs // or nil
 
-	testMaybeReconfigHook func()                        // for tests; if non-nil, fires if maybeReconfigWireguardLocked called
-	testDiscoChangedHook  func(map[key.NodePublic]bool) // for tests; if non-nil, fires after assembling discoChanged map
+	testMaybeReconfigHook func() // for tests; if non-nil, fires if maybeReconfigWireguardLocked called
 
 	// isLocalAddr reports the whether an IP is assigned to the local
 	// tunnel interface. It's used to reflect local packets
@@ -161,10 +159,6 @@ type userspaceEngine struct {
 	// rewritten.
 	wgPeerLookup syncs.AtomicValue[func(wgString string) (tsString string, ok bool)]
 
-	// tsmpLearnedDisco tracks per node key if a peer disco key was learned via TSMP.
-	// wgLock must be held when using this map.
-	tsmpLearnedDisco map[key.NodePublic]key.DiscoPublic
-
 	// Lock ordering: magicsock.Conn.mu, wgLock, then mu.
 }
 
@@ -766,6 +760,13 @@ func (e *userspaceEngine) SyncDevicePeer(k key.NodePublic) {
 	}
 }
 
+// ResetDevicePeer implements [Engine.ResetDevicePeer].
+func (e *userspaceEngine) ResetDevicePeer(k key.NodePublic) {
+	e.wgLock.Lock()
+	defer e.wgLock.Unlock()
+	e.wgdev.RemovePeer(k.Raw32())
+}
+
 // SetPeerByIPPacketFunc installs a callback used by wireguard-go to look up
 // which peer should handle an outbound packet by destination IP.
 //
@@ -851,12 +852,6 @@ func (e *userspaceEngine) ResetAndStop() (*Status, error) {
 	return e.getStatus()
 }
 
-func (e *userspaceEngine) PatchDiscoKey(pub key.NodePublic, disco key.DiscoPublic) {
-	e.wgLock.Lock()
-	defer e.wgLock.Unlock()
-	mak.Set(&e.tsmpLearnedDisco, pub, disco)
-}
-
 func (e *userspaceEngine) Reconfig(cfg *wgcfg.Config, routerCfg *router.Config, dnsCfg *dns.Config) error {
 	if routerCfg == nil {
 		panic("routerCfg must not be nil")
@@ -871,11 +866,6 @@ func (e *userspaceEngine) Reconfig(cfg *wgcfg.Config, routerCfg *router.Config,
 	defer e.wgLock.Unlock()
 	e.tundev.SetWGConfig(cfg)
 
-	peerSet := make(set.Set[key.NodePublic], len(cfg.Peers))
-	for _, p := range cfg.Peers {
-		peerSet.Add(p.PublicKey)
-	}
-
 	e.mu.Lock()
 	self := e.selfNode
 	e.mu.Unlock()
@@ -928,63 +918,6 @@ func (e *userspaceEngine) Reconfig(cfg *wgcfg.Config, routerCfg *router.Config,
 		e.isDNSIPOverTailscale.Store(ipset.NewContainsIPFunc(views.SliceOf(dnsIPsOverTailscale(dnsCfg, routerCfg))))
 	}
 
-	// See if any peers have changed disco keys, which means they've restarted.
-	// If so, remove the peer from wireguard-go to flush its session key,
-	// then let the PeerLookupFunc re-create it on demand.
-	discoChanged := make(map[key.NodePublic]bool)
-	if engineChanged {
-		prevEP := make(map[key.NodePublic]key.DiscoPublic)
-		for i := range e.lastCfgFull.Peers {
-			if p := &e.lastCfgFull.Peers[i]; !p.DiscoKey.IsZero() {
-				prevEP[p.PublicKey] = p.DiscoKey
-			}
-		}
-		for i := range cfg.Peers {
-			p := &cfg.Peers[i]
-			if p.DiscoKey.IsZero() {
-				continue
-			}
-
-			pub := p.PublicKey
-
-			if old, ok := prevEP[pub]; ok && old != p.DiscoKey {
-				// If the disco key was learned via TSMP, we do not need to reset the
-				// wireguard config as the new key was received over an existing wireguard
-				// connection.
-				if discoTSMP, okTSMP := e.tsmpLearnedDisco[p.PublicKey]; okTSMP {
-					// Key matches, remove entry from map.
-					delete(e.tsmpLearnedDisco, p.PublicKey)
-					if discoTSMP == p.DiscoKey {
-						e.logf("wgengine: Skipping reconfig (TSMP key): %s changed from %q to %q",
-							pub.ShortString(), old, p.DiscoKey)
-						// Skip session clear.
-						continue
-					}
-
-					// The new disco key does not match what we received via
-					// TSMP for this peer. This is unexpected, though possible
-					// if processing a change in a large netmap ends up taking
-					// longer than the 2 second timeout in
-					// [controlClient.mapRoutineState.UpdateNetmapDelta], or if
-					// the context is cancelled mid update. Log the event, and reset
-					// the connection as it is possibly a stale entry in the map
-					// instead of a TSMP disco key update that led us here.
-					e.logf("wgengine: [unexpected] Reconfig: using TSMP key for %s (control stale): tsmp=%q control=%q old=%q",
-						pub.ShortString(), discoTSMP, p.DiscoKey, old)
-					metricTSMPLearnedKeyMismatch.Add(1)
-				}
-
-				discoChanged[pub] = true
-				e.logf("wgengine: Reconfig: %s changed from %q to %q", pub.ShortString(), old, p.DiscoKey)
-			}
-		}
-	}
-
-	// For tests, what disco connections needs to be changed.
-	if e.testDiscoChangedHook != nil {
-		e.testDiscoChangedHook(discoChanged)
-	}
-
 	if !e.lastCfgFull.PrivateKey.Equal(cfg.PrivateKey) {
 		// Tell magicsock about the new (or initial) private key
 		// (which is needed by DERP) before wgdev gets it, as wgdev
@@ -1008,19 +941,6 @@ func (e *userspaceEngine) Reconfig(cfg *wgcfg.Config, routerCfg *router.Config,
 		if err := e.maybeReconfigWireguardLocked(); err != nil {
 			return err
 		}
-		// Now that we've reconfigured wireguard-go, remove any peers with
-		// changed disco keys to flush their session keys, and let them be
-		// re-created on demand by the PeerLookupFunc.
-		for pub := range discoChanged {
-			e.wgdev.RemovePeer(pub.Raw32())
-		}
-	}
-
-	// Cleanup map of tsmp marks for peers that no longer exists in config.
-	for nodeKey := range e.tsmpLearnedDisco {
-		if !peerSet.Contains(nodeKey) {
-			delete(e.tsmpLearnedDisco, nodeKey)
-		}
 	}
 
 	if routerChanged {
@@ -1692,8 +1612,6 @@ var (
 
 	metricTSMPDiscoKeyAdvertisementSent  = clientmetric.NewCounter("magicsock_tsmp_disco_key_advertisement_sent")
 	metricTSMPDiscoKeyAdvertisementError = clientmetric.NewCounter("magicsock_tsmp_disco_key_advertisement_error")
-
-	metricTSMPLearnedKeyMismatch = clientmetric.NewCounter("magicsock_tsmp_learned_key_mismatch")
 )
 
 func (e *userspaceEngine) InstallCaptureHook(cb packet.CaptureCallback) {

wgengine/userspace_test.go

@@ -125,167 +125,6 @@ func TestUserspaceEngineReconfig(t *testing.T) {
 	}
 }
 
-func TestUserspaceEngineTSMPLearned(t *testing.T) {
-	bus := eventbustest.NewBus(t)
-
-	ht := health.NewTracker(bus)
-	reg := new(usermetric.Registry)
-	e, err := NewFakeUserspaceEngine(t.Logf, 0, ht, reg, bus)
-	if err != nil {
-		t.Fatal(err)
-	}
-	t.Cleanup(e.Close)
-	ue := e.(*userspaceEngine)
-
-	discoChangedChan := make(chan map[key.NodePublic]bool, 1)
-	ue.testDiscoChangedHook = func(m map[key.NodePublic]bool) {
-		discoChangedChan <- m
-	}
-
-	routerCfg := &router.Config{}
-
-	keyChanges := []struct {
-		tsmp  bool
-		inMap bool
-	}{
-		{tsmp: false, inMap: false},
-		{tsmp: true, inMap: false},
-		{tsmp: false, inMap: true},
-	}
-
-	nkHex := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
-	for _, change := range keyChanges {
-		oldDisco := key.NewDisco()
-		nm := &netmap.NetworkMap{
-			Peers: nodeViews([]*tailcfg.Node{
-				{
-					ID:       1,
-					Key:      nkFromHex(nkHex),
-					DiscoKey: oldDisco.Public(),
-				},
-			}),
-		}
-		nk, err := key.ParseNodePublicUntyped(mem.S(nkHex))
-		if err != nil {
-			t.Fatal(err)
-		}
-		e.SetSelfNode(nm.SelfNode)
-
-		newDisco := key.NewDisco()
-		cfg := &wgcfg.Config{
-			Peers: []wgcfg.Peer{
-				{
-					PublicKey: nk,
-					DiscoKey:  newDisco.Public(),
-				},
-			},
-		}
-
-		if change.tsmp {
-			ue.PatchDiscoKey(nk, newDisco.Public())
-		}
-		err = e.Reconfig(cfg, routerCfg, &dns.Config{})
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		changeMap := <-discoChangedChan
-
-		if _, ok := changeMap[nk]; ok != change.inMap {
-			t.Fatalf("expect key %v in map %v to be %t, got %t", nk, changeMap,
-				change.inMap, ok)
-		}
-	}
-}
-
-func TestUserspaceEngineTSMPLearnedMismatch(t *testing.T) {
-	bus := eventbustest.NewBus(t)
-
-	ht := health.NewTracker(bus)
-	reg := new(usermetric.Registry)
-	e, err := NewFakeUserspaceEngine(t.Logf, 0, ht, reg, bus)
-	if err != nil {
-		t.Fatal(err)
-	}
-	t.Cleanup(e.Close)
-	ue := e.(*userspaceEngine)
-
-	discoChangedChan := make(chan map[key.NodePublic]bool, 1)
-	ue.testDiscoChangedHook = func(m map[key.NodePublic]bool) {
-		discoChangedChan <- m
-	}
-
-	routerCfg := &router.Config{}
-	var metricValue int64 = 0
-
-	keyChanges := []struct {
-		tsmp     bool
-		inMap    bool
-		wrongKey bool
-	}{
-		{tsmp: false, inMap: false, wrongKey: false},
-		{tsmp: true, inMap: false, wrongKey: false},
-		{tsmp: true, inMap: true, wrongKey: true},
-		{tsmp: false, inMap: true, wrongKey: false},
-	}
-
-	nkHex := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
-	for _, change := range keyChanges {
-		oldDisco := key.NewDisco()
-		nm := &netmap.NetworkMap{
-			Peers: nodeViews([]*tailcfg.Node{
-				{
-					ID:       1,
-					Key:      nkFromHex(nkHex),
-					DiscoKey: oldDisco.Public(),
-				},
-			}),
-		}
-		nk, err := key.ParseNodePublicUntyped(mem.S(nkHex))
-		if err != nil {
-			t.Fatal(err)
-		}
-		e.SetSelfNode(nm.SelfNode)
-
-		newDisco := key.NewDisco()
-		cfg := &wgcfg.Config{
-			Peers: []wgcfg.Peer{
-				{
-					PublicKey: nk,
-					DiscoKey:  newDisco.Public(),
-				},
-			},
-		}
-
-		tsmpKey := newDisco.Public()
-		if change.tsmp {
-			if change.wrongKey {
-				tsmpKey = key.NewDisco().Public()
-			}
-			ue.PatchDiscoKey(nk, tsmpKey)
-		}
-		err = e.Reconfig(cfg, routerCfg, &dns.Config{})
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		changeMap := <-discoChangedChan
-
-		if _, ok := changeMap[nk]; ok != change.inMap {
-			t.Fatalf("expect key %v in map %v to be %t, got %t", nk, changeMap,
-				change.inMap, ok)
-		}
-
-		metric := metricTSMPLearnedKeyMismatch.Value()
-		delta := metric - metricValue
-		metricValue = metric
-
-		if change.wrongKey && delta != 1 {
-			t.Fatalf("expected a delta of 1, got %d", delta)
-		}
-	}
-}
-
 func TestUserspaceEnginePortReconfig(t *testing.T) {
 	flakytest.Mark(t, "https://github.com/tailscale/tailscale/issues/2855")
 	const defaultPort = 49983

wgengine/wgengine.go

@@ -198,6 +198,16 @@ type Engine interface {
 	// It is a no-op if no config source is installed.
 	SyncDevicePeer(key.NodePublic)
 
+	// ResetDevicePeer removes the peer from the WireGuard device,
+	// discarding any session key material and in-flight handshake
+	// state. If the peer is still known to the config source installed
+	// via [Engine.SetPeerConfigFunc], it is lazily re-created on demand
+	// with fresh state.
+	//
+	// LocalBackend calls it when a peer's disco key changes, which
+	// means the peer restarted and its old sessions are dead.
+	ResetDevicePeer(key.NodePublic)
+
 	// SetNetLogSource installs the [NetLogSource] consulted by the
 	// engine's network flow logger for node lookups and the current
 	// audit logging identity.