diff --git a/hscontrol/state/persist_test.go b/hscontrol/state/persist_test.go index e29de839..95064ea4 100644 --- a/hscontrol/state/persist_test.go +++ b/hscontrol/state/persist_test.go @@ -1,6 +1,7 @@ package state import ( + "errors" "net/netip" "testing" "time" @@ -10,6 +11,7 @@ import ( "github.com/juanfont/headscale/hscontrol/util" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "gorm.io/gorm" "tailscale.com/tailcfg" "tailscale.com/types/key" ) @@ -472,3 +474,55 @@ func TestPreAuthKeyReauthRejectsNodeKeyClaimedByAnotherMachine(t *testing.T) { require.Equal(t, victimMachine.Public(), owner.MachineKey(), "victim's NodeKey index entry must be untouched") } + +var errInjectedNodeUpdate = errors.New("injected node update failure") + +// TestPreAuthKeyReauthRevertsNodeStoreOnDBFailure ensures a failed database +// write during pre-auth-key re-registration does not leave the NodeStore +// holding a node key that was never persisted: a restart would reload the old +// row and the client's current key would no longer resolve, locking it out. +func TestPreAuthKeyReauthRevertsNodeStoreOnDBFailure(t *testing.T) { + dbPath := t.TempDir() + "/headscale.db" + cfg := persistTestConfig(dbPath) + + s, err := NewState(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = s.Close() }) + + user := s.CreateUserForTest("reauth-user") + + pak, err := s.CreatePreAuthKey(user.TypedID(), true, false, nil, nil) + require.NoError(t, err) + + machineKey := key.NewMachine() + regReq := tailcfg.RegisterRequest{ + Auth: &tailcfg.RegisterResponseAuth{AuthKey: pak.Key}, + NodeKey: key.NewNode().Public(), + Hostinfo: &tailcfg.Hostinfo{Hostname: "reauth-node"}, + Expiry: time.Now().Add(24 * time.Hour), + } + node, _, err := s.HandleNodeFromPreAuthKey(regReq, machineKey.Public()) + require.NoError(t, err) + + origNodeKey := node.NodeKey() + + // Fail the node row update so the re-registration's database write errors + // after the NodeStore has already been mutated. + require.NoError(t, s.db.DB.Callback().Update().Before("gorm:update"). + Register("fail_node_update", func(tx *gorm.DB) { + if tx.Statement.Table == "nodes" { + _ = tx.AddError(errInjectedNodeUpdate) + } + })) + + reReg := regReq + reReg.NodeKey = key.NewNode().Public() // rotate -> NodeStore mutation, then DB write fails + _, _, err = s.HandleNodeFromPreAuthKey(reReg, machineKey.Public()) + require.NoError(t, s.db.DB.Callback().Update().Remove("fail_node_update")) + require.Error(t, err, "re-registration must fail when the database write fails") + + got, ok := s.nodeStore.GetNode(node.ID()) + require.True(t, ok) + require.Equal(t, origNodeKey, got.NodeKey(), + "NodeStore must revert to the persisted node key when the write fails") +} diff --git a/hscontrol/state/state.go b/hscontrol/state/state.go index f2803fde..07f1c24e 100644 --- a/hscontrol/state/state.go +++ b/hscontrol/state/state.go @@ -2321,6 +2321,12 @@ func (s *State) HandleNodeFromPreAuthKey( return types.NodeView{}, change.Change{}, ErrNodeKeyInUse } + // Snapshot the pre-update node so the NodeStore can be rolled back if + // the database write below fails. The view points at the immutable + // pre-update snapshot (UpdateNode swaps in a new one), so this stays + // valid after the mutation. + priorNode := existingNodeSameUser.AsStruct() + // Update existing node - NodeStore first, then database updatedNodeView, ok := s.nodeStore.UpdateNode(existingNodeSameUser.ID(), func(node *types.Node) { node.NodeKey = regReq.NodeKey @@ -2401,6 +2407,13 @@ func (s *State) HandleNodeFromPreAuthKey( return nil, nil //nolint:nilnil // intentional: transaction success }) if err != nil { + // The NodeStore was updated before the database write. Roll it back + // so it does not advertise a registration the database rejected + // (e.g. a node key that a restart would not reload). + if priorNode != nil { + s.nodeStore.PutNode(*priorNode) + } + return types.NodeView{}, change.Change{}, fmt.Errorf("writing node to database: %w", err) }