From 69c69d35e29c86b40ef1f26313a3e814bce33146 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Sat, 17 Oct 2026 09:00:31 +0000 Subject: [PATCH] state: undo a failed re-registration without reverting newer writes --- hscontrol/state/auth_tagged_expiry_test.go | 234 +++++++++++++++++++++ hscontrol/state/state.go | 73 +++++-- 2 files changed, 290 insertions(+), 17 deletions(-) diff --git a/hscontrol/state/auth_tagged_expiry_test.go b/hscontrol/state/auth_tagged_expiry_test.go index 6658877a9..3885e79a4 100644 --- a/hscontrol/state/auth_tagged_expiry_test.go +++ b/hscontrol/state/auth_tagged_expiry_test.go @@ -4,7 +4,9 @@ import ( "context" "database/sql" "fmt" + "net/netip" "runtime" + "slices" "strings" "sync" "sync/atomic" @@ -2480,6 +2482,238 @@ func TestPersistNodeIncludesCurrentExpiry(t *testing.T) { 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 ec7c91438..28dacdeba 100644 --- a/hscontrol/state/state.go +++ b/hscontrol/state/state.go @@ -17,6 +17,7 @@ import ( "fmt" "maps" "net/netip" + "reflect" "slices" "strconv" "strings" @@ -673,7 +674,6 @@ func (s *State) persistNodeWith( return nil }) - if err != nil { return types.NodeView{}, fmt.Errorf("saving node: %w", err) } @@ -3153,9 +3153,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 ) @@ -3166,6 +3168,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 { @@ -3308,21 +3311,7 @@ func (s *State) reregisterNodeWithPAK( 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) } @@ -3339,6 +3328,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].