From 09aba4e4887e943caaeaca67a8ca13f6e6d05362 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Sat, 17 Oct 2026 08:53:19 +0000 Subject: [PATCH] state: write node rows from the NodeStore under persistMu --- hscontrol/state/auth_tagged_expiry_test.go | 156 ++++++++++++++++++++- hscontrol/state/state.go | 79 +++++------ 2 files changed, 194 insertions(+), 41 deletions(-) diff --git a/hscontrol/state/auth_tagged_expiry_test.go b/hscontrol/state/auth_tagged_expiry_test.go index 77159c6c2..85241241a 100644 --- a/hscontrol/state/auth_tagged_expiry_test.go +++ b/hscontrol/state/auth_tagged_expiry_test.go @@ -1,10 +1,13 @@ package state import ( + "context" + "database/sql" "fmt" "runtime" "strings" "sync" + "sync/atomic" "testing" "time" @@ -2234,8 +2237,6 @@ func TestPAKReregisterRollbackKeepsPriorWrite(t *testing.T) { // TestPAKReregisterConcurrentAdminExpire races a same-key restart with a spent // key against an admin expiry. Either serial order leaves the node expired. -// Only the NodeStore is checked: the order of the two database writes is not -// serialised with it. func TestPAKReregisterConcurrentAdminExpire(t *testing.T) { s := newRetagTestState(t) user := s.CreateUserForTest("pakuser") @@ -2276,6 +2277,16 @@ func TestPAKReregisterConcurrentAdminExpire(t *testing.T) { n, ok := s.GetNodeByID(c.node.ID()) require.True(t, ok) require.True(t, n.IsExpired(), "iteration %d: the spent key revived an expired node", i) + + stored, err := s.db.GetNodeByID(c.node.ID()) + require.NoError(t, err) + require.True( + t, + stored.IsExpired(), + "iteration %d: the row lost the admin expiry", + i, + ) + require.True(t, n.Expiry().Get().Equal(*stored.Expiry)) } } @@ -2328,6 +2339,147 @@ func TestPreAuthKeyRegistrationOwnsHostinfo(t *testing.T) { } } +// beforeNodeTransaction runs once before opening the database transaction so +// another writer can publish its NodeStore update without sharing a SQL lock. +type beforeNodeTransaction struct { + gorm.ConnPool + + fired atomic.Bool + before func() +} + +func (p *beforeNodeTransaction) BeginTx( + ctx context.Context, + opts *sql.TxOptions, +) (gorm.ConnPool, error) { + if p.fired.CompareAndSwap(false, true) { + p.before() + } + + beginner, ok := p.ConnPool.(gorm.TxBeginner) + if !ok { + return nil, gorm.ErrInvalidTransaction + } + + return beginner.BeginTx(ctx, opts) +} + +func TestReauthPersistenceKeepsConcurrentExpiry(t *testing.T) { + for _, interactive := range []bool{false, true} { + t.Run(fmt.Sprintf("interactive=%t", interactive), func(t *testing.T) { + dbPath, s, _ := persistTestSetup(t) + t.Cleanup(func() { _ = s.Close() }) + + user := s.CreateUserForTest("reauth-persist-user") + k, err := s.CreatePreAuthKey(user.TypedID(), true, false, nil, nil) + require.NoError(t, err) + c := registerForPAKRereg(t, s, k) + expiry := time.Now().Add(-time.Minute) + expireDone := make(chan struct{}) + + var expireErr error + + pool := s.db.DB.ConnPool + barrier := &beforeNodeTransaction{ConnPool: pool, before: func() { + go func() { + defer close(expireDone) + + _, _, expireErr = s.SetNodeExpiry(c.node.ID(), &expiry) + }() + // Expiry either finishes, or waits for the registration's persist lock. + // In both cases its NodeStore update must precede the registration's SQL. + buf := make([]byte, 1<<20) + + require.Eventually(t, func() bool { + select { + case <-expireDone: + return true + default: + } + + stacks := string(buf[:runtime.Stack(buf, true)]) + for g := range strings.SplitSeq(stacks, "\n\n") { + if strings.Contains(g, ".SetNodeExpiry(") && + strings.Contains(g, ".persistNodeWith(") { + return true + } + } + + return false + }, 5*time.Second, time.Millisecond) + }} + s.db.DB.ConnPool = barrier + s.db.DB.Statement.ConnPool = barrier + + restore := func() { + s.db.DB.ConnPool = pool + s.db.DB.Statement.ConnPool = pool + } + defer restore() + + if interactive { + authID := types.MustAuthID() + require.NoError(t, s.SetAuthCacheEntry( + authID, + types.NewRegisterAuthRequest(&types.RegistrationData{ + MachineKey: c.machineKey, + NodeKey: c.regReq.NodeKey, + Hostname: "h1", + Hostinfo: c.regReq.Hostinfo, + Expiry: &c.regReq.Expiry, + }), + )) + _, _, err = s.HandleNodeFromAuthPath( + authID, + types.UserID(user.ID), + nil, + util.RegisterMethodOIDC, + ) + } else { + _, _, err = s.HandleNodeFromPreAuthKey(c.regReq, c.machineKey) + } + + require.NoError(t, err) + <-expireDone + require.NoError(t, expireErr) + restore() + + ns, row := c.fields(t) + require.Equal(t, expiry.UnixNano(), ns.Expiry) + assert.Equal( + t, + ns.Expiry, + row.Expiry, + "the row must retain the concurrent admin expiry", + ) + require.NoError(t, s.Close()) + reopened := persistTestReopen(t, dbPath) + reloaded, ok := reopened.GetNodeByID(c.node.ID()) + require.True(t, ok) + assert.True(t, reloaded.IsExpired(), "a restart must keep the node expired") + }) + } +} + +func TestPersistNodeIncludesCurrentExpiry(t *testing.T) { + _, s, nodeID := persistTestSetup(t) + t.Cleanup(func() { _ = s.Close() }) + + stale, ok := s.GetNodeByID(nodeID) + require.True(t, ok) + + expiry := time.Now().Add(-time.Minute) + _, ok = s.nodeStore.UpdateNode(nodeID, func(n *types.Node) { n.Expiry = &expiry }) + require.True(t, ok) + + _, err := s.persistNode(stale) + require.NoError(t, err) + stored, err := s.db.GetNodeByID(nodeID) + require.NoError(t, err) + require.NotNil(t, stored.Expiry) + require.True(t, expiry.Equal(*stored.Expiry)) +} + // TestUserKeyOnTaggedMachine: a user key on a tagged machine re-authenticates // the tagged node only when it speaks for it (the user an admin tagged it // from); another user's key wakes it, gets a device of its own beside it with diff --git a/hscontrol/state/state.go b/hscontrol/state/state.go index ee7767d30..1b242e09e 100644 --- a/hscontrol/state/state.go +++ b/hscontrol/state/state.go @@ -105,9 +105,6 @@ var ErrNodeNameNotUnique = errors.New("node name is not unique") // PreAuthKey references after a key has been deleted (#2862). // - User: GORM association, not a real column. // - IsOnline: runtime-only field (gorm:"-"). -// -// Expiry is included here but may be omitted at call sites that must -// not touch it (see persistNodeAndRefreshPolicy). var nodeUpdateColumns = []string{ "MachineKey", "NodeKey", @@ -677,6 +674,16 @@ func (s *State) ListAllUsers() ([]types.User, error) { // Batch callers (e.g. autoApproveNodes) use it to write many rows and then // trigger a single policy rebuild instead of one per node. func (s *State) persistNode(node types.NodeView) (types.NodeView, error) { + return s.persistNodeWith(node, nodeUpdateColumns, nil) +} + +// persistNodeWith writes the current node and any related change in one +// transaction. inTx runs inside it, with persistMu held. +func (s *State) persistNodeWith( + node types.NodeView, + columns []string, + inTx func(*gorm.DB) error, +) (types.NodeView, error) { if !node.Valid() { return types.NodeView{}, ErrInvalidNodeView } @@ -693,11 +700,10 @@ func (s *State) persistNode(node types.NodeView) (types.NodeView, error) { // ephemeral logout) between the caller's update and this persist: a missing // node means we must not re-insert it. s.persistMu.Lock() + defer s.persistMu.Unlock() fresh, exists := s.nodeStore.GetNode(node.ID()) if !exists { - s.persistMu.Unlock() - log.Warn(). EmbedObject(node). Bool("is_ephemeral", node.IsEphemeral()). @@ -710,10 +716,18 @@ func (s *State) persistNode(node types.NodeView) (types.NodeView, error) { // Explicitly select all node columns so GORM includes nil/zero-value // fields (e.g. UserID=nil when converting a user-owned node to tagged). - // Omit "Expiry" here: expiry is only updated through explicit - // SetNodeExpiry calls or re-registration, not during MapRequest updates. - err := s.db.DB.Select(nodeUpdateColumns).Omit("Expiry").Updates(nodePtr).Error - s.persistMu.Unlock() + err := s.db.Write(func(tx *gorm.DB) error { + err := tx.Select(columns).Updates(nodePtr).Error + if err != nil { + return err + } + + if inTx != nil { + return inTx(tx) + } + + return nil + }) if err != nil { return types.NodeView{}, fmt.Errorf("saving node: %w", err) @@ -1116,8 +1130,7 @@ func (s *State) SetNodeExpiry(nodeID types.NodeID, expiry *time.Time) (types.Nod recompute = change.PolicyChange() } - // Persist expiry change to database directly since persistNodeAndRefreshPolicy omits expiry. - err := s.db.NodeSetExpiry(nodeID, expiry) + _, err := s.persistNode(n) if err != nil { c := nodeWriteFailed(nodeID, s.policyChangeSince(genBefore).Merge(recompute)) @@ -2341,14 +2354,7 @@ func (s *State) applyAuthNodeUpdate(params authNodeUpdateParams) (types.NodeView updateColumns = append(slices.Clone(nodeUpdateColumns), "AuthKeyID") } - _, err := hsdb.Write(s.db.DB, func(tx *gorm.DB) (*types.Node, error) { - err := tx.Select(updateColumns).Updates(updatedNodeView.AsStruct()).Error - if err != nil { - return nil, fmt.Errorf("saving node: %w", err) - } - - return nil, nil //nolint:nilnil // side-effect only write - }) + _, err := s.persistNodeWith(updatedNodeView, updateColumns, nil) if err != nil { return types.NodeView{}, false, err } @@ -3369,37 +3375,32 @@ func (s *State) reregisterNodeWithPAK( return types.NodeView{}, false, authErr } - _, err := hsdb.Write(s.db.DB, func(tx *gorm.DB) (*types.Node, error) { - // Explicitly select all node columns so GORM includes nil/zero-value fields - // (see nodeUpdateColumns comment). AuthKeyID is excluded on MapRequest - // so a deleted key's stale reference is not persisted, but - // re-registration presents a key just loaded from the database, so - // its ID must be persisted here — otherwise a restart reloads the old - // key and any key-scoped property (e.g. Ephemeral) silently reverts. - // A wake-up keeps the prior key, which may have been deleted (#2862). - reregColumns := nodeUpdateColumns - if reauth { - reregColumns = append(slices.Clone(nodeUpdateColumns), "AuthKeyID") - } - - err := tx.Select(reregColumns).Updates(updatedNodeView.AsStruct()).Error - if err != nil { - return nil, fmt.Errorf("saving node: %w", err) - } + // Explicitly select all node columns so GORM includes nil/zero-value fields + // (see nodeUpdateColumns comment). AuthKeyID is excluded on MapRequest + // so a deleted key's stale reference is not persisted, but + // re-registration presents a key just loaded from the database, so + // its ID must be persisted here — otherwise a restart reloads the old + // key and any key-scoped property (e.g. Ephemeral) silently reverts. + // A wake-up keeps the prior key, which may have been deleted (#2862). + reregColumns := nodeUpdateColumns + if reauth { + reregColumns = append(slices.Clone(nodeUpdateColumns), "AuthKeyID") + } + _, err := s.persistNodeWith(updatedNodeView, reregColumns, func(tx *gorm.DB) error { // Only mark the key used on the *first* registration. On // re-registration the same key is already used and the // atomic compare-and-set in [hsdb.UsePreAuthKey] would otherwise // reject it as "authkey already used", locking out containers // that restart with the same one-shot key. if reauth && !pak.Reusable && !pak.Used { - err = hsdb.UsePreAuthKey(tx, pak) + err := hsdb.UsePreAuthKey(tx, pak) if err != nil { - return nil, fmt.Errorf("using pre auth key: %w", err) + return fmt.Errorf("using pre auth key: %w", err) } } - return nil, nil //nolint:nilnil // intentional: transaction success + return nil }) if err != nil { // Restore the registration fields from the node the writer replaced.