diff --git a/hscontrol/servertest/connect_race_test.go b/hscontrol/servertest/connect_race_test.go new file mode 100644 index 00000000..71ae9ff6 --- /dev/null +++ b/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 + } + } +} diff --git a/hscontrol/state/state.go b/hscontrol/state/state.go index b7bc7432..d6bd27ff 100644 --- a/hscontrol/state/state.go +++ b/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{}