control/controlclient,net/tstun,wgengine/magicsock: fix handling of zero keys in TSMP (#20508)

Updates tailscale/corp#45042 Signed-off-by: Claus Lensbøl <claus@tailscale.com>
2026-07-17 14:02:53 -04:00
parent c1edf7f458
commit 82a381e54b
6 changed files with 121 additions and 49 deletions
@@ -4495,6 +4495,11 @@ func (c *Conn) PeerRelays() set.Set[netip.Addr] {
 // node is the Tailscale tailcfg.NodeView of the peer that sent the update.
 func (c *Conn) HandleDiscoKeyAdvertisement(node tailcfg.NodeView, update packet.TSMPDiscoKeyAdvertisement) {
 	discoKey := update.Key
+	if discoKey.IsZero() {
+		c.logf("[v1] magicsock: received zero disco key update from %v", node.StableID())
+		return
+	}
+
 	c.logf("magicsock: received disco key update %v from %v", discoKey.ShortString(), node.StableID())
 	metricTSMPDiscoKeyAdvertisementReceived.Add(1)

@@ -4618,46 +4618,73 @@ func TestRotateDiscoKeyMultipleTimes(t *testing.T) {
 }

 func TestReceiveTSMPDiscoKeyAdvertisement(t *testing.T) {
-	conn := newTestConn(t)
-	t.Cleanup(func() { conn.Close() })
-
-	peerKey := key.NewNode().Public()
-	ep := &endpoint{
-		nodeID:    1,
-		publicKey: peerKey,
-		nodeAddr:  netip.MustParseAddr("100.64.0.1"),
-	}
-	discoKey := key.NewDisco().Public()
-	ep.disco.Store(&endpointDisco{
-		key:   discoKey,
-		short: discoKey.ShortString(),
-	})
-	ep.c = conn
-	conn.mu.Lock()
-	nodeView := (&tailcfg.Node{
-		Key: ep.publicKey,
-		Addresses: []netip.Prefix{
-			netip.MustParsePrefix("100.64.0.1/32"),
+	tests := []struct {
+		name       string
+		newKeyFunc func() key.DiscoPublic
+		wantUpdate bool
+	}{
+		{
+			name:       "normal_key_change",
+			newKeyFunc: key.NewDisco().Public,
+			wantUpdate: true,
+		},
+		{
+			name: "zero_key_change",
+			newKeyFunc: func() key.DiscoPublic {
+				return key.DiscoPublic{}
+			},
+			wantUpdate: false,
 		},
-	}).View()
-	conn.peersByID = map[tailcfg.NodeID]tailcfg.NodeView{nodeView.ID(): nodeView}
-	conn.mu.Unlock()
-
-	conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{})
-
-	if ep.discoShort() != discoKey.ShortString() {
-		t.Errorf("Original disco key %s, does not match %s", discoKey.ShortString(), ep.discoShort())
 	}

-	newDiscoKey := key.NewDisco().Public()
-	tka := packet.TSMPDiscoKeyAdvertisement{
-		Src: netip.MustParseAddr("100.64.0.1"),
-		Key: newDiscoKey,
-	}
-	conn.HandleDiscoKeyAdvertisement(nodeView, tka)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			conn := newTestConn(t)
+			t.Cleanup(func() { conn.Close() })

-	if ep.disco.Load().short != newDiscoKey.ShortString() {
-		t.Errorf("New disco key %s, does not match %s", newDiscoKey.ShortString(), ep.disco.Load().short)
+			peerKey := key.NewNode().Public()
+			ep := &endpoint{
+				nodeID:    1,
+				publicKey: peerKey,
+				nodeAddr:  netip.MustParseAddr("100.64.0.1"),
+			}
+			discoKey := key.NewDisco().Public()
+			ep.disco.Store(&endpointDisco{
+				key:   discoKey,
+				short: discoKey.ShortString(),
+			})
+			ep.c = conn
+			conn.mu.Lock()
+			nodeView := (&tailcfg.Node{
+				Key: ep.publicKey,
+				Addresses: []netip.Prefix{
+					netip.MustParsePrefix("100.64.0.1/32"),
+				},
+			}).View()
+			conn.peersByID = map[tailcfg.NodeID]tailcfg.NodeView{nodeView.ID(): nodeView}
+			conn.mu.Unlock()
+
+			conn.peerMap.upsertEndpoint(ep, key.DiscoPublic{})
+
+			if ep.discoShort() != discoKey.ShortString() {
+				t.Errorf("Original disco key %s, does not match %s", discoKey.ShortString(), ep.discoShort())
+			}
+
+			newDiscoKey := tt.newKeyFunc()
+			tka := packet.TSMPDiscoKeyAdvertisement{
+				Src: netip.MustParseAddr("100.64.0.1"),
+				Key: newDiscoKey,
+			}
+			conn.HandleDiscoKeyAdvertisement(nodeView, tka)
+			wantDiscoKey := discoKey
+			if tt.wantUpdate {
+				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)
+			}
+		})
 	}
 }