wgengine/magicsock: add facilities to support multiple disco keys (#20927)

Instead of accessing the key and string directly, put it behind a method
to make it easier to use that as a proxy when we add multiple keys
(control and TSMP origin).

Instead of only updating a single origin for disco keys, teach peermap
how to work with multiple disco key sources so a key can originate from
either control or TSMP.

This sets the work up for switching dynamically between disco keys
later. As of this commit, the keys still arrive for the most part via
the controlClient, but this commit sets us up for:
 - Switching dynamically between received keys.
 - Route TSMP keys directly into magicsock, not via controlClient.
 - Potentially revert controlClient to the state before any TSMP
   changes, making it single threaded.
 - Use switching keys as trigger for optimistic WG handshakes, instead
   of tearing down the full connection and setting it up again.

Updates #20494

Signed-off-by: Claus Lensbøl <claus@tailscale.com>

wgengine/magicsock/endpoint.go

@@ -371,11 +371,56 @@ func (de *endpoint) setProbeUDPLifetimeConfigLocked(desired *ProbeUDPLifetimeCon
 	p.resetCycleEndpointLocked()
 }
 
-// endpointDisco is the current disco key and short string for an endpoint. This
-// structure is immutable.
+// endpointDisco is the current disco key and short string for an endpoint for
+// keys learned both from controlClient and via TSMP. Only one key is active at
+// a time for sending. Currently the controlClient learned key is always
+// considered active as TSMP writes keys through this route.
+//
+// This structure is immutable.
 type endpointDisco struct {
-	key   key.DiscoPublic // for discovery messages.
-	short string          // ShortString of discoKey.
+	controlKey   key.DiscoPublic // key learned via control for disco messages.
+	tsmpKey      key.DiscoPublic // key learned via TSMP for disco messages.
+	controlShort string          // ShortString of control learned key.
+	tsmpShort    string          // ShortString of TSMP learned key.
+	tsmpActive   bool
+}
+
+// key returns the disco key currently regarded as active or a zero key if
+// endpointDisco is nil.
+func (e *endpointDisco) key() key.DiscoPublic {
+	if e == nil {
+		return key.DiscoPublic{}
+	}
+	if e.tsmpActive {
+		return e.tsmpKey
+	}
+	return e.controlKey
+}
+
+func (e *endpointDisco) keyFromControl() key.DiscoPublic {
+	if e == nil {
+		return key.DiscoPublic{}
+	}
+	return e.controlKey
+}
+
+func (e *endpointDisco) keyFromTSMP() key.DiscoPublic {
+	if e == nil {
+		return key.DiscoPublic{}
+	}
+	return e.tsmpKey
+}
+
+// shortString returns the ShortString of the key currently regarded as active
+// or an empty string if endpointDisco is nil.
+func (e *endpointDisco) shortString() string {
+	if e == nil {
+		return ""
+	}
+	if e.tsmpActive {
+		return e.tsmpShort
+	}
+	return e.controlShort
 }
 
 type sentPing struct {
@@ -544,11 +589,7 @@ func (de *endpoint) noteRecvActivity(src epAddr, now mono.Time) bool {
 }
 
 func (de *endpoint) discoShort() string {
-	var short string
-	if d := de.disco.Load(); d != nil {
-		short = d.short
-	}
-	return short
+	return de.disco.Load().shortString()
 }
 
 // String exists purely so wireguard-go internals can log.Printf("%v")
@@ -706,7 +747,7 @@ func (de *endpoint) maybeProbeUDPLifetimeLocked() (afterInactivityFor time.Durat
 	// shuffling probing probability where the local node ends up with a large
 	// key value lexicographically relative to the other nodes it tends to
 	// communicate with. If de's disco key changes, the cycle will reset.
-	if de.c.discoAtomic.Public().Compare(epDisco.key) >= 0 {
+	if de.c.discoAtomic.Public().Compare(epDisco.key()) >= 0 {
 		// lower disco pub key node probes higher
 		return afterInactivityFor, false
 	}
@@ -1351,9 +1392,8 @@ func (de *endpoint) startDiscoPingLocked(ep epAddr, now mono.Time, purpose disco
 		if purpose == pingHeartbeatForUDPLifetime && de.probeUDPLifetime != nil {
 			de.probeUDPLifetime.lastTxID = txid
 		}
-		go de.sendDiscoPing(ep, epDisco.key, txid, s, logLevel)
+		go de.sendDiscoPing(ep, epDisco.key(), txid, s, logLevel)
 	}
-
 }
 
 // sendDiscoPingsLocked starts pinging all of ep's direct endpoints.
@@ -1476,21 +1516,73 @@ func (de *endpoint) setLastPing(ipp netip.AddrPort, now mono.Time) {
 	state.lastPing = now
 }
 
-// updateDiscoKey replaces the disco key for de. If the key is a zero value key,
-// set the key to nil.
+// updateDiscoKey replaces the controlClient learned disco key for de.
+// Update only the control-provided key, leaving any existing TSMP key as-is.
+// If the new control key is zero, switch back to using the TSMP key if it
+// exists; otherwise mark the new control key as preferred.
+// Should both keys be zero, nil out the saved key.
+// The loop here ensures another update did not occur during the interval
+// between load and store (based on pointer identity).
 func (de *endpoint) updateDiscoKey(key key.DiscoPublic) {
-	if key.IsZero() {
-		de.disco.Store(nil)
-	} else {
-		de.disco.Store(&endpointDisco{
-			key:   key,
-			short: key.ShortString(),
-		})
+	epDisco := &endpointDisco{}
+	for {
+		old := de.disco.Load()
+		if old != nil {
+			epDisco.tsmpKey = old.tsmpKey
+			epDisco.tsmpShort = old.tsmpShort
+			epDisco.tsmpActive = key.IsZero()
+		}
+		if !key.IsZero() {
+			epDisco.controlKey = key
+			epDisco.controlShort = key.ShortString()
+			epDisco.tsmpActive = false
+		}
+		// We have no key material, nil out key.
+		if epDisco.controlKey.IsZero() && epDisco.tsmpKey.IsZero() {
+			epDisco = nil
+		}
+		if de.disco.CompareAndSwap(old, epDisco) {
+			return
+		}
+	}
+}
+
+// updateTSMPDiscoKey replaces the TSMP learned disco key for de.
+// Update only the TSMP-provided key, leaving any existing control key as-is.
+// If the new TSMP key is zero, switch back to using the control key if it
+// exists; otherwise mark the new TSMP key as preferred.
+// Should both keys be zero, nil out the saved key.
+// The loop here ensures another update did not occur during the interval
+// between load and store (based on pointer identity).
+func (de *endpoint) updateTSMPDiscoKey(key key.DiscoPublic) {
+	epDisco := &endpointDisco{}
+	for {
+		old := de.disco.Load()
+		if old != nil {
+			epDisco.controlKey = old.controlKey
+			epDisco.controlShort = old.controlShort
+			epDisco.tsmpActive = !key.IsZero()
+		}
+		if !key.IsZero() {
+			epDisco.tsmpKey = key
+			epDisco.tsmpShort = key.ShortString()
+			epDisco.tsmpActive = true
+		}
+		// We have no key material, nil out key.
+		if epDisco.controlKey.IsZero() && epDisco.tsmpKey.IsZero() {
+			epDisco = nil
+		}
+		if de.disco.CompareAndSwap(old, epDisco) {
+			// Fall out of the loop if the swap was successful (no change was made
+			// since the load).
+			break
+		}
 	}
 }
 
 // updateFromNode updates the endpoint based on a tailcfg.Node from a NetMap
-// update.
+// update. The node is assumed to originate from the control client from the
+// perspective of discoKey management.
 func (de *endpoint) updateFromNode(n tailcfg.NodeView, heartbeatDisabled bool, probeUDPLifetimeEnabled bool) {
 	if !n.Valid() {
 		panic("nil node when updating endpoint")
@@ -1506,11 +1598,7 @@ func (de *endpoint) updateFromNode(n tailcfg.NodeView, heartbeatDisabled bool, p
 	}
 	de.expired = n.Expired()
 
-	epDisco := de.disco.Load()
-	var discoKey key.DiscoPublic
-	if epDisco != nil {
-		discoKey = epDisco.key
-	}
+	discoKey := de.disco.Load().keyFromControl()
 
 	if discoKey != n.DiscoKey() {
 		de.c.logf("[v1] magicsock: disco: node %s changed from %s to %s", de.publicKey.ShortString(), discoKey, n.DiscoKey())

wgengine/magicsock/endpoint_test.go

@@ -309,10 +309,7 @@ func Test_endpoint_maybeProbeUDPLifetimeLocked(t *testing.T) {
 				bestAddr: tt.bestAddr,
 			}
 			if tt.remoteDisco != nil {
-				remote := &endpointDisco{
-					key: *tt.remoteDisco,
-				}
-				de.disco.Store(remote)
+				de.updateDiscoKey(*tt.remoteDisco)
 			}
 			p := tt.probeUDPLifetimeFn()
 			de.probeUDPLifetime = p
@@ -585,7 +582,7 @@ func Test_endpoint_sendDiscoPingsLocked_neverDirectUDP(t *testing.T) {
 				sentPing:      make(map[stun.TxID]sentPing),
 				endpointState: make(map[netip.AddrPort]*endpointState),
 			}
-			de.disco.Store(&endpointDisco{key: key.NewDisco().Public()})
+			de.updateDiscoKey(key.NewDisco().Public())
 			de.endpointState[directAddr] = &endpointState{}
 			de.sendDiscoPingsLocked(now, true)
 
@@ -650,7 +647,7 @@ func Test_endpoint_updateFromNodeAfterDiscoKeyChange(t *testing.T) {
 				debugUpdates:       ringlog.New[EndpointChange](10),
 			}
 			de.lastUDPRelayPathDiscovery = mono.Now()
-			de.disco.Store(&endpointDisco{key: oldKey, short: oldKey.ShortString()})
+			de.updateDiscoKey(oldKey)
 
 			incomingKey := oldKey
 			if tc.keyChanges {
@@ -800,3 +797,59 @@ func Test_endpoint_handlePongConnLocked(t *testing.T) {
 		})
 	}
 }
+
+// TestUpdateFromNodeUsesControlKeyForComparison verifies that updateFromNode
+// compares the netmap-provided disco key against the endpoint's control-learned
+// key (keyFromControl), not the currently active key (key()).
+// Some of this is scaffold for later changes.
+func TestUpdateFromNodeUsesControlKeyForComparison(t *testing.T) {
+	dk1 := key.NewDisco().Public() // initial control key
+	dk2 := key.NewDisco().Public() // TSMP key
+
+	tests := []struct {
+		name           string
+		netmapDiscoKey key.DiscoPublic
+		wantControlKey key.DiscoPublic
+		wantActiveKey  key.DiscoPublic
+		wantTsmpActive bool
+	}{
+		{
+			name:           "control_catches_up_to_tsmp_key",
+			netmapDiscoKey: dk2,
+			wantControlKey: dk2,
+			wantActiveKey:  dk2,
+			wantTsmpActive: false,
+		},
+		{
+			name:           "unchanged_netmap_key_preserves_active",
+			netmapDiscoKey: dk1,
+			wantControlKey: dk1,
+			wantActiveKey:  dk2,
+			wantTsmpActive: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			de := &endpoint{c: &Conn{logf: func(msg string, args ...any) {}}}
+			de.updateDiscoKey(dk1)
+			de.updateTSMPDiscoKey(dk2)
+
+			// Calls updateDiscoKey() internally and sets the endpoint `disco` field.
+			de.updateFromNode(
+				(&tailcfg.Node{Key: de.publicKey, DiscoKey: tt.netmapDiscoKey}).View(),
+				false, false)
+
+			epDisco := de.disco.Load()
+			if got := epDisco.keyFromControl(); got != tt.wantControlKey {
+				t.Errorf("keyFromControl: got %v, want %v", got, tt.wantControlKey)
+			}
+			if got := epDisco.tsmpActive; got != tt.wantTsmpActive {
+				t.Errorf("tsmpActive: got %t, want %t", got, tt.wantTsmpActive)
+			}
+			if got := epDisco.key(); got != tt.wantActiveKey {
+				t.Errorf("key(): got %v, want %v", got, tt.wantActiveKey)
+			}
+		})
+	}
+}

wgengine/magicsock/magicsock.go

@@ -629,14 +629,14 @@ func (c *Conn) onUDPRelayAllocResp(allocResp UDPRelayAllocResp) {
 	if disco == nil {
 		return
 	}
-	if disco.key.Compare(allocResp.ReqRxFromDiscoKey) != 0 {
+	if disco.key().Compare(allocResp.ReqRxFromDiscoKey) != 0 {
 		return
 	}
 	ep.mu.Lock()
 	defer ep.mu.Unlock()
 	derpAddr := ep.derpAddr
 	if derpAddr.IsValid() {
-		go c.sendDiscoMessage(epAddr{ap: derpAddr}, ep.publicKey, disco.key, allocResp.Message, discoVerboseLog)
+		go c.sendDiscoMessage(epAddr{ap: derpAddr}, ep.publicKey, disco.key(), allocResp.Message, discoVerboseLog)
 	}
 }
 
@@ -2378,7 +2378,8 @@ func (c *Conn) handleDiscoMessage(msg []byte, src epAddr, shouldBeRelayHandshake
 		if epDisco == nil {
 			return
 		}
-		if epDisco.key != di.discoKey {
+		// TODO(cmol): Switch active keys based on what we see here
+		if epDisco.key() != di.discoKey {
 			if isVia {
 				metricRecvDiscoCallMeMaybeViaBadDisco.Add(1)
 			} else {
@@ -2403,13 +2404,13 @@ func (c *Conn) handleDiscoMessage(msg []byte, src epAddr, shouldBeRelayHandshake
 		}
 		if isVia {
 			c.dlogf("[v1] magicsock: disco: %v<-%v via %v (%v, %v)  got call-me-maybe-via, %d endpoints",
-				c.discoAtomic.Short(), epDisco.short, via.ServerDisco.ShortString(),
+				c.discoAtomic.Short(), epDisco.shortString(), via.ServerDisco.ShortString(),
 				ep.publicKey.ShortString(), derpStr(src.String()),
 				len(via.AddrPorts))
 			c.relayManager.handleCallMeMaybeVia(ep, lastBest, lastBestIsTrusted, via)
 		} else {
 			c.dlogf("[v1] magicsock: disco: %v<-%v (%v, %v)  got call-me-maybe, %d endpoints",
-				c.discoAtomic.Short(), epDisco.short,
+				c.discoAtomic.Short(), epDisco.shortString(),
 				ep.publicKey.ShortString(), derpStr(src.String()),
 				len(cmm.MyNumber))
 			go ep.handleCallMeMaybe(cmm)
@@ -2443,7 +2444,8 @@ func (c *Conn) handleDiscoMessage(msg []byte, src epAddr, shouldBeRelayHandshake
 		if epDisco == nil {
 			return
 		}
-		if epDisco.key != di.discoKey {
+		// TODO(cmol): Switch active keys based on what we see here
+		if epDisco.key() != di.discoKey {
 			if isResp {
 				metricRecvDiscoAllocUDPRelayEndpointResponseBadDisco.Add(1)
 			} else {
@@ -2455,7 +2457,7 @@ func (c *Conn) handleDiscoMessage(msg []byte, src epAddr, shouldBeRelayHandshake
 
 		if isResp {
 			c.dlogf("[v1] magicsock: disco: %v<-%v (%v, %v) got %s, %d endpoints",
-				c.discoAtomic.Short(), epDisco.short,
+				c.discoAtomic.Short(), epDisco.shortString(),
 				ep.publicKey.ShortString(), derpStr(src.String()),
 				msgType,
 				len(resp.AddrPorts))
@@ -2469,7 +2471,7 @@ func (c *Conn) handleDiscoMessage(msg []byte, src epAddr, shouldBeRelayHandshake
 			return
 		} else {
 			c.dlogf("[v1] magicsock: disco: %v<-%v (%v, %v) got %s disco[0]=%v disco[1]=%v",
-				c.discoAtomic.Short(), epDisco.short,
+				c.discoAtomic.Short(), epDisco.shortString(),
 				ep.publicKey.ShortString(), derpStr(src.String()),
 				msgType,
 				req.ClientDisco[0].ShortString(), req.ClientDisco[1].ShortString())
@@ -2507,9 +2509,10 @@ func (c *Conn) handleDiscoMessage(msg []byte, src epAddr, shouldBeRelayHandshake
 // c.mu must be held.
 func (c *Conn) unambiguousNodeKeyOfPingLocked(dm *disco.Ping, dk key.DiscoPublic, derpNodeSrc key.NodePublic) (nk key.NodePublic, ok bool) {
 	if !derpNodeSrc.IsZero() {
+		// TODO(cmol): Switch active keys based on what we see here
 		if ep, ok := c.peerMap.endpointForNodeKey(derpNodeSrc); ok {
 			epDisco := ep.disco.Load()
-			if epDisco != nil && epDisco.key == dk {
+			if epDisco != nil && epDisco.key() == dk {
 				return derpNodeSrc, true
 			}
 		}
@@ -2519,7 +2522,8 @@ func (c *Conn) unambiguousNodeKeyOfPingLocked(dm *disco.Ping, dk key.DiscoPublic
 	if !dm.NodeKey.IsZero() {
 		if ep, ok := c.peerMap.endpointForNodeKey(dm.NodeKey); ok {
 			epDisco := ep.disco.Load()
-			if epDisco != nil && epDisco.key == dk {
+			// TODO(cmol): Switch active keys based on what we see here
+			if epDisco != nil && epDisco.key() == dk {
 				return dm.NodeKey, true
 			}
 		}
@@ -2527,6 +2531,9 @@ func (c *Conn) unambiguousNodeKeyOfPingLocked(dm *disco.Ping, dk key.DiscoPublic
 
 	// If there's exactly 1 node in our netmap with DiscoKey dk,
 	// then it's not ambiguous which node key dm was from.
+	c.peerMap.nodesMu.RLock()
+	defer c.peerMap.nodesMu.RUnlock()
+
 	if set := c.peerMap.nodesOfDisco[dk]; len(set) == 1 {
 		for nk = range set {
 			return nk, true
@@ -2657,7 +2664,7 @@ func (c *Conn) enqueueCallMeMaybe(derpAddr netip.AddrPort, de *endpoint) {
 		c.dlogf("[v1] magicsock: want call-me-maybe but endpoints stale; restunning")
 
 		mak.Set(&c.onEndpointRefreshed, de, func() {
-			c.dlogf("[v1] magicsock: STUN done; sending call-me-maybe to %v %v", epDisco.short, de.publicKey.ShortString())
+			c.dlogf("[v1] magicsock: STUN done; sending call-me-maybe to %v %v", epDisco.shortString(), de.publicKey.ShortString())
 			c.enqueueCallMeMaybe(derpAddr, de)
 		})
 		// TODO(bradfitz): make a new 'reSTUNQuickly' method
@@ -2676,12 +2683,12 @@ func (c *Conn) enqueueCallMeMaybe(derpAddr netip.AddrPort, de *endpoint) {
 	for _, ep := range c.lastEndpoints {
 		eps = append(eps, ep.Addr)
 	}
-	go de.c.sendDiscoMessage(epAddr{ap: derpAddr}, de.publicKey, epDisco.key, &disco.CallMeMaybe{MyNumber: eps}, discoLog)
+	go de.c.sendDiscoMessage(epAddr{ap: derpAddr}, de.publicKey, epDisco.key(), &disco.CallMeMaybe{MyNumber: eps}, discoLog)
 	if debugSendCallMeUnknownPeer() {
 		// Send a callMeMaybe packet to a non-existent peer
 		unknownKey := key.NewNode().Public()
 		c.logf("magicsock: sending CallMeMaybe to unknown peer per TS_DEBUG_SEND_CALLME_UNKNOWN_PEER")
-		go de.c.sendDiscoMessage(epAddr{ap: derpAddr}, unknownKey, epDisco.key, &disco.CallMeMaybe{MyNumber: eps}, discoLog)
+		go de.c.sendDiscoMessage(epAddr{ap: derpAddr}, unknownKey, epDisco.key(), &disco.CallMeMaybe{MyNumber: eps}, discoLog)
 	}
 }
 
@@ -3219,10 +3226,13 @@ func (c *Conn) upsertPeerLocked(n tailcfg.NodeView, flags debugFlags, entriesPer
 		}
 		var oldDiscoKey key.DiscoPublic
 		if epDisco := ep.disco.Load(); epDisco != nil {
-			oldDiscoKey = epDisco.key
+			// Upserted peers originates from control. Compare with the discoKey
+			// learned from control.
+			oldDiscoKey = epDisco.keyFromControl()
 		}
 		ep.updateFromNode(n, flags.heartbeatDisabled, flags.probeUDPLifetimeOn)
-		c.peerMap.upsertEndpoint(ep, oldDiscoKey) // maybe update discokey mappings in peerMap
+		// Maybe update the control learned discokey mappings in peerMap.
+		c.peerMap.upsertEndpoint(ep, oldDiscoKey, false)
 		return
 	}
 
@@ -3285,7 +3295,7 @@ func (c *Conn) upsertPeerLocked(n tailcfg.NodeView, flags debugFlags, entriesPer
 	}
 
 	ep.updateFromNode(n, flags.heartbeatDisabled, flags.probeUDPLifetimeOn)
-	c.peerMap.upsertEndpoint(ep, key.DiscoPublic{})
+	c.peerMap.upsertEndpoint(ep, key.DiscoPublic{}, false)
 }
 
 // UpsertPeer adds or updates a single peer in c. It is the efficient
@@ -4520,7 +4530,8 @@ func (c *Conn) HandleDiscoKeyAdvertisement(node tailcfg.NodeView, update packet.
 
 	oldDiscoKey := key.DiscoPublic{}
 	if epDisco := ep.disco.Load(); epDisco != nil {
-		oldDiscoKey = epDisco.key
+		// Compare with the known key (could be a zero key) learned via TSMP.
+		oldDiscoKey = epDisco.keyFromTSMP()
 	}
 	// If the key did not change, count it and return.
 	if oldDiscoKey.Compare(discoKey) == 0 {
@@ -4529,8 +4540,8 @@ func (c *Conn) HandleDiscoKeyAdvertisement(node tailcfg.NodeView, update packet.
 		return
 	}
 	c.discoInfoForKnownPeerLocked(discoKey)
-	ep.updateDiscoKey(discoKey)
-	c.peerMap.upsertEndpoint(ep, oldDiscoKey)
+	ep.updateTSMPDiscoKey(discoKey)
+	c.peerMap.upsertEndpoint(ep, oldDiscoKey, true)
 	if !oldDiscoKey.IsZero() && !c.peerMap.knownPeerDiscoKey(oldDiscoKey) {
 		delete(c.discoInfo, oldDiscoKey)
 	}

wgengine/magicsock/magicsock_test.go

@@ -1878,8 +1878,8 @@ func TestSetNetworkMapChangingNodeKey(t *testing.T) {
 	if deDisco == nil {
 		t.Fatalf("discoEndpoint disco is nil")
 	}
-	if deDisco.key != discoKey {
-		t.Errorf("discoKey = %v; want %v", deDisco.key, discoKey)
+	if deDisco.key() != discoKey {
+		t.Errorf("discoKey = %v; want %v", deDisco.key(), discoKey)
 	}
 	if _, ok := conn.peerMap.endpointForNodeKey(nodeKey1); ok {
 		t.Errorf("didn't expect to find node for key1")
@@ -4195,9 +4195,7 @@ func TestConn_receiveIP(t *testing.T) {
 			publicKey:  key.NewNode().Public(),
 			lastRecvWG: lastRecvWG,
 		}
-		ep.disco.Store(&endpointDisco{
-			key: key.NewDisco().Public(),
-		})
+		ep.updateDiscoKey(key.NewDisco().Public())
 		return ep
 	}
 
@@ -4352,7 +4350,7 @@ func TestConn_receiveIP(t *testing.T) {
 					t.Fatal("unexpected tt.wantEndpointType concrete type")
 				}
 				insertEPIntoPeerMap.c = c
-				c.peerMap.upsertEndpoint(insertEPIntoPeerMap, key.DiscoPublic{})
+				c.peerMap.upsertEndpoint(insertEPIntoPeerMap, key.DiscoPublic{}, false)
 				c.peerMap.setNodeKeyForEpAddr(tt.peerMapEpAddr, insertEPIntoPeerMap.publicKey)
 			}
 
@@ -4474,9 +4472,7 @@ func Test_lazyEndpoint_InitiationMessagePublicKey(t *testing.T) {
 				nodeID:    1,
 				publicKey: key.NewNode().Public(),
 			}
-			ep.disco.Store(&endpointDisco{
-				key: key.NewDisco().Public(),
-			})
+			ep.updateDiscoKey(key.NewDisco().Public())
 
 			conn := newConn(t.Logf)
 			ep.c = conn
@@ -4485,7 +4481,7 @@ func Test_lazyEndpoint_InitiationMessagePublicKey(t *testing.T) {
 			if tt.callWithPeerMapKey {
 				copy(pubKey[:], ep.publicKey.AppendTo(nil))
 			}
-			conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{})
+			conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{}, false)
 
 			le := &lazyEndpoint{
 				c: conn,
@@ -4536,9 +4532,7 @@ func Test_lazyEndpoint_FromPeer(t *testing.T) {
 				nodeID:    1,
 				publicKey: key.NewNode().Public(),
 			}
-			ep.disco.Store(&endpointDisco{
-				key: key.NewDisco().Public(),
-			})
+			ep.updateDiscoKey(key.NewDisco().Public())
 			conn := newConn(t.Logf)
 			ep.c = conn
 
@@ -4546,7 +4540,7 @@ func Test_lazyEndpoint_FromPeer(t *testing.T) {
 			if tt.callWithPeerMapKey {
 				copy(pubKey[:], ep.publicKey.AppendTo(nil))
 			}
-			conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{})
+			conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{}, false)
 
 			le := &lazyEndpoint{
 				c:   conn,
@@ -4672,10 +4666,7 @@ func TestReceiveTSMPDiscoKeyAdvertisement(t *testing.T) {
 				nodeAddr:  netip.MustParseAddr("100.64.0.1"),
 			}
 			discoKey := key.NewDisco().Public()
-			ep.disco.Store(&endpointDisco{
-				key:   discoKey,
-				short: discoKey.ShortString(),
-			})
+			ep.updateDiscoKey(discoKey)
 			ep.c = conn
 			conn.mu.Lock()
 			nodeView := (&tailcfg.Node{
@@ -4687,7 +4678,7 @@ func TestReceiveTSMPDiscoKeyAdvertisement(t *testing.T) {
 			conn.peersByID = map[tailcfg.NodeID]tailcfg.NodeView{nodeView.ID(): nodeView}
 			conn.mu.Unlock()
 
-			conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{})
+			conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{}, true)
 
 			if ep.discoShort() != discoKey.ShortString() {
 				t.Errorf("Original disco key %s, does not match %s", discoKey.ShortString(), ep.discoShort())
@@ -4704,8 +4695,8 @@ func TestReceiveTSMPDiscoKeyAdvertisement(t *testing.T) {
 				wantDiscoKey = newDiscoKey
 			}
 
-			if ep.disco.Load().short != wantDiscoKey.ShortString() {
-				t.Errorf("New disco key %s, does not match %s", newDiscoKey.ShortString(), ep.disco.Load().short)
+			if ep.disco.Load().shortString() != wantDiscoKey.ShortString() {
+				t.Errorf("New disco key %s, does not match %s", newDiscoKey.ShortString(), ep.disco.Load().shortString())
 			}
 		})
 	}
@@ -4743,10 +4734,7 @@ func TestPriorityMessageForPeer(t *testing.T) {
 	}
 
 	discoKey := key.NewDisco().Public()
-	ep.disco.Store(&endpointDisco{
-		key:   discoKey,
-		short: discoKey.ShortString(),
-	})
+	ep.updateDiscoKey(discoKey)
 
 	ep.c = conn
 
@@ -4756,7 +4744,7 @@ func TestPriorityMessageForPeer(t *testing.T) {
 	}
 
 	conn.mu.Lock()
-	conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{})
+	conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{}, false)
 	conn.mu.Unlock()
 
 	// Test isWireguardOnly.
@@ -4841,10 +4829,7 @@ func BenchmarkPriorityMessageForPeer(b *testing.B) {
 				}
 
 				discoKey := key.NewDisco().Public()
-				ep.disco.Store(&endpointDisco{
-					key:   discoKey,
-					short: discoKey.ShortString(),
-				})
+				ep.updateDiscoKey(discoKey)
 
 				ep.c = conn
 				nodeView := (&tailcfg.Node{
@@ -4857,7 +4842,7 @@ func BenchmarkPriorityMessageForPeer(b *testing.B) {
 				}).View()
 				peersByID[nodeID] = nodeView
 				conn.mu.Lock()
-				conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{})
+				conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{}, false)
 				conn.mu.Unlock()
 			}
 
@@ -4871,3 +4856,68 @@ func BenchmarkPriorityMessageForPeer(b *testing.B) {
 		})
 	}
 }
+
+// TestHandleDiscoKeyAdvertisementControlKeyPreservedInPeermap verifies that
+// receiving a TSMP disco key advertisement does not evict the peer's existing
+// control key from nodesOfDisco.
+func TestHandleDiscoKeyAdvertisementControlKeyPreservedInPeermap(t *testing.T) {
+	conn := newTestConn(t)
+
+	nk := key.NewNode().Public()
+	dk1 := key.NewDisco().Public() // initial control key
+	dk2 := key.NewDisco().Public() // TSMP key
+
+	peer := &tailcfg.Node{
+		ID:        1,
+		Key:       nk,
+		DiscoKey:  dk1,
+		Addresses: []netip.Prefix{netip.MustParsePrefix("100.64.0.1/32")},
+	}
+	conn.SetNetworkMap(tailcfg.NodeView{}, nodeViews([]*tailcfg.Node{peer}))
+
+	conn.HandleDiscoKeyAdvertisement(peer.View(),
+		packet.TSMPDiscoKeyAdvertisement{Key: dk2})
+
+	if !conn.peerMap.knownPeerDiscoKey(dk1) {
+		t.Error("control key dk1 should still be in peermap")
+	}
+	if !conn.peerMap.knownPeerDiscoKey(dk2) {
+		t.Error("TSMP key dk2 should be in peermap after TSMP advertisement")
+	}
+}
+
+// TestUpsertPeerTSMPKeyPreservedOnControlUpdate verifies that a control-plane
+// UpsertPeer call does not evict the TSMP key from nodesOfDisco, but does
+// evict the old control learned key.
+func TestUpsertPeerTSMPKeyPreservedOnControlUpdate(t *testing.T) {
+	conn := newTestConn(t)
+
+	nk := key.NewNode().Public()
+	dk1 := key.NewDisco().Public() // control key
+	dk2 := key.NewDisco().Public() // TSMP key
+	dk3 := key.NewDisco().Public() // new control key
+
+	peer := &tailcfg.Node{
+		ID:        1,
+		Key:       nk,
+		DiscoKey:  dk1,
+		Addresses: []netip.Prefix{netip.MustParsePrefix("100.64.0.1/32")},
+	}
+	conn.SetNetworkMap(tailcfg.NodeView{}, nodeViews([]*tailcfg.Node{peer}))
+
+	conn.HandleDiscoKeyAdvertisement(peer.View(),
+		packet.TSMPDiscoKeyAdvertisement{Key: dk2})
+
+	peer.DiscoKey = dk3
+	conn.UpsertPeer(peer.View())
+
+	if !conn.peerMap.knownPeerDiscoKey(dk2) {
+		t.Error("TSMP key dk2 should still be in peermap")
+	}
+	if !conn.peerMap.knownPeerDiscoKey(dk3) {
+		t.Error("control key dk3 should be in peermap")
+	}
+	if conn.peerMap.knownPeerDiscoKey(dk1) {
+		t.Error("control key dk1 shoud not be in peermap")
+	}
+}

wgengine/magicsock/peermap.go

@@ -4,6 +4,8 @@
 package magicsock
 
 import (
+	"sync"
+
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/key"
 	"tailscale.com/util/set"
@@ -48,6 +50,7 @@ type peerMap struct {
 	// byEpAddr. That issue is being tracked in http://go/corp/29422.
 	relayEpAddrByNodeKey map[key.NodePublic]epAddr
 
+	nodesMu sync.RWMutex // protects nodesOfDisco
 	// nodesOfDisco contains the set of nodes that are using a
 	// DiscoKey. Usually those sets will be just one node.
 	nodesOfDisco map[key.DiscoPublic]set.Set[key.NodePublic]
@@ -74,6 +77,9 @@ func (m *peerMap) nodeCount() int {
 // knownPeerDiscoKey reports whether there exists any peer with the disco key
 // dk.
 func (m *peerMap) knownPeerDiscoKey(dk key.DiscoPublic) bool {
+	m.nodesMu.RLock()
+	defer m.nodesMu.RUnlock()
+
 	_, ok := m.nodesOfDisco[dk]
 	return ok
 }
@@ -117,8 +123,11 @@ func (m *peerMap) forEachEndpoint(f func(ep *endpoint)) {
 
 // forEachEndpointWithDiscoKey invokes f on every endpoint in m that has the
 // provided DiscoKey until f returns false or there are no endpoints left to
-// iterate.
+// iterate. f must call back into peerMap and mutate [nodesOfDisco].
 func (m *peerMap) forEachEndpointWithDiscoKey(dk key.DiscoPublic, f func(*endpoint) (keepGoing bool)) {
+	m.nodesMu.RLock()
+	defer m.nodesMu.RUnlock()
+
 	for nk := range m.nodesOfDisco[dk] {
 		pi, ok := m.byNodeKey[nk]
 		if !ok {
@@ -138,7 +147,9 @@ func (m *peerMap) forEachEndpointWithDiscoKey(dk key.DiscoPublic, f func(*endpoi
 // upsertEndpoint stores endpoint in the peerInfo for
 // ep.publicKey, and updates indexes. m must already have a
 // tailcfg.Node for ep.publicKey.
-func (m *peerMap) upsertEndpoint(ep *endpoint, oldDiscoKey key.DiscoPublic) {
+func (m *peerMap) upsertEndpoint(ep *endpoint, oldDiscoKey key.DiscoPublic,
+	fromTSMP bool,
+) {
 	if ep.nodeID == 0 {
 		panic("internal error: upsertEndpoint called with zero NodeID")
 	}
@@ -149,13 +160,14 @@ func (m *peerMap) upsertEndpoint(ep *endpoint, oldDiscoKey key.DiscoPublic) {
 	}
 	m.byNodeID[ep.nodeID] = pi
 
+	// Load key and make the comparison based on the source of the new key.
 	epDisco := ep.disco.Load()
-	if epDisco == nil || oldDiscoKey != epDisco.key {
-		s := m.nodesOfDisco[oldDiscoKey]
-		delete(s, ep.publicKey)
-		if len(s) == 0 {
-			delete(m.nodesOfDisco, oldDiscoKey)
-		}
+	epKey := epDisco.keyFromControl()
+	if fromTSMP {
+		epKey = epDisco.keyFromTSMP()
+	}
+	if epDisco == nil || oldDiscoKey != epKey {
+		m.cleanFromNodesOfDisco(oldDiscoKey, ep.publicKey)
 	}
 	if ep.isWireguardOnly {
 		// If the peer is a WireGuard only peer, add all of its endpoints.
@@ -170,12 +182,22 @@ func (m *peerMap) upsertEndpoint(ep *endpoint, oldDiscoKey key.DiscoPublic) {
 		}
 		return
 	}
-	discoSet := m.nodesOfDisco[epDisco.key]
+
+	// There is no need to insert a new node key under a zero disco key as it
+	// wastes space and will risk stale entries just sitting there in a long list
+	// under the zero key.
+	if epKey.IsZero() {
+		return
+	}
+
+	m.nodesMu.Lock()
+	discoSet := m.nodesOfDisco[epKey]
 	if discoSet == nil {
 		discoSet = set.Set[key.NodePublic]{}
-		m.nodesOfDisco[epDisco.key] = discoSet
+		m.nodesOfDisco[epKey] = discoSet
 	}
 	discoSet.Add(ep.publicKey)
+	m.nodesMu.Unlock()
 }
 
 // setNodeKeyForEpAddr makes future peer lookups by addr return the
@@ -218,12 +240,16 @@ func (m *peerMap) deleteEndpoint(ep *endpoint) {
 
 	pi := m.byNodeKey[ep.publicKey]
 	if epDisco != nil {
-		s := m.nodesOfDisco[epDisco.key]
-		delete(s, ep.publicKey)
-		if len(s) == 0 {
-			delete(m.nodesOfDisco, epDisco.key)
+		for _, discoKey := range []key.DiscoPublic{
+			epDisco.keyFromControl(), epDisco.keyFromTSMP(),
+		} {
+			if discoKey.IsZero() {
+				continue
+			}
+			m.cleanFromNodesOfDisco(discoKey, ep.publicKey)
 		}
 	}
+
 	delete(m.byNodeKey, ep.publicKey)
 	if was, ok := m.byNodeID[ep.nodeID]; ok && was.ep == ep {
 		delete(m.byNodeID, ep.nodeID)
@@ -238,3 +264,13 @@ func (m *peerMap) deleteEndpoint(ep *endpoint) {
 	}
 	delete(m.relayEpAddrByNodeKey, ep.publicKey)
 }
+
+func (m *peerMap) cleanFromNodesOfDisco(disco key.DiscoPublic, node key.NodePublic) {
+	m.nodesMu.Lock()
+	defer m.nodesMu.Unlock()
+	s := m.nodesOfDisco[disco]
+	delete(s, node)
+	if len(s) == 0 {
+		delete(m.nodesOfDisco, disco)
+	}
+}

wgengine/magicsock/peermap_test.go

@@ -18,9 +18,8 @@ func Test_peerMap_oneRelayEpAddrPerNK(t *testing.T) {
 		nodeID:    1,
 		publicKey: nk,
 	}
-	ed := &endpointDisco{key: key.NewDisco().Public()}
-	ep.disco.Store(ed)
-	pm.upsertEndpoint(ep, key.DiscoPublic{})
+	ep.updateDiscoKey(key.NewDisco().Public())
+	pm.upsertEndpoint(ep, key.DiscoPublic{}, false)
 	vni := packet.VirtualNetworkID{}
 	vni.Set(1)
 	relayEpAddrA := epAddr{ap: netip.MustParseAddrPort("127.0.0.1:1"), vni: vni}
@@ -43,16 +42,16 @@ func Test_peerMap_nodesOfDisco_upsertCleansOldKey(t *testing.T) {
 	discoK2 := key.NewDisco().Public()
 
 	ep := &endpoint{nodeID: 1, publicKey: nk}
-	ep.disco.Store(&endpointDisco{key: discoK1})
-	pm.upsertEndpoint(ep, key.DiscoPublic{}) // insert with K1
+	ep.updateDiscoKey(discoK1)
+	pm.upsertEndpoint(ep, key.DiscoPublic{}, false) // insert with K1
 
 	if !pm.knownPeerDiscoKey(discoK1) {
 		t.Fatal("expected K1 to be known after initial upsert")
 	}
 
 	// Rotate disco
-	ep.disco.Store(&endpointDisco{key: discoK2})
-	pm.upsertEndpoint(ep, discoK1)
+	ep.updateDiscoKey(discoK2)
+	pm.upsertEndpoint(ep, discoK1, false)
 
 	if pm.knownPeerDiscoKey(discoK1) {
 		t.Error("old disco key K1 is still known after rotation")
@@ -77,8 +76,8 @@ func Test_peerMap_nodesOfDisco_deleteCleansKey(t *testing.T) {
 		c:             conn,
 		endpointState: map[netip.AddrPort]*endpointState{},
 	}
-	ep.disco.Store(&endpointDisco{key: dk})
-	pm.upsertEndpoint(ep, key.DiscoPublic{})
+	ep.updateDiscoKey(dk)
+	pm.upsertEndpoint(ep, key.DiscoPublic{}, false)
 
 	if !pm.knownPeerDiscoKey(dk) {
 		t.Fatal("expected disco key to be known after upsert")
@@ -108,8 +107,8 @@ func Test_peerMap_nodesOfDisco_sharedDiscoKey(t *testing.T) {
 		c:             conn,
 		endpointState: map[netip.AddrPort]*endpointState{},
 	}
-	ep1.disco.Store(&endpointDisco{key: dk})
-	pm.upsertEndpoint(ep1, key.DiscoPublic{})
+	ep1.updateDiscoKey(dk)
+	pm.upsertEndpoint(ep1, key.DiscoPublic{}, false)
 
 	ep2 := &endpoint{
 		nodeID:        2,
@@ -117,8 +116,8 @@ func Test_peerMap_nodesOfDisco_sharedDiscoKey(t *testing.T) {
 		c:             conn,
 		endpointState: map[netip.AddrPort]*endpointState{},
 	}
-	ep2.disco.Store(&endpointDisco{key: dk})
-	pm.upsertEndpoint(ep2, key.DiscoPublic{})
+	ep2.updateDiscoKey(dk)
+	pm.upsertEndpoint(ep2, key.DiscoPublic{}, false)
 
 	pm.deleteEndpoint(ep1)
 
@@ -132,3 +131,85 @@ func Test_peerMap_nodesOfDisco_sharedDiscoKey(t *testing.T) {
 		t.Error("disco key should be unknown after both peers removed")
 	}
 }
+
+func TestPeerMapNodesOfDiscoMultipleKeySourcesDeleteEndpoint(t *testing.T) {
+	pm := newPeerMap()
+	nk := key.NewNode().Public()
+	dk1 := key.NewDisco().Public()
+	dk2 := key.NewDisco().Public()
+
+	conn := newTestConn(t)
+
+	ep := &endpoint{
+		nodeID:        1,
+		publicKey:     nk,
+		c:             conn,
+		endpointState: map[netip.AddrPort]*endpointState{},
+	}
+	ep.updateDiscoKey(dk1)
+	pm.upsertEndpoint(ep, key.DiscoPublic{}, false)
+	ep.updateTSMPDiscoKey(dk2)
+	pm.upsertEndpoint(ep, key.DiscoPublic{}, true)
+
+	if !pm.knownPeerDiscoKey(dk1) {
+		t.Error("disco key 1 should be known")
+	}
+	if !pm.knownPeerDiscoKey(dk2) {
+		t.Error("disco key 2 should be known")
+	}
+
+	// Delete endpoint, nothing should be known.
+	pm.deleteEndpoint(ep)
+	for d, s := range pm.nodesOfDisco {
+		for n := range s {
+			if n == nk {
+				t.Errorf("node should be unknown, found nk: %v, under dk: %v", n, d)
+			}
+		}
+	}
+}
+
+func TestPeerMapNodesOfDiscoMultipleKeySourcesUpsertEndpoint(t *testing.T) {
+	pm := newPeerMap()
+	nk := key.NewNode().Public()
+	dk1 := key.NewDisco().Public()
+	dk2 := key.NewDisco().Public()
+
+	conn := newTestConn(t)
+
+	ep := &endpoint{
+		nodeID:        1,
+		publicKey:     nk,
+		c:             conn,
+		endpointState: map[netip.AddrPort]*endpointState{},
+	}
+	ep.updateDiscoKey(dk1)
+	pm.upsertEndpoint(ep, key.DiscoPublic{}, false)
+	ep.updateTSMPDiscoKey(dk2)
+	pm.upsertEndpoint(ep, key.DiscoPublic{}, true)
+
+	if !pm.knownPeerDiscoKey(dk1) {
+		t.Error("disco key 1 should be known")
+	}
+	if !pm.knownPeerDiscoKey(dk2) {
+		t.Error("disco key 2 should be known")
+	}
+
+	// Delete control learned key, tsmp key should still be known.
+	ep.updateDiscoKey(key.DiscoPublic{})
+	pm.upsertEndpoint(ep, dk1, false)
+	if !pm.knownPeerDiscoKey(dk2) {
+		t.Error("disco key 2 should be known")
+	}
+
+	// Delete TSMP learned key, nothing should be known.
+	ep.updateTSMPDiscoKey(key.DiscoPublic{})
+	pm.upsertEndpoint(ep, dk2, true)
+	for d, s := range pm.nodesOfDisco {
+		for n := range s {
+			if n == nk {
+				t.Errorf("node should be unknown, found nk: %v, under dk: %v", n, d)
+			}
+		}
+	}
+}

wgengine/magicsock/relaymanager.go

@@ -869,7 +869,7 @@ func (r *relayManager) sendCallMeMaybeVia(ep *endpoint, se udprelay.ServerEndpoi
 			AddrPorts:           se.AddrPorts,
 		},
 	}
-	ep.c.sendDiscoMessage(epAddr{ap: derpAddr}, ep.publicKey, epDisco.key, callMeMaybeVia, discoVerboseLog)
+	ep.c.sendDiscoMessage(epAddr{ap: derpAddr}, ep.publicKey, epDisco.key(), callMeMaybeVia, discoVerboseLog)
 }
 
 func (r *relayManager) handshakeServerEndpoint(work *relayHandshakeWork) {
@@ -892,7 +892,7 @@ func (r *relayManager) handshakeServerEndpoint(work *relayHandshakeWork) {
 	common := disco.BindUDPRelayEndpointCommon{
 		VNI:        work.se.VNI,
 		Generation: work.handshakeGen,
-		RemoteKey:  epDisco.key,
+		RemoteKey:  epDisco.key(),
 	}
 
 	work.dlogf("[v1] magicsock: relayManager: starting handshake addrPorts=%v",
@@ -946,7 +946,7 @@ func (r *relayManager) handshakeServerEndpoint(work *relayHandshakeWork) {
 				answer.Challenge = *withAnswer
 				ep.c.sendDiscoMessage(epAddr{ap: to, vni: vni}, key.NodePublic{}, work.se.ServerDisco, answer, discoVerboseLog)
 			}
-			ep.c.sendDiscoMessage(epAddr{ap: to, vni: vni}, key.NodePublic{}, epDisco.key, ping, discoVerboseLog)
+			ep.c.sendDiscoMessage(epAddr{ap: to, vni: vni}, key.NodePublic{}, epDisco.key(), ping, discoVerboseLog)
 		}()
 	}
 
@@ -954,7 +954,7 @@ func (r *relayManager) handshakeServerEndpoint(work *relayHandshakeWork) {
 		if common.VNI != work.se.VNI {
 			return errors.New("mismatching VNI")
 		}
-		if common.RemoteKey.Compare(epDisco.key) != 0 {
+		if common.RemoteKey.Compare(epDisco.key()) != 0 {
 			return errors.New("mismatching RemoteKey")
 		}
 		return nil
@@ -1099,7 +1099,7 @@ func (r *relayManager) allocateAllServersRunLoop(wlb endpointWithLastBest) {
 	if remoteDisco == nil {
 		return
 	}
-	discoKeys := key.NewSortedPairOfDiscoPublic(wlb.ep.c.discoAtomic.Public(), remoteDisco.key)
+	discoKeys := key.NewSortedPairOfDiscoPublic(wlb.ep.c.discoAtomic.Public(), remoteDisco.key())
 	for _, v := range r.serversByNodeKey {
 		byDiscoKeys, ok := r.allocWorkByDiscoKeysByServerNodeKey[v.nodeKey]
 		if !ok {