db: bound legacy hash verification and reload verified credentials

Limit the wait for a verification slot and reload the credential after
verification so concurrent revocation or deletion remains effective.

hscontrol/db/oauth_test.go

@@ -128,6 +128,73 @@ func TestVerifySecretConcurrent(t *testing.T) {
 	}
 }
 
+// TODO(kradalby): remove in 0.32 with bcrypt/Argon2id support.
+func TestVerifySecretGivesUpWhenLimiterBusy(t *testing.T) {
+	const secret = "s3cr3t"
+
+	hash, err := bcrypt.GenerateFromPassword([]byte(secret), bcrypt.MinCost)
+	require.NoError(t, err)
+
+	previousWait := legacyHashWait
+	legacyHashWait = 10 * time.Millisecond
+
+	t.Cleanup(func() { legacyHashWait = previousWait })
+
+	for range cap(legacyHashLimiter) {
+		legacyHashLimiter <- struct{}{}
+	}
+
+	release := sync.OnceFunc(func() {
+		for range cap(legacyHashLimiter) {
+			<-legacyHashLimiter
+		}
+	})
+
+	var wg sync.WaitGroup
+
+	defer func() {
+		release()
+		wg.Wait()
+	}()
+
+	_, err = verifySecret(hashSecret(secret), secret)
+	require.NoError(t, err)
+
+	type result struct {
+		needsRehash bool
+		err         error
+	}
+
+	const n = 8
+
+	results := make(chan result, n)
+	for range n {
+		wg.Go(func() {
+			needsRehash, err := verifySecret(hash, secret)
+			results <- result{needsRehash: needsRehash, err: err}
+		})
+	}
+
+	deadline := time.After(time.Second)
+
+	for range n {
+		select {
+		case got := <-results:
+			require.ErrorIs(t, got.err, errLegacyHashBusy)
+			assert.False(t, got.needsRehash)
+		case <-deadline:
+			t.Fatal("legacy verification did not give up while the limiter was full")
+		}
+	}
+
+	release()
+	wg.Wait()
+
+	needsRehash, err := verifySecret(hash, secret)
+	require.NoError(t, err)
+	assert.True(t, needsRehash)
+}
+
 // TestGenerateSecret verifies the unified key shape
 // <prefix><identifier(12)>-<secret(64)> and that the stored hash verifies the
 // secret without needing a rehash.

hscontrol/db/preauth_keys_test.go

@@ -3,6 +3,7 @@ package db
 import (
 	"slices"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -449,3 +450,105 @@ func TestPreAuthKeyLazyRehashesBcrypt(t *testing.T) {
 	_, err = db.GetPreAuthKey(keyStr)
 	require.NoError(t, err)
 }
+
+// TODO(kradalby): remove in 0.32 with bcrypt/Argon2id support.
+func TestPreAuthKeyInvalidatedWhileQueuedOnLegacyLimiter(t *testing.T) {
+	tests := []struct {
+		name       string
+		invalidate func(*HSDatabase, uint64) error
+		wantErr    types.PAKError
+	}{
+		{
+			name:       "expired",
+			invalidate: (*HSDatabase).ExpirePreAuthKey,
+			wantErr:    "authkey expired",
+		},
+		{
+			name:       "revoked",
+			invalidate: (*HSDatabase).RevokePreAuthKey,
+			wantErr:    "authkey revoked",
+		},
+		{name: "deleted", invalidate: (*HSDatabase).DeletePreAuthKey},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			db, err := newSQLiteTestDB()
+			require.NoError(t, err)
+			t.Cleanup(func() { require.NoError(t, db.Close()) })
+
+			user := db.CreateUserForTest("queued-user")
+			prefix := "abcdefghijkl"
+			secret := strings.Repeat("b", 64)
+
+			hash, err := bcrypt.GenerateFromPassword([]byte(secret), bcrypt.MinCost)
+			require.NoError(t, err)
+
+			cred := types.Credential{
+				Kind:       types.CredentialPreAuthKey,
+				Identifier: prefix,
+				Hash:       hash,
+				UserID:     &user.ID,
+				Reusable:   true,
+			}
+			require.NoError(t, db.DB.Create(&cred).Error)
+
+			read := make(chan struct{})
+			signalRead := sync.OnceFunc(func() { close(read) })
+
+			require.NoError(t, db.DB.Callback().Query().After("gorm:query").
+				Register("signal_credential_read", func(tx *gorm.DB) {
+					if tx.Statement.Table == "credentials" {
+						signalRead()
+					}
+				}))
+
+			for range cap(legacyHashLimiter) {
+				legacyHashLimiter <- struct{}{}
+			}
+
+			release := sync.OnceFunc(func() {
+				for range cap(legacyHashLimiter) {
+					<-legacyHashLimiter
+				}
+			})
+
+			var (
+				wg     sync.WaitGroup
+				pak    *types.PreAuthKey
+				getErr error
+			)
+
+			defer func() {
+				release()
+				wg.Wait()
+			}()
+
+			wg.Go(func() {
+				pak, getErr = db.GetPreAuthKey(authKeyPrefix + prefix + "-" + secret)
+			})
+
+			select {
+			case <-read:
+			case <-time.After(5 * time.Second):
+				t.Fatal("credential lookup did not read the row")
+			}
+
+			require.NoError(t, tt.invalidate(db, cred.ID))
+			release()
+			wg.Wait()
+
+			if tt.name == "deleted" {
+				require.ErrorIs(t, getErr, ErrPreAuthKeyNotFound)
+
+				return
+			}
+
+			require.NoError(t, getErr)
+			require.NotNil(t, pak)
+			require.NotNil(t, pak.User)
+			assert.Equal(t, user.ID, pak.User.ID)
+			require.ErrorIs(t, pak.Validate(), tt.wantErr)
+		})
+	}
+}

hscontrol/db/secret.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"runtime"
 	"strings"
+	"time"
 
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/rs/zerolog/log"
@@ -57,7 +58,11 @@ var (
 // only run until every stored hash has been upgraded to SHA-256.
 //
 // TODO(kradalby): remove in 0.32 with bcrypt/Argon2id support.
-var legacyHashLimiter = make(chan struct{}, max(2, runtime.GOMAXPROCS(0)))
+var (
+	legacyHashLimiter = make(chan struct{}, max(2, runtime.GOMAXPROCS(0)))
+	legacyHashWait    = 5 * time.Second
+	errLegacyHashBusy = errors.New("legacy hash verification is busy")
+)
 
 // generateSecret builds a new credential string prefix+identifier+"-"+secret,
 // returning the full string (shown ONCE to the user), the public identifier
@@ -95,6 +100,9 @@ func authenticateCredential(
 		query = query.Preload(p)
 	}
 
+	// Session lets the legacy reload below reuse the preloads.
+	query = query.Session(&gorm.Session{})
+
 	err := query.First(&cred, "kind = ? AND identifier = ?", kind, identifier).Error
 	if errors.Is(err, gorm.ErrRecordNotFound) {
 		return nil, notFound
@@ -117,14 +125,26 @@ func authenticateCredential(
 	if needsRehash {
 		// Best effort: on failure the legacy hash stays and the next
 		// authentication retries.
-		cred.Hash = hashSecret(secret)
-
 		err := tx.Model(&types.Credential{}).Where("id = ?", cred.ID).
-			Update("hash", cred.Hash).Error
+			Update("hash", hashSecret(secret)).Error
 		if err != nil {
 			log.Warn().Err(err).Uint64("credential", cred.ID).
 				Msg("upgrading legacy credential hash")
 		}
+
+		// Validity may change while legacy verification waits for a slot.
+		var current types.Credential
+
+		err = query.First(&current, "id = ?", cred.ID).Error
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, notFound
+		}
+
+		if err != nil {
+			return nil, fmt.Errorf("reloading %s credential %d: %w", kind, cred.ID, err)
+		}
+
+		return &current, nil
 	}
 
 	return &cred, nil
@@ -151,7 +171,12 @@ func verifySecret(encoded []byte, secret string) (bool, error) {
 		return false, nil
 	}
 
-	legacyHashLimiter <- struct{}{}
+	select {
+	case legacyHashLimiter <- struct{}{}:
+	case <-time.After(legacyHashWait):
+		return false, errLegacyHashBusy
+	}
+
 	defer func() { <-legacyHashLimiter }()
 
 	if bytes.HasPrefix(encoded, []byte("$argon2id$")) {