diff --git a/hscontrol/db/preauth_keys.go b/hscontrol/db/preauth_keys.go index 0d0eba13..d2fb0265 100644 --- a/hscontrol/db/preauth_keys.go +++ b/hscontrol/db/preauth_keys.go @@ -331,11 +331,22 @@ func (hsdb *HSDatabase) DeletePreAuthKey(id uint64) error { }) } -// UsePreAuthKey marks a PreAuthKey as used. +// UsePreAuthKey atomically marks a PreAuthKey as used. The UPDATE is +// guarded by `used = false` so two concurrent registrations racing for +// the same single-use key cannot both succeed: the first commits and +// the second returns PAKError("authkey already used"). Without the +// guard the previous code (Update("used", true) with no WHERE) would +// silently let both transactions claim the key. func UsePreAuthKey(tx *gorm.DB, k *types.PreAuthKey) error { - err := tx.Model(k).Update("used", true).Error - if err != nil { - return fmt.Errorf("updating key used status in database: %w", err) + res := tx.Model(&types.PreAuthKey{}). + Where("id = ? AND used = ?", k.ID, false). + Update("used", true) + if res.Error != nil { + return fmt.Errorf("updating key used status in database: %w", res.Error) + } + + if res.RowsAffected == 0 { + return types.PAKError("authkey already used") } k.Used = true diff --git a/hscontrol/db/preauth_keys_test.go b/hscontrol/db/preauth_keys_test.go index 2f28d449..25bec17e 100644 --- a/hscontrol/db/preauth_keys_test.go +++ b/hscontrol/db/preauth_keys_test.go @@ -11,6 +11,7 @@ import ( "github.com/juanfont/headscale/hscontrol/util" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "gorm.io/gorm" ) func TestCreatePreAuthKey(t *testing.T) { @@ -444,3 +445,45 @@ func TestMultipleLegacyKeysAllowed(t *testing.T) { require.Error(t, err, "duplicate non-empty prefix should be rejected") assert.Contains(t, err.Error(), "UNIQUE constraint failed", "should fail with UNIQUE constraint error") } + +// TestUsePreAuthKeyAtomicCAS verifies that UsePreAuthKey is an atomic +// compare-and-set: a second call against an already-used key reports +// PAKError("authkey already used") rather than silently succeeding. +func TestUsePreAuthKeyAtomicCAS(t *testing.T) { + db, err := newSQLiteTestDB() + require.NoError(t, err) + + user, err := db.CreateUser(types.User{Name: "atomic-cas"}) + require.NoError(t, err) + + pakNew, err := db.CreatePreAuthKey(user.TypedID(), false /* reusable */, false, nil, nil) + require.NoError(t, err) + + pak, err := db.GetPreAuthKey(pakNew.Key) + require.NoError(t, err) + require.False(t, pak.Reusable, "test sanity: key must be single-use") + + // First Use should commit cleanly. + err = db.Write(func(tx *gorm.DB) error { + return UsePreAuthKey(tx, pak) + }) + require.NoError(t, err, "first UsePreAuthKey should succeed") + + // Reload from disk to drop the in-memory Used=true the first call + // set on the struct, simulating a second concurrent transaction + // that loaded the same row before the first one committed. + stale, err := db.GetPreAuthKey(pakNew.Key) + require.NoError(t, err) + + stale.Used = false + + err = db.Write(func(tx *gorm.DB) error { + return UsePreAuthKey(tx, stale) + }) + require.Error(t, err, "second UsePreAuthKey on the same single-use key must fail") + + var pakErr types.PAKError + require.ErrorAs(t, err, &pakErr, + "second UsePreAuthKey error must be a PAKError, got: %v", err) + assert.Equal(t, "authkey already used", pakErr.Error()) +} diff --git a/hscontrol/state/state.go b/hscontrol/state/state.go index 39d7e36a..909bab63 100644 --- a/hscontrol/state/state.go +++ b/hscontrol/state/state.go @@ -2100,7 +2100,13 @@ func (s *State) HandleNodeFromPreAuthKey( return nil, fmt.Errorf("saving node: %w", err) } - if !pak.Reusable { + // 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 UsePreAuthKey would otherwise + // reject it as "authkey already used". This is the path + // behind issue #2830 where containers restart with the + // same one-shot key. + if !pak.Reusable && !pak.Used { err = hsdb.UsePreAuthKey(tx, pak) if err != nil { return nil, fmt.Errorf("using pre auth key: %w", err)