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.