state: serialise registration per machine key

Concurrent registrations of one machine key each saw no existing node and
created their own, duplicating nodes and IPs. Hold a per-machine lock across
the find-then-create section.

Updates #3312

hscontrol/state/persist_test.go

@@ -3,6 +3,7 @@ package state
 import (
 	"errors"
 	"net/netip"
+	"sync"
 	"testing"
 	"time"
 
@@ -526,3 +527,58 @@ func TestPreAuthKeyReauthRevertsNodeStoreOnDBFailure(t *testing.T) {
 	require.Equal(t, origNodeKey, got.NodeKey(),
 		"NodeStore must revert to the persisted node key when the write fails")
 }
+
+// TestConcurrentPreAuthKeyRegistrationSameMachineKey ensures concurrent
+// registrations of the same machine key resolve to a single node. Without
+// serialising the find-then-create section, each request sees "no existing
+// node" and creates its own, leaving duplicate nodes and IP allocations for
+// one machine.
+func TestConcurrentPreAuthKeyRegistrationSameMachineKey(t *testing.T) {
+	dbPath := t.TempDir() + "/headscale.db"
+	cfg := persistTestConfig(dbPath)
+
+	s, err := NewState(cfg)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = s.Close() })
+
+	user := s.CreateUserForTest("concurrent-user")
+
+	pak, err := s.CreatePreAuthKey(user.TypedID(), true, false, nil, nil)
+	require.NoError(t, err)
+
+	machineKey := key.NewMachine()
+
+	const n = 12
+
+	var wg sync.WaitGroup
+
+	start := make(chan struct{})
+	errs := make(chan error, n)
+
+	for range n {
+		wg.Go(func() {
+			regReq := tailcfg.RegisterRequest{
+				Auth:     &tailcfg.RegisterResponseAuth{AuthKey: pak.Key},
+				NodeKey:  key.NewNode().Public(),
+				Hostinfo: &tailcfg.Hostinfo{Hostname: "concurrent-node"},
+				Expiry:   time.Now().Add(24 * time.Hour),
+			}
+
+			<-start
+
+			_, _, err := s.HandleNodeFromPreAuthKey(regReq, machineKey.Public())
+			errs <- err
+		})
+	}
+
+	close(start)
+	wg.Wait()
+	close(errs)
+
+	for err := range errs {
+		require.NoError(t, err)
+	}
+
+	require.Equal(t, 1, s.ListNodes().Len(),
+		"concurrent registrations of one machine key must yield a single node")
+}

hscontrol/state/state.go

@@ -32,6 +32,7 @@ import (
 	"github.com/juanfont/headscale/hscontrol/types/change"
 	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/juanfont/headscale/hscontrol/util/zlog/zf"
+	"github.com/puzpuzpuz/xsync/v4"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 	"gorm.io/gorm"
@@ -179,6 +180,22 @@ type State struct {
 	// persistNodeToDB so the database row always converges on [NodeStore]
 	// rather than being clobbered by a stale caller snapshot.
 	persistMu sync.Mutex
+
+	// 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]
+}
+
+// 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()
+
+	return mu.Unlock
 }
 
 // NewState creates and initializes a new [State] instance, setting up the database,
@@ -272,7 +289,8 @@ func NewState(cfg *types.Config) (*State, error) {
 		nodeStore: nodeStore,
 		pings:     newPingTracker(),
 
-		sshCheckAuth: make(map[sshCheckPair]time.Time),
+		sshCheckAuth:  make(map[sshCheckPair]time.Time),
+		registerLocks: xsync.NewMap[key.MachinePublic, *sync.Mutex](),
 	}, nil
 }
 
@@ -2028,6 +2046,11 @@ func (s *State) HandleNodeFromAuthPath(
 
 	// Lookup existing nodes
 	machineKey := regData.MachineKey
+
+	// Serialise registration for this machine so concurrent auth callbacks
+	// resolve to a single node rather than racing the find-then-create section.
+	defer s.lockRegistration(machineKey)()
+
 	existingNodeSameUser, _ := s.nodeStore.GetNodeByMachineKey(machineKey, types.UserID(user.ID))
 	existingNodeAnyUser, _ := s.nodeStore.GetNodeByMachineKeyAnyUser(machineKey)
 
@@ -2204,6 +2227,10 @@ func (s *State) HandleNodeFromPreAuthKey(
 	regReq tailcfg.RegisterRequest,
 	machineKey key.MachinePublic,
 ) (types.NodeView, change.Change, error) {
+	// Serialise registration for this machine so concurrent restarts resolve
+	// to a single node rather than racing the find-then-create section.
+	defer s.lockRegistration(machineKey)()
+
 	pak, err := s.GetPreAuthKey(regReq.Auth.AuthKey)
 	if err != nil {
 		return types.NodeView{}, change.Change{}, err