state: serialise Connect/Disconnect with per-node lock

The generation check added in b09af384 closed the most obvious
window for stale Disconnect calls but left a residual TOCTOU:
the read of the gen counter was not atomic with the subsequent
NodeStore.UpdateNode and primaryRoutes.SetRoutes mutations. A
concurrent Connect that bumped the gen and rewrote primary
routes could be silently overtaken by a stale Disconnect's
SetRoutes(empty), leaving the node online in NodeStore but with
no primary route — the persistent broken state described in
the issue.

Replace the bare sync.Map[NodeID]*atomic.Uint64 with a
nodeLock pairing the counter with a sync.Mutex. Connect bumps
gen and runs its mutations under the lock; Disconnect acquires
the lock first, re-checks gen, and only proceeds on a match.

Per-node, so unrelated nodes do not contend; same-node steady
state has near-zero contention (one Connect at session start,
one Disconnect ~10s after grace period). Holding the lock
across persistNodeToDB is safe — that path goes through
nodeStore.RebuildPeerMaps which queues onto the writer
goroutine but is invoked from the caller, so no re-entrance.

TestConnectDisconnectRace drives the race directly via State,
deterministically reproducing the empty-primary state on main
and passing under the lock.

Fixes #3203

hscontrol/servertest/connect_race_test.go

@@ -0,0 +1,149 @@
+package servertest_test
+
+import (
+	"context"
+	"net/netip"
+	"slices"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/juanfont/headscale/hscontrol/servertest"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"tailscale.com/tailcfg"
+)
+
+// TestConnectDisconnectRace targets the residual TOCTOU window in
+// state.Disconnect: the connectGeneration check at state.go:644 is not
+// atomic with the subsequent NodeStore.UpdateNode and
+// primaryRoutes.SetRoutes calls. A new Connect that runs between the
+// gen check and the mutations can have its effects overwritten by the
+// stale Disconnect's SetRoutes(empty).
+//
+// The poll.go grace-period flow protects against the most common case
+// (RemoveNode + stillConnected). Connect/Disconnect on State directly
+// bypasses that protection and should still leave the state consistent
+// — if it doesn't, that is the bug behind issue #3203.
+//
+// Run with -race to also catch any data race exposed.
+func TestConnectDisconnectRace(t *testing.T) {
+	srv := servertest.NewServer(t)
+	user := srv.CreateUser(t, "race-user")
+
+	route := netip.MustParsePrefix("10.0.0.0/24")
+
+	// Use NewClient to get a node fully registered + Connected via the
+	// real noise/poll path. After this, NodeStore + primaryRoutes already
+	// have the node, and Connect has been called once.
+	//
+	// Only c2 advertises the route. PrimaryRoutes preserves a current
+	// primary across changes (anti-flap, see primary.go), so if both
+	// nodes were advertising, c1 (lower NodeID) would stay primary and
+	// the test could never observe the route slipping out of c2's
+	// PrimaryRoutes — it would never have been there in the first place.
+	c1 := servertest.NewClient(t, srv, "race-r1", servertest.WithUser(user))
+	c2 := servertest.NewClient(t, srv, "race-r2", servertest.WithUser(user))
+
+	c1.WaitForPeers(t, 1, 10*time.Second)
+
+	c2.Direct().SetHostinfo(&tailcfg.Hostinfo{
+		BackendLogID: "servertest-race-r2",
+		Hostname:     "race-r2",
+		RoutableIPs:  []netip.Prefix{route},
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	_ = c2.Direct().SendUpdate(ctx)
+
+	cancel()
+
+	r2ID := findNodeID(t, srv, "race-r2")
+
+	_, ch, err := srv.State().SetApprovedRoutes(r2ID, []netip.Prefix{route})
+	require.NoError(t, err)
+	srv.App.Change(ch)
+
+	// Wait for advertisement + approval to be reflected as a primary
+	// route assignment in PrimaryRoutes; otherwise we'd be racing the
+	// initial steady-state setup, not the Connect/Disconnect window.
+	require.Eventually(t, func() bool {
+		return slices.Contains(srv.State().GetNodePrimaryRoutes(r2ID), route)
+	}, 10*time.Second, 50*time.Millisecond,
+		"primary route should be assigned to r2 before driving the race")
+
+	// Drive the race repeatedly. Each iteration:
+	//   1. Call Connect(id) to obtain a fresh gen — this stands in for
+	//      a session that "owns" the node.
+	//   2. Spawn a goroutine that issues Disconnect(id, gen) — the
+	//      stale deferred disconnect.
+	//   3. Concurrently spawn a goroutine that issues Connect(id) —
+	//      the new session arriving.
+	//   4. After both finish, check the state is consistent: the node
+	//      should be online and primaryRoutes should hold the approved
+	//      route for it.
+	//
+	// The two goroutines synchronise on a barrier so they start
+	// approximately simultaneously, maximising the chance of hitting the
+	// TOCTOU window.
+	const iterations = 100
+
+	for i := range iterations {
+		// Establish a "current session" with a known gen for r2.
+		_, gen := srv.State().Connect(r2ID)
+
+		var wg sync.WaitGroup
+
+		start := make(chan struct{})
+
+		wg.Add(2)
+
+		go func() {
+			defer wg.Done()
+
+			<-start
+
+			_, _ = srv.State().Disconnect(r2ID, gen)
+		}()
+		go func() {
+			defer wg.Done()
+
+			<-start
+
+			_, _ = srv.State().Connect(r2ID)
+		}()
+
+		close(start)
+		wg.Wait()
+
+		// Post-condition: the node should be ONLINE (the new Connect's
+		// effect must dominate, because the stale Disconnect ran with
+		// an older gen and should have been a no-op — or its effects
+		// must not have overtaken the new Connect's writes).
+		nv, ok := srv.State().GetNodeByID(r2ID)
+		if !assert.True(t, ok, "iteration %d: node should exist", i) {
+			continue
+		}
+
+		online, known := nv.IsOnline().GetOk()
+		if !assert.True(t, known, "iteration %d: online status should be known", i) {
+			continue
+		}
+
+		assert.True(t, online,
+			"iteration %d: node should be ONLINE after concurrent Connect+Disconnect (gen=%d)",
+			i, gen)
+
+		// The approved route must still be reflected as a primary for r2.
+		primary := srv.State().GetNodePrimaryRoutes(r2ID)
+		assert.True(t, slices.Contains(primary, route),
+			"iteration %d: r2 should hold primary for %s after concurrent Connect+Disconnect, got %v",
+			i, route, primary)
+
+		if t.Failed() {
+			t.Logf("primaryRoutes state at failure: %s",
+				srv.State().PrimaryRoutes())
+			return
+		}
+	}
+}

hscontrol/state/state.go

@@ -120,6 +120,14 @@ type sshCheckPair struct {
 	Dst types.NodeID
 }
 
+// nodeLock pairs a per-node Connect generation counter with a mutex
+// that serialises Connect/Disconnect transitions for that node.
+// See the State.nodeLocks field for the rationale (issue #3203).
+type nodeLock struct {
+	mu  sync.Mutex
+	gen atomic.Uint64
+}
+
 // State manages Headscale's core state, coordinating between database, policy management,
 // IP allocation, and DERP routing. All methods are thread-safe.
 type State struct {
@@ -149,11 +157,21 @@ type State struct {
 	// primaryRoutes tracks primary route assignments for nodes
 	primaryRoutes *routes.PrimaryRoutes
 
-	// connectGen tracks a per-node monotonic generation counter so stale
-	// Disconnect() calls from old poll sessions are rejected. Connect()
-	// increments the counter and returns the current value; Disconnect()
-	// only proceeds when the generation it carries matches the latest.
-	connectGen sync.Map // types.NodeID → *atomic.Uint64
+	// nodeLocks pairs a per-node monotonic Connect generation counter with
+	// a mutex that serialises Connect/Disconnect transitions for that
+	// node. Connect bumps gen and rewrites primaryRoutes for id;
+	// Disconnect re-checks gen UNDER the lock and only proceeds when its
+	// captured generation still matches. This closes the TOCTOU window
+	// between the gen check and the NodeStore + primaryRoutes mutations
+	// that allowed a stale Disconnect to overwrite a fresh Connect's
+	// primary route assignment (issue #3203).
+	//
+	// TODO(#3203 follow-up): if a similar race is found in
+	// SetApprovedRoutes, UpdateNodeFromMapRequest's route path, or
+	// SetNodeHealthy, those mutators should acquire this same per-node
+	// lock. Each requires its own reproduction test before being brought
+	// under the lock.
+	nodeLocks sync.Map // types.NodeID → *nodeLock
 
 	// pings tracks pending ping requests and their response channels.
 	pings *pingTracker
@@ -563,11 +581,22 @@ func (s *State) DeleteNode(node types.NodeView) (change.Change, error) {
 // It returns the list of changes and a generation number. The generation number
 // must be passed to Disconnect() so that stale disconnects from old poll sessions
 // are rejected (see the grace period logic in poll.go).
+//
+// The per-node lock is held across the entire body so that Connect's
+// gen bump, NodeStore mutation, and primaryRoutes.SetRoutes are atomic
+// with respect to a concurrent Disconnect on the same node. Without
+// this, a stale Disconnect that passed its gen check could overwrite
+// the fresh Connect's online status and approved-route assignment
+// (issue #3203).
 func (s *State) Connect(id types.NodeID) ([]change.Change, uint64) {
+	nl := s.nodeLockFor(id)
+	nl.mu.Lock()
+	defer nl.mu.Unlock()
+
 	// Increment the connect generation for this node. This ensures that any
 	// in-flight Disconnect() from a previous session will see a stale generation
 	// and become a no-op.
-	gen := s.nextConnectGen(id)
+	gen := nl.gen.Add(1)
 
 	// Update online status in NodeStore before creating change notification
 	// so the NodeStore already reflects the correct state when other nodes
@@ -596,41 +625,22 @@ func (s *State) Connect(id types.NodeID) ([]change.Change, uint64) {
 		c = append(c, change.NodeAdded(id))
 	}
 
-	// Mirror Disconnect: a node coming online may (re)enable cap/relay
-	// grants targeting it, reintroduce identity-based aliases that
-	// resolve to its tags/IPs, and so on. Always trigger a PolicyChange
-	// so peers can recompute their netmap and pick up any policy
-	// elements that depend on this node being present.
+	// Coming online may re-enable cap/relay grants and identity-based
+	// aliases targeting this node, so peers need a fresh netmap.
 	c = append(c, change.PolicyChange())
 
 	return c, gen
 }
 
-// nextConnectGen atomically increments and returns the connect generation for a node.
-func (s *State) nextConnectGen(id types.NodeID) uint64 {
-	val, _ := s.connectGen.LoadOrStore(id, &atomic.Uint64{})
-
-	counter, ok := val.(*atomic.Uint64)
-	if !ok {
-		return 0
+// nodeLockFor returns the per-node lock for id, creating it on first
+// use. The returned *nodeLock is stable for the lifetime of the State —
+// callers may cache the reference across operations on the same node.
+func (s *State) nodeLockFor(id types.NodeID) *nodeLock {
+	if v, ok := s.nodeLocks.Load(id); ok {
+		return v.(*nodeLock)
 	}
-
-	return counter.Add(1)
-}
-
-// connectGeneration returns the current connect generation for a node.
-func (s *State) connectGeneration(id types.NodeID) uint64 {
-	val, ok := s.connectGen.Load(id)
-	if !ok {
-		return 0
-	}
-
-	counter, ok := val.(*atomic.Uint64)
-	if !ok {
-		return 0
-	}
-
-	return counter.Load()
+	v, _ := s.nodeLocks.LoadOrStore(id, &nodeLock{})
+	return v.(*nodeLock)
 }
 
 // Disconnect marks a node as disconnected and updates its primary routes in the state.
@@ -638,10 +648,22 @@ func (s *State) connectGeneration(id types.NodeID) uint64 {
 // been called since the session that is disconnecting, the generation will not match
 // and this call becomes a no-op, preventing stale disconnects from overwriting the
 // online status set by a newer session.
+//
+// The per-node lock is acquired before the gen check so that the check
+// and the subsequent NodeStore + primaryRoutes mutations are atomic
+// with respect to a concurrent Connect on the same node. A stale
+// Disconnect that observed an out-of-date gen must not be allowed to
+// proceed even partially — see issue #3203.
 func (s *State) Disconnect(id types.NodeID, gen uint64) ([]change.Change, error) {
+	nl := s.nodeLockFor(id)
+	nl.mu.Lock()
+	defer nl.mu.Unlock()
+
 	// Check if this disconnect is stale. A newer Connect() will have incremented
 	// the generation, so if ours doesn't match, a newer session owns this node.
-	if current := s.connectGeneration(id); current != gen {
+	// The check happens UNDER the lock so a concurrent Connect cannot bump gen
+	// between this read and the mutations below.
+	if current := nl.gen.Load(); current != gen {
 		log.Debug().
 			Uint64("disconnect_gen", gen).
 			Uint64("current_gen", current).
@@ -653,7 +675,6 @@ func (s *State) Disconnect(id types.NodeID, gen uint64) ([]change.Change, error)
 	node, ok := s.nodeStore.UpdateNode(id, func(n *types.Node) {
 		now := time.Now()
 		n.LastSeen = &now
-		// NodeStore is the source of truth for all node state including online status.
 		n.IsOnline = new(false)
 	})
 
@@ -663,12 +684,10 @@ func (s *State) Disconnect(id types.NodeID, gen uint64) ([]change.Change, error)
 
 	log.Info().EmbedObject(node).Msg("node disconnected")
 
-	// Special error handling for disconnect - we log errors but continue
-	// because NodeStore is already updated and we need to notify peers
+	// Persist LastSeen best-effort: NodeStore already reflects offline
+	// and peers still need the change notifications below.
 	_, c, err := s.persistNodeToDB(node)
 	if err != nil {
-		// Log error but don't fail the disconnection - NodeStore is already updated
-		// and we need to send change notifications to peers
 		log.Error().Err(err).EmbedObject(node).Msg("failed to update last seen in database")
 
 		c = change.Change{}