state: write node rows from the NodeStore under persistMu

2026-11-04 04:34:07 +09:00 · 2026-10-30 14:55:36 +00:00
1 parent b4f970f93c
commit 09aba4e488
@@ -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
@@ -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.