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>
2026-07-20 13:12:58 -04:00
parent be0e460a20
commit b14f7b7543
3 changed files with 110 additions and 2 deletions
@@ -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")
+	}
+}