From bc7fe6b8ebb1c1d4527b70ac78c47857f40b0120 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Wed, 30 Sep 2026 15:29:58 +0000
Subject: [PATCH] state: test NodeKey index agrees with node list

Covers put, rotation and payload writes on the reuse path, and delete.
---
 hscontrol/state/node_store_test.go | 130 +++++++++++++++++++++++++++++
 1 file changed, 130 insertions(+)

diff --git a/hscontrol/state/node_store_test.go b/hscontrol/state/node_store_test.go
index a62a8148f..f405eba1d 100644
--- a/hscontrol/state/node_store_test.go
+++ b/hscontrol/state/node_store_test.go
@@ -2439,3 +2439,133 @@ func TestPolicyCachesSurviveOldViewDuringBuild(t *testing.T) {
 		})
 	}
 }
+
+// nodeKeyIndexMismatch describes where snap's NodeKey index disagrees with its
+// node list, or returns "" when every listed node, and nothing else, is indexed.
+func nodeKeyIndexMismatch(snap *Snapshot) string {
+	if len(snap.nodesByNodeKey) != len(snap.allNodes) {
+		return fmt.Sprintf("index holds %d keys, list holds %d nodes",
+			len(snap.nodesByNodeKey), len(snap.allNodes))
+	}
+
+	for _, nv := range snap.allNodes {
+		got, ok := snap.nodesByNodeKey[nv.NodeKey()]
+		if !ok || got.ID() != nv.ID() {
+			return fmt.Sprintf("node %d: key indexed=%t id=%d", nv.ID(), ok, got.ID())
+		}
+	}
+
+	return ""
+}
+
+// TestNodeStoreNodeKeyIndexMatchesNodes pins the index DERP admission reads:
+// every published snapshot indexes exactly its listed nodes' keys, and a key
+// rotated away or deleted stops resolving, also on writes that reuse the
+// previous peer map.
+func TestNodeStoreNodeKeyIndexMatchesNodes(t *testing.T) {
+	var peersCalls atomic.Int64
+
+	countingPeersFunc := func(nodes []types.NodeView) map[types.NodeID][]types.NodeID {
+		peersCalls.Add(1)
+
+		return allowAllPeersFunc(nodes)
+	}
+
+	store := NewNodeStore(nil, countingPeersFunc, TestBatchSize, TestBatchTimeout)
+	store.Start()
+
+	defer store.Stop()
+
+	// Readers load snapshots while the writes below publish them, so a torn
+	// or stale index surfaces here and under -race.
+	var (
+		firstMismatch atomic.Pointer[string]
+		stop          atomic.Bool
+		readers       sync.WaitGroup
+	)
+
+	for range 4 {
+		readers.Go(func() {
+			for !stop.Load() {
+				msg := nodeKeyIndexMismatch(store.data.Load())
+				if msg != "" {
+					firstMismatch.CompareAndSwap(nil, &msg)
+				}
+			}
+		})
+	}
+
+	stopReaders := func() {
+		stop.Store(true)
+		readers.Wait()
+	}
+	t.Cleanup(stopReaders)
+
+	requireResolves := func(t *testing.T, k key.NodePublic, want types.NodeID) {
+		t.Helper()
+
+		nv, ok := store.GetNodeByNodeKey(k)
+		require.True(t, ok, "key of node %d must resolve", want)
+		require.True(t, nv.Valid())
+		require.Equal(t, want, nv.ID())
+	}
+
+	requireRefused := func(t *testing.T, k key.NodePublic) {
+		t.Helper()
+
+		_, ok := store.GetNodeByNodeKey(k)
+		require.False(t, ok, "key must not resolve")
+	}
+
+	n1 := createTestNode(1, 1, "user1", "node1")
+	n2 := createTestNode(2, 1, "user1", "node2")
+	n3 := createTestNode(3, 2, "user2", "node3")
+
+	for _, n := range []types.Node{n1, n2, n3} {
+		store.PutNode(n)
+	}
+
+	requireResolves(t, n1.NodeKey, 1)
+	requireResolves(t, n2.NodeKey, 2)
+	requireResolves(t, n3.NodeKey, 3)
+
+	peersCalls.Store(0)
+
+	// Rotation is not a peer-map input, so it takes the reuse path; repeat it
+	// so the readers overlap many publishes.
+	current := n1.NodeKey
+	for range 50 {
+		prev := current
+		current = key.NewNode().Public()
+
+		_, ok := store.UpdateNode(1, func(n *types.Node) { n.NodeKey = current })
+		require.True(t, ok)
+
+		requireRefused(t, prev)
+		requireResolves(t, current, 1)
+	}
+
+	store.UpdateNode(2, func(n *types.Node) { n.Hostname = "node2-renamed" })
+
+	nv, ok := store.GetNodeByNodeKey(n2.NodeKey)
+	require.True(t, ok)
+	require.Equal(t, "node2-renamed", nv.Hostname(),
+		"payload-only write must republish the indexed view")
+
+	require.Zero(t, peersCalls.Load(),
+		"rotation and payload writes must take the peer-map reuse path")
+
+	store.DeleteNode(3)
+	requireRefused(t, n3.NodeKey)
+	requireResolves(t, current, 1)
+	requireResolves(t, n2.NodeKey, 2)
+
+	stopReaders()
+
+	msg := firstMismatch.Load()
+	if msg != nil {
+		t.Fatalf("snapshot index disagreed with node list: %s", *msg)
+	}
+
+	require.Empty(t, nodeKeyIndexMismatch(store.data.Load()))
+}