db: reload the credential row after a legacy hash verification

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

@@ -95,6 +95,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 +120,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