wgengine/magicsock: properly clean up peer disco maps (#20543)

Updates tailscale/corp#45124
Updates tailscale/corp#45128

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

wgengine/magicsock/magicsock.go

@@ -4532,6 +4532,9 @@ func (c *Conn) HandleDiscoKeyAdvertisement(node tailcfg.NodeView, update packet.
 	c.discoInfoForKnownPeerLocked(discoKey)
 	ep.updateDiscoKey(discoKey)
 	c.peerMap.upsertEndpoint(ep, oldDiscoKey)
+	if !oldDiscoKey.IsZero() && !c.peerMap.knownPeerDiscoKey(oldDiscoKey) {
+		delete(c.discoInfo, oldDiscoKey)
+	}
 	c.logf("magicsock: updated disco key for peer %v to %v", nodeKey.ShortString(), discoKey.ShortString())
 	metricTSMPDiscoKeyAdvertisementApplied.Add(1)
 }

wgengine/magicsock/peermap.go

@@ -151,7 +151,11 @@ func (m *peerMap) upsertEndpoint(ep *endpoint, oldDiscoKey key.DiscoPublic) {
 
 	epDisco := ep.disco.Load()
 	if epDisco == nil || oldDiscoKey != epDisco.key {
-		delete(m.nodesOfDisco[oldDiscoKey], ep.publicKey)
+		s := m.nodesOfDisco[oldDiscoKey]
+		delete(s, ep.publicKey)
+		if len(s) == 0 {
+			delete(m.nodesOfDisco, oldDiscoKey)
+		}
 	}
 	if ep.isWireguardOnly {
 		// If the peer is a WireGuard only peer, add all of its endpoints.
@@ -214,7 +218,11 @@ func (m *peerMap) deleteEndpoint(ep *endpoint) {
 
 	pi := m.byNodeKey[ep.publicKey]
 	if epDisco != nil {
-		delete(m.nodesOfDisco[epDisco.key], ep.publicKey)
+		s := m.nodesOfDisco[epDisco.key]
+		delete(s, ep.publicKey)
+		if len(s) == 0 {
+			delete(m.nodesOfDisco, epDisco.key)
+		}
 	}
 	delete(m.byNodeKey, ep.publicKey)
 	if was, ok := m.byNodeID[ep.nodeID]; ok && was.ep == ep {

wgengine/magicsock/peermap_test.go

@@ -35,3 +35,100 @@ func Test_peerMap_oneRelayEpAddrPerNK(t *testing.T) {
 		t.Fatalf("expected relay epAddr %v, got: %v", relayEpAddrB, got)
 	}
 }
+
+func Test_peerMap_nodesOfDisco_upsertCleansOldKey(t *testing.T) {
+	pm := newPeerMap()
+	nk := key.NewNode().Public()
+	discoK1 := key.NewDisco().Public()
+	discoK2 := key.NewDisco().Public()
+
+	ep := &endpoint{nodeID: 1, publicKey: nk}
+	ep.disco.Store(&endpointDisco{key: discoK1})
+	pm.upsertEndpoint(ep, key.DiscoPublic{}) // 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)
+
+	if pm.knownPeerDiscoKey(discoK1) {
+		t.Error("old disco key K1 is still known after rotation")
+	}
+	if old, ok := pm.nodesOfDisco[discoK1]; ok {
+		t.Errorf("old disco key K1 should be absent from nodesOfDisco, but entry %v remains", old)
+	}
+	if !pm.knownPeerDiscoKey(discoK2) {
+		t.Error("new disco key K2 should be known after rotation")
+	}
+}
+
+func Test_peerMap_nodesOfDisco_deleteCleansKey(t *testing.T) {
+	pm := newPeerMap()
+	nk := key.NewNode().Public()
+	dk := key.NewDisco().Public()
+
+	conn := newTestConn(t)
+	ep := &endpoint{
+		nodeID:        1,
+		publicKey:     nk,
+		c:             conn,
+		endpointState: map[netip.AddrPort]*endpointState{},
+	}
+	ep.disco.Store(&endpointDisco{key: dk})
+	pm.upsertEndpoint(ep, key.DiscoPublic{})
+
+	if !pm.knownPeerDiscoKey(dk) {
+		t.Fatal("expected disco key to be known after upsert")
+	}
+
+	pm.deleteEndpoint(ep)
+
+	if pm.knownPeerDiscoKey(dk) {
+		t.Error("disco key is still known after deletion")
+	}
+	if s, ok := pm.nodesOfDisco[dk]; ok {
+		t.Errorf("nodesOfDisco for deleted key: found %v, want absent", s)
+	}
+}
+
+func Test_peerMap_nodesOfDisco_sharedDiscoKey(t *testing.T) {
+	pm := newPeerMap()
+	nk1 := key.NewNode().Public()
+	nk2 := key.NewNode().Public()
+	dk := key.NewDisco().Public()
+
+	conn := newTestConn(t)
+
+	ep1 := &endpoint{
+		nodeID:        1,
+		publicKey:     nk1,
+		c:             conn,
+		endpointState: map[netip.AddrPort]*endpointState{},
+	}
+	ep1.disco.Store(&endpointDisco{key: dk})
+	pm.upsertEndpoint(ep1, key.DiscoPublic{})
+
+	ep2 := &endpoint{
+		nodeID:        2,
+		publicKey:     nk2,
+		c:             conn,
+		endpointState: map[netip.AddrPort]*endpointState{},
+	}
+	ep2.disco.Store(&endpointDisco{key: dk})
+	pm.upsertEndpoint(ep2, key.DiscoPublic{})
+
+	pm.deleteEndpoint(ep1)
+
+	if !pm.knownPeerDiscoKey(dk) {
+		t.Error("shared disco key should still be known after one of two peers is removed")
+	}
+
+	pm.deleteEndpoint(ep2)
+
+	if pm.knownPeerDiscoKey(dk) {
+		t.Error("disco key should be unknown after both peers removed")
+	}
+}