db: guard UsePreAuthKey with WHERE used=false

Add a row-level check so concurrent registrations with the same
single-use key cannot both succeed. Skip the call on
re-registration where the key is already marked used (#2830).

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

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())
+}

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)