state: make any-user machine-key lookup deterministic

The lookup returned the first map match, so re-auth branch choice varied
with map order once a machine key had more than one node. Prefer the tagged
node, else the lowest node ID.

Updates #3312
This commit is contained in:
Kristoffer Dalby
2026-06-15 11:07:26 +00:00
committed by Kristoffer Dalby
parent a73d38bb3f
commit fd08b8fa8c
2 changed files with 71 additions and 6 deletions
+15 -6
View File
@@ -804,12 +804,15 @@ func (s *NodeStore) GetNodeByMachineKey(machineKey key.MachinePublic, userID typ
return types.NodeView{}, false return types.NodeView{}, false
} }
// GetNodeByMachineKeyAnyUser returns the first node with the given machine key, // GetNodeByMachineKeyAnyUser returns a node with the given machine key,
// regardless of which user it belongs to. This is useful for scenarios like // regardless of which user it belongs to. This is useful for scenarios like
// transferring a node to a different user when re-authenticating with a // transferring a node to a different user when re-authenticating with a
// different user's auth key. // different user's auth key.
// If multiple nodes exist with the same machine key (different users), the //
// first one found is returned (order is not guaranteed). // When more than one node shares the machine key (e.g. a tagged node and a
// user-owned node), the choice is deterministic: the tagged node is preferred,
// otherwise the node with the lowest ID. A stable pick keeps re-auth branch
// selection (convert vs create) from depending on map iteration order.
func (s *NodeStore) GetNodeByMachineKeyAnyUser(machineKey key.MachinePublic) (types.NodeView, bool) { func (s *NodeStore) GetNodeByMachineKeyAnyUser(machineKey key.MachinePublic) (types.NodeView, bool) {
timer := prometheus.NewTimer(nodeStoreOperationDuration.WithLabelValues("get_by_machine_key_any_user")) timer := prometheus.NewTimer(nodeStoreOperationDuration.WithLabelValues("get_by_machine_key_any_user"))
defer timer.ObserveDuration() defer timer.ObserveDuration()
@@ -817,14 +820,20 @@ func (s *NodeStore) GetNodeByMachineKeyAnyUser(machineKey key.MachinePublic) (ty
nodeStoreOperations.WithLabelValues("get_by_machine_key_any_user").Inc() nodeStoreOperations.WithLabelValues("get_by_machine_key_any_user").Inc()
snapshot := s.data.Load() snapshot := s.data.Load()
var best types.NodeView
if userMap, exists := snapshot.nodesByMachineKey[machineKey]; exists { if userMap, exists := snapshot.nodesByMachineKey[machineKey]; exists {
// Return the first node found (order not guaranteed due to map iteration)
for _, node := range userMap { for _, node := range userMap {
return node, true if !best.Valid() ||
(node.IsTagged() && !best.IsTagged()) ||
(node.IsTagged() == best.IsTagged() && node.ID() < best.ID()) {
best = node
}
} }
} }
return types.NodeView{}, false return best, best.Valid()
} }
// DebugString returns debug information about the [NodeStore]. // DebugString returns debug information about the [NodeStore].
+56
View File
@@ -1321,3 +1321,59 @@ func TestRebuildPeerMapsWithChangedPeersFunc(t *testing.T) {
assert.Equal(t, 1, peers1.Len(), "ListPeers for node1 should return 1") assert.Equal(t, 1, peers1.Len(), "ListPeers for node1 should return 1")
assert.Equal(t, 1, peers2.Len(), "ListPeers for node2 should return 1") assert.Equal(t, 1, peers2.Len(), "ListPeers for node2 should return 1")
} }
// TestGetNodeByMachineKeyAnyUserDeterministic ensures the any-user machine-key
// lookup returns a stable, well-defined node when more than one node shares a
// machine key: the tagged node if present, otherwise the lowest node ID. A
// nondeterministic pick makes re-auth branch choice (convert vs create) depend
// on map iteration order.
func TestGetNodeByMachineKeyAnyUserDeterministic(t *testing.T) {
mk := key.NewMachine().Public()
assertStablePick := func(t *testing.T, store *NodeStore, want types.NodeID) {
t.Helper()
for range 50 {
got, ok := store.GetNodeByMachineKeyAnyUser(mk)
require.True(t, ok)
require.Equal(t, want, got.ID(), "pick must be stable and well-defined")
}
}
t.Run("prefers lowest node ID", func(t *testing.T) {
store := NewNodeStore(nil, allowAllPeersFunc, TestBatchSize, TestBatchTimeout)
store.Start()
defer store.Stop()
n2 := createTestNode(2, 2, "user2", "node2")
n2.MachineKey = mk
n1 := createTestNode(1, 1, "user1", "node1")
n1.MachineKey = mk
store.PutNode(n2)
store.PutNode(n1)
assertStablePick(t, store, 1)
})
t.Run("prefers tagged node", func(t *testing.T) {
store := NewNodeStore(nil, allowAllPeersFunc, TestBatchSize, TestBatchTimeout)
store.Start()
defer store.Stop()
owned := createTestNode(1, 1, "user1", "node1")
owned.MachineKey = mk
tagged := createTestNode(3, 3, "user3", "node3")
tagged.MachineKey = mk
tagged.UserID = nil
tagged.User = nil
tagged.Tags = []string{"tag:foo"}
store.PutNode(owned)
store.PutNode(tagged)
assertStablePick(t, store, 3)
})
}