From c9e9f728572d5c068f8578a088e327ad378538e9 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Wed, 28 Oct 2026 13:38:50 +0000 Subject: [PATCH] state: drop registration locks once unused The lock is taken before any credential check, so client-chosen machine keys grew the map forever. Ref-count entries; delete at zero. Fixes ANT-2026-ZC18AF82 --- CHANGELOG.md | 1 + hscontrol/state/persist_test.go | 57 +++++++++++++++++++++++++++++++++ hscontrol/state/state.go | 49 ++++++++++++++++++++++++---- 3 files changed, 100 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index daac7c3cc..36f123641 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -165,6 +165,7 @@ the ACME DNS-01 challenge through a built-in DNS provider: `acmedns`, - A stalled request body, or a Noise tunnel with no open request, is closed after 30 seconds - A registration follow-up from a machine other than the one that started the registration is refused before it waits - Only the newest follow-up waits on a pending registration or SSH check; older ones are answered at once +- Registration locks are freed once no registration uses them ## 0.29.7 (202x-xx-xx) diff --git a/hscontrol/state/persist_test.go b/hscontrol/state/persist_test.go index 182e9e8af..4394d1173 100644 --- a/hscontrol/state/persist_test.go +++ b/hscontrol/state/persist_test.go @@ -4,6 +4,7 @@ import ( "encoding/json" "errors" "net/netip" + "runtime" "slices" "strings" "sync" @@ -797,6 +798,62 @@ func TestConcurrentPreAuthKeyRegistrationSameMachineKey(t *testing.T) { require.Equal(t, 1, s.ListNodes().Len(), "concurrent registrations of one machine key must yield a single node") + require.Zero(t, s.registerLocks.Size(), + "registration locks must not outlive the registrations using them") +} + +// TestRejectedRegistrationLeavesNoRegisterLock pins that a refused +// registration keeps nothing. The lock is taken before the pre-auth key is +// checked, so a client without a credential reaches it with as many machine +// keys as it likes. +func TestRejectedRegistrationLeavesNoRegisterLock(t *testing.T) { + s, err := NewState(persistTestConfig(t.TempDir() + "/headscale.db")) + require.NoError(t, err) + t.Cleanup(func() { _ = s.Close() }) + + for range 100 { + _, _, err := s.HandleNodeFromPreAuthKey(tailcfg.RegisterRequest{ + Auth: &tailcfg.RegisterResponseAuth{AuthKey: "not-a-key"}, + NodeKey: key.NewNode().Public(), + }, key.NewMachine().Public()) + require.Error(t, err) + } + + require.Zero(t, s.registerLocks.Size()) +} + +// TestLockRegistrationSerialisesSameMachineKey pins that dropping unused +// locks never lets two registrations of one machine key run at once. +func TestLockRegistrationSerialisesSameMachineKey(t *testing.T) { + s, err := NewState(persistTestConfig(t.TempDir() + "/headscale.db")) + require.NoError(t, err) + t.Cleanup(func() { _ = s.Close() }) + + machineKey := key.NewMachine().Public() + + var ( + wg sync.WaitGroup + holders atomic.Int32 + overlap atomic.Bool + ) + + for range 64 { + wg.Go(func() { + defer s.lockRegistration(machineKey)() + + if holders.Add(1) != 1 { + overlap.Store(true) + } + + runtime.Gosched() + holders.Add(-1) + }) + } + + wg.Wait() + + require.False(t, overlap.Load(), "two registrations held the lock at once") + require.Zero(t, s.registerLocks.Size()) } // TestUpdatePolicyManagerUsersUnchangedKeepsSnapshot ensures re-sending the diff --git a/hscontrol/state/state.go b/hscontrol/state/state.go index b19f55900..a48a53699 100644 --- a/hscontrol/state/state.go +++ b/hscontrol/state/state.go @@ -209,18 +209,53 @@ type State struct { // registerLocks serialises registration per machine key so concurrent // registrations of the same machine resolve to a single node instead of // racing the find-then-create section and each creating their own. - // ponytail: entries are never pruned; bounded by distinct machine keys - // seen, add cleanup on node delete only if it ever matters. - registerLocks *xsync.Map[key.MachinePublic, *sync.Mutex] + // An entry lives only while a registration holds or waits for it: the + // lock is taken before the request is authenticated, and the client + // chooses the machine key. + registerLocks *xsync.Map[key.MachinePublic, *registerLock] +} + +// registerLock is the registration lock of one machine key. refs counts the +// holder and the waiters, and is only changed inside [xsync.Map.Compute]. +type registerLock struct { + mu sync.Mutex + refs int } // lockRegistration serialises registration for a single machine key and // returns the unlock function. func (s *State) lockRegistration(machineKey key.MachinePublic) func() { - mu, _ := s.registerLocks.LoadOrStore(machineKey, &sync.Mutex{}) - mu.Lock() + // Count this registration before blocking on the mutex, so the entry + // cannot be deleted and replaced while it waits. + lock, _ := s.registerLocks.Compute( + machineKey, + func(existing *registerLock, loaded bool) (*registerLock, xsync.ComputeOp) { + if !loaded { + existing = ®isterLock{} + } - return mu.Unlock + existing.refs++ + + return existing, xsync.UpdateOp + }, + ) + lock.mu.Lock() + + return func() { + lock.mu.Unlock() + + s.registerLocks.Compute( + machineKey, + func(existing *registerLock, _ bool) (*registerLock, xsync.ComputeOp) { + existing.refs-- + if existing.refs == 0 { + return existing, xsync.DeleteOp + } + + return existing, xsync.CancelOp + }, + ) + } } // NewState creates and initializes a new [State] instance, setting up the database, @@ -338,7 +373,7 @@ func NewState(cfg *types.Config) (*State, error) { pings: newPingTracker(), sshCheckAuth: make(map[sshCheckPair]time.Time), - registerLocks: xsync.NewMap[key.MachinePublic, *sync.Mutex](), + registerLocks: xsync.NewMap[key.MachinePublic, *registerLock](), } // Surface nodes whose stored data would break map generation (e.g. an