control/controlclient: filter out disco updates from full map (#19220)

When getting a full map from control, disco keys for the nodes will also
be delivered. When communicating with a peer that is running without
being connected to control, and having that connection running based on
a TSMP learned disco key, we need to avoid overwriting the disco key for
that peer with the stale one control knows about.

Add a filter that filteres out keys from control, and replace them with
the TSMP learned disco keys.

Updates #12639

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

control/controlclient/map.go

@@ -295,6 +295,7 @@ func (ms *mapSession) handleNonKeepAliveMapResponse(ctx context.Context, resp *t
 	ms.patchifyPeersChanged(resp)
 
 	ms.removeUnwantedDiscoUpdates(resp)
+	ms.removeUnwantedDiscoUpdatesFromFullNetmapUpdate(resp)
 
 	ms.updateStateFromResponse(resp)
 
@@ -404,6 +405,9 @@ type updateStats struct {
 // where the node is offline and has last been seen before the recorded last seen.
 func (ms *mapSession) removeUnwantedDiscoUpdates(resp *tailcfg.MapResponse) {
 	existingMap := ms.netmap()
+	if existingMap == nil {
+		return
+	}
 	acceptedDiscoUpdates := resp.PeersChangedPatch[:0]
 
 	for _, change := range resp.PeersChangedPatch {
@@ -442,6 +446,62 @@ func (ms *mapSession) removeUnwantedDiscoUpdates(resp *tailcfg.MapResponse) {
 	resp.PeersChangedPatch = acceptedDiscoUpdates
 }
 
+// removeUnwantedDiscoUpdatesFromFullNetmapUpdate makes a pass over the full
+// set of peers in an update, usually only received when getting a full netmap
+// from control at startup. If the pass finds a peer with a disco key where the
+// local netmap has a newer key learned via TSMP, overwrite the update with the
+// key from TSMP.
+func (ms *mapSession) removeUnwantedDiscoUpdatesFromFullNetmapUpdate(resp *tailcfg.MapResponse) {
+	if len(resp.Peers) == 0 {
+		return
+	}
+	existingMap := ms.netmap()
+	if existingMap == nil {
+		return
+	}
+	for _, peer := range resp.Peers {
+		if peer.DiscoKey.IsZero() {
+			continue
+		}
+
+		// Accept if:
+		// - peer is new
+		peerIdx := existingMap.PeerIndexByNodeID(peer.ID)
+		if peerIdx < 0 {
+			continue
+		}
+
+		// Accept if:
+		// - disco key has not changed
+		existingNode := existingMap.Peers[peerIdx]
+		if existingNode.DiscoKey() == peer.DiscoKey {
+			continue
+		}
+
+		// Accept if:
+		// - key has changed but peer is online
+		if peer.Online != nil && *peer.Online {
+			continue
+		}
+
+		// Accept if:
+		// - there's no last seen on the existing node
+		existingLastSeen, ok := existingNode.LastSeen().GetOk()
+		if !ok {
+			continue
+		}
+
+		// Accept if:
+		// - last seen on on control is higher
+		if peer.LastSeen != nil && peer.LastSeen.After(existingLastSeen) {
+			continue
+		}
+
+		// Overwrite the key in the full netmap update.
+		peer.DiscoKey = existingNode.DiscoKey()
+	}
+}
+
 // updateStateFromResponse updates ms from res. It takes ownership of res.
 func (ms *mapSession) updateStateFromResponse(resp *tailcfg.MapResponse) {
 	ms.updatePeersStateFromResponse(resp)

control/controlclient/map_test.go

@@ -743,6 +743,7 @@ func TestUpdateDiscoForNodeCallback(t *testing.T) {
 				nu.lastTSMPKey, nu.lastTSMPDisco)
 		}
 	})
+	// Even though key stays in list of update, the updater only triggers on TSMP.
 	t.Run("key_not_wired_through_to_updater", func(t *testing.T) {
 		nu := &rememberLastNetmapUpdater{
 			done: make(chan any, 1),
@@ -776,6 +777,7 @@ func TestUpdateDiscoForNodeCallback(t *testing.T) {
 				DiscoKey: &newKey,
 			}},
 		}
+		// Not TSMP Path, just regular injection path.
 		ms.HandleNonKeepAliveMapResponse(t.Context(), resp)
 		<-nu.done
 
@@ -786,6 +788,119 @@ func TestUpdateDiscoForNodeCallback(t *testing.T) {
 	})
 }
 
+func TestUpdateDiscoForNodeCallbackWithFullNetmap(t *testing.T) {
+	now := time.Now()
+	oldTime := time.Unix(1, 0)
+
+	tests := []struct {
+		name            string
+		initialOnline   bool
+		initialLastSeen time.Time
+		updateOnline    bool
+		updateLastSeen  time.Time
+		expectNewDisco  bool
+	}{
+		{
+			name:            "disco key wired through when newer lastSeen",
+			initialOnline:   false,
+			initialLastSeen: oldTime,
+			updateOnline:    false,
+			updateLastSeen:  now,
+			expectNewDisco:  true,
+		},
+		{
+			name:            "disco key NOT wired through when older lastSeen",
+			initialOnline:   false,
+			initialLastSeen: now,
+			updateOnline:    false,
+			updateLastSeen:  oldTime,
+			expectNewDisco:  false,
+		},
+		{
+			name:            "disco key wired through when newer lastSeen, going offline",
+			initialOnline:   true,
+			initialLastSeen: oldTime,
+			updateOnline:    false,
+			updateLastSeen:  now,
+			expectNewDisco:  true,
+		},
+		{
+			name:            "online flip with newer lastSeen",
+			initialOnline:   false,
+			initialLastSeen: oldTime,
+			updateOnline:    true,
+			updateLastSeen:  now,
+			expectNewDisco:  true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			nu := &rememberLastNetmapUpdater{
+				done: make(chan any, 1),
+			}
+			ms := newTestMapSession(t, nu)
+
+			oldKey := key.NewDisco()
+
+			// Initial node
+			node := tailcfg.Node{
+				ID:       1,
+				Key:      key.NewNode().Public(),
+				DiscoKey: oldKey.Public(),
+				Online:   new(tt.initialOnline),
+				LastSeen: new(tt.initialLastSeen),
+				Name:     "host.network.ts.net",
+			}
+
+			if nm := ms.netmapForResponse(&tailcfg.MapResponse{
+				Peers: []*tailcfg.Node{&node},
+			}); len(nm.Peers) != 1 {
+				t.Fatalf("node not inserted")
+			}
+
+			newKey := key.NewDisco()
+
+			// Updated node
+			newNode := tailcfg.Node{
+				ID:       1,
+				Key:      node.Key,
+				DiscoKey: newKey.Public(),
+				Online:   new(tt.updateOnline),
+				LastSeen: new(tt.updateLastSeen),
+				Name:     "host.network.ts.net",
+			}
+
+			ms.HandleNonKeepAliveMapResponse(t.Context(), &tailcfg.MapResponse{
+				Node: &newNode,
+				Peers: []*tailcfg.Node{
+					&newNode,
+				},
+			})
+			<-nu.done
+
+			newMap := nu.last
+			if n := len(newMap.Peers); n != 1 {
+				t.Fatalf("netmap not right length, got %d, expected %d", n, 1)
+			}
+
+			peer := newMap.Peers[0]
+
+			expectedDisco := oldKey.Public()
+			if tt.expectNewDisco {
+				expectedDisco = newKey.Public()
+			}
+
+			if peer.Key() != node.Key || peer.DiscoKey() != expectedDisco {
+				t.Fatalf("expected [%s]=%s, got [%s]=%s",
+					node.Key, expectedDisco,
+					peer.Key(), peer.DiscoKey(),
+				)
+			}
+		})
+	}
+}
+
 func first[T any](s []T) T {
 	if len(s) == 0 {
 		var zero T