diff --git a/hscontrol/state/auth_tagged_expiry_test.go b/hscontrol/state/auth_tagged_expiry_test.go index 77159c6c2..64c1ec75f 100644 --- a/hscontrol/state/auth_tagged_expiry_test.go +++ b/hscontrol/state/auth_tagged_expiry_test.go @@ -1,10 +1,15 @@ package state import ( + "context" + "database/sql" "fmt" + "net/netip" "runtime" + "slices" "strings" "sync" + "sync/atomic" "testing" "time" @@ -2234,8 +2239,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 +2279,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 +2341,379 @@ 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)) +} + +// TestPAKReregisterRollbackKeepsSameBatchWrites: admin writes queued behind a +// failing re-registration share its NodeStore batch, so the rollback must +// compare against what the registration wrote, not the node after the batch. +func TestPAKReregisterRollbackKeepsSameBatchWrites(t *testing.T) { + s := newRetagTestState(t) + user := s.CreateUserForTest("rollback-user") + _, err := s.SetPolicy([]byte(`{"tagOwners":{ + "tag:prod":["rollback-user@"],"tag:quarantine":["rollback-user@"]}}`)) + require.NoError(t, err) + originalKey, err := s.CreatePreAuthKey(user.TypedID(), true, false, nil, nil) + require.NoError(t, err) + a := registerForPAKRereg(t, s, originalKey) + b := registerForPAKRereg(t, s, originalKey) + + // Both rows are read unused; a spends the key, so b fails in the database. + sharedKey, err := s.CreatePreAuthKey(nil, false, false, nil, []string{"tag:prod"}) + require.NoError(t, err) + pakA, err := s.GetPreAuthKeyByID(sharedKey.ID) + require.NoError(t, err) + pakB, err := s.GetPreAuthKeyByID(sharedKey.ID) + require.NoError(t, err) + + reqA := a.regReq + reqA.Auth = &tailcfg.RegisterResponseAuth{AuthKey: sharedKey.Key} + _, err = a.reregister(a.node, pakA, reqA) + require.NoError(t, err) + + reqB := b.regReq + reqB.Auth = &tailcfg.RegisterResponseAuth{AuthKey: sharedKey.Key} + reqB.NodeKey = key.NewNode().Public() + + // Hold the writer so the registration and both admin writes queue behind it + // and land in one batch, in this order. + entered, release := make(chan struct{}), make(chan struct{}) + + releaseOnce := sync.OnceFunc(func() { close(release) }) + defer releaseOnce() + + go s.nodeStore.UpdateNode(a.node.ID(), func(*types.Node) { + close(entered) + <-release + }) + + <-entered + + regDone := make(chan error, 1) + + go func() { + _, err := b.reregister(b.node, pakB, reqB) + regDone <- err + }() + + waitParkedOnWriteQueue(t, "reregisterNodeWithPAK") + + tagsDone := make(chan error, 1) + + go func() { + _, _, err := s.SetNodeTags(b.node.ID(), []string{"tag:quarantine"}) + tagsDone <- err + }() + + waitParkedOnWriteQueue(t, "SetNodeTags") + + expiry := time.Now().Add(-time.Minute) + expiryDone := make(chan error, 1) + + go func() { + _, _, err := s.SetNodeExpiry(b.node.ID(), &expiry) + expiryDone <- err + }() + + waitParkedOnWriteQueue(t, "SetNodeExpiry") + + releaseOnce() + + require.ErrorIs(t, <-regDone, types.PAKError("authkey already used")) + require.NoError(t, <-tagsDone) + require.NoError(t, <-expiryDone) + + ns, row := b.fields(t) + assert.Equal(t, []string{"tag:quarantine"}, ns.Tags, "admin tags must survive") + assert.Equal(t, expiry.UnixNano(), ns.Expiry, "admin expiry must survive") + assert.Equal(t, b.regReq.NodeKey, ns.NodeKey, "registration key must roll back") + assert.Equal(t, ns, row, "row must match NodeStore") +} + +func TestPAKReregisterRollbackKeepsConcurrentWrites(t *testing.T) { + for _, update := range []string{"persisted_grant", "admin", "map"} { + t.Run(update, func(t *testing.T) { + dbPath, s, _ := persistTestSetup(t) + t.Cleanup(func() { _ = s.Close() }) + + user := s.CreateUserForTest("rollback-user") + _, err := s.SetPolicy([]byte(`{"tagOwners":{ + "tag:prod":["rollback-user@"],"tag:quarantine":["rollback-user@"]}}`)) + require.NoError(t, err) + originalKey, err := s.CreatePreAuthKey(user.TypedID(), true, false, nil, nil) + require.NoError(t, err) + a := registerForPAKRereg(t, s, originalKey) + b := registerForPAKRereg(t, s, originalKey) + sharedKey, err := s.CreatePreAuthKey(nil, false, false, nil, []string{"tag:prod"}) + require.NoError(t, err) + pakA, err := s.GetPreAuthKeyByID(sharedKey.ID) + require.NoError(t, err) + pakB, err := s.GetPreAuthKeyByID(sharedKey.ID) + require.NoError(t, err) + require.False(t, pakA.Used) + require.False(t, pakB.Used) + + reqA := a.regReq + reqA.Auth = &tailcfg.RegisterResponseAuth{AuthKey: sharedKey.Key} + _, err = a.reregister(a.node, pakA, reqA) + require.NoError(t, err) + + reqB := b.regReq + reqB.Auth = &tailcfg.RegisterResponseAuth{AuthKey: sharedKey.Key} + reqB.NodeKey = key.NewNode().Public() + want := pakFieldsOf(b.node.AsStruct()) + expiry := time.Now().Add(-time.Minute) + endpoint := netip.MustParseAddrPort("192.0.2.1:1234") + + var pending []<-chan error + + pool := s.db.DB.ConnPool + barrier := &beforeNodeTransaction{ConnPool: pool, before: func() { + published, ok := s.GetNodeByID(b.node.ID()) + require.True(t, ok) + require.Equal(t, []string{"tag:prod"}, published.Tags().AsSlice()) + // Another persist can write the published grant before registration + // takes persistMu. Reproduce that row at the transaction boundary. + require.NoError(t, s.db.DB.Select(nodeUpdateColumns).Updates(published.AsStruct()).Error) + + switch update { + case "admin": + tagsDone := make(chan error, 1) + pending = append(pending, tagsDone) + + go func() { + _, _, err := s.SetNodeTags(b.node.ID(), []string{"tag:quarantine"}) + tagsDone <- err + }() + + require.Eventually(t, func() bool { + n, ok := s.GetNodeByID(b.node.ID()) + return ok && slices.Equal(n.Tags().AsSlice(), []string{"tag:quarantine"}) + }, 5*time.Second, time.Millisecond) + + expiryDone := make(chan error, 1) + pending = append(pending, expiryDone) + + go func() { + _, _, err := s.SetNodeExpiry(b.node.ID(), &expiry) + expiryDone <- err + }() + + require.Eventually(t, func() bool { + n, ok := s.GetNodeByID(b.node.ID()) + return ok && n.IsExpired() + }, 5*time.Second, time.Millisecond) + + want.Tags = []string{"tag:quarantine"} + want.UserID = nil + want.Expiry = expiry.UnixNano() + case "map": + mapDone := make(chan error, 1) + pending = append(pending, mapDone) + + go func() { + _, err := s.UpdateNodeFromMapRequest(b.node.ID(), tailcfg.MapRequest{ + NodeKey: reqB.NodeKey, + Endpoints: []netip.AddrPort{endpoint}, + Hostinfo: &tailcfg.Hostinfo{ + Hostname: "current-host", + NetInfo: &tailcfg.NetInfo{PreferredDERP: 7}, + }, + }) + mapDone <- err + }() + + require.Eventually(t, func() bool { + n, ok := s.GetNodeByID(b.node.ID()) + return ok && n.Hostname() == "current-host" + }, 5*time.Second, time.Millisecond) + + want.Hostname = "current-host" + } + }} + 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() + + _, err = b.reregister(b.node, pakB, reqB) + require.ErrorIs(t, err, types.PAKError("authkey already used")) + + for _, done := range pending { + require.NoError(t, <-done) + } + + restore() + + ns, row := b.fields(t) + assert.Equal(t, want, ns, "rollback must preserve newer writes") + assert.Equal(t, want, row, "rollback must repair the persisted grant") + + if update == "map" { + current, ok := s.GetNodeByID(b.node.ID()) + require.True(t, ok) + require.True(t, current.Hostinfo().NetInfo().Valid()) + assert.Equal(t, tailcfg.DERPRegionID(7), current.Hostinfo().NetInfo().PreferredDERP()) + assert.Equal(t, []netip.AddrPort{endpoint}, current.Endpoints().AsSlice()) + + stored, err := s.db.GetNodeByID(b.node.ID()) + require.NoError(t, err) + assert.Equal(t, "current-host", stored.Hostinfo.Hostname) + assert.Equal(t, tailcfg.DERPRegionID(7), stored.Hostinfo.NetInfo.PreferredDERP) + assert.Equal(t, []netip.AddrPort{endpoint}, stored.Endpoints.List()) + } + + require.NoError(t, s.Close()) + reopened := persistTestReopen(t, dbPath) + reloaded, ok := reopened.GetNodeByID(b.node.ID()) + require.True(t, ok) + assert.Equal(t, want, pakFieldsOf(reloaded.AsStruct()), "restart must keep the rollback") + }) + } +} + // 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..6c8ee8cf8 100644 --- a/hscontrol/state/state.go +++ b/hscontrol/state/state.go @@ -17,6 +17,7 @@ import ( "fmt" "maps" "net/netip" + "reflect" "slices" "strconv" "strings" @@ -105,9 +106,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 +675,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 +701,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,11 +717,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 } @@ -3242,9 +3248,11 @@ func (s *State) reregisterNodeWithPAK( } // prior is the node the writer replaced: the NetInfo source and the - // rollback target if the database write below fails. + // rollback target if the database write below fails. written is what this + // write left; the view after the batch can hold later writes. var ( prior *types.Node + written *types.Node authErr error reauth bool ) @@ -3255,6 +3263,7 @@ func (s *State) reregisterNodeWithPAK( now := time.Now() prior = node.Clone() + defer func() { written = node.Clone() }() reauth, authErr = pakReregister(node.View(), pak, regReq.NodeKey, now) if authErr != nil { @@ -3369,54 +3378,35 @@ 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. - // Sessions, endpoints and health may have moved since that snapshot; - // preserve them. LastSeen stays because the node did contact us. - s.nodeStore.UpdateNode(prior.ID, func(n *types.Node) { - n.NodeKey = prior.NodeKey - n.Hostname = prior.Hostname - n.Hostinfo = prior.Hostinfo - n.RegisterMethod = prior.RegisterMethod - n.Tags = prior.Tags - n.UserID = prior.UserID - n.User = prior.User - n.Expiry = prior.Expiry - n.AuthKey = prior.AuthKey - n.AuthKeyID = prior.AuthKeyID - }) + s.undoReregistration(prior, written) return types.NodeView{}, false, fmt.Errorf("writing node to database: %w", err) } @@ -3433,6 +3423,56 @@ func (s *State) reregisterNodeWithPAK( return updatedNodeView, onlyKeysChanged(before, updatedNodeView, time.Now()), nil } +// undoReregistration puts back the fields of a failed re-registration that +// still hold what it wrote, so later admin and map writes survive, and saves +// the result. LastSeen stays because the node did contact us; sessions, +// endpoints and health are not rolled back. +func (s *State) undoReregistration(prior, written *types.Node) { + restored, exists := s.nodeStore.UpdateNode(prior.ID, func(n *types.Node) { + if n.NodeKey == written.NodeKey { + n.NodeKey = prior.NodeKey + } + + if n.Hostname == written.Hostname { + n.Hostname = prior.Hostname + } + + if reflect.DeepEqual(n.Hostinfo, written.Hostinfo) { + n.Hostinfo = prior.Hostinfo + } + + if n.RegisterMethod == written.RegisterMethod { + n.RegisterMethod = prior.RegisterMethod + } + + if slices.Equal(n.Tags, written.Tags) && + reflect.DeepEqual(n.UserID, written.UserID) && + reflect.DeepEqual(n.User, written.User) { + n.Tags = prior.Tags + n.UserID = prior.UserID + n.User = prior.User + } + + if reflect.DeepEqual(n.Expiry, written.Expiry) { + n.Expiry = prior.Expiry + } + + if reflect.DeepEqual(n.AuthKeyID, written.AuthKeyID) && + reflect.DeepEqual(n.AuthKey, written.AuthKey) { + n.AuthKey = prior.AuthKey + n.AuthKeyID = prior.AuthKeyID + } + }) + if !exists { + return + } + + _, err := s.persistNode(restored) + if err != nil { + log.Error().Err(err).Uint64(zf.NodeID, prior.ID.Uint64()).Msg("Persisting re-registration rollback failed") + } +} + // reauthChange returns the [change.Change] to broadcast after an authentication // that updated or created a node. keysOnly marks a same-user relogin that // changed nothing peers read besides its keys; see [onlyKeysChanged].