From 10b1e9995d29bdd778df170923cb5d0137f232e0 Mon Sep 17 00:00:00 2001 From: Hintay Date: Thu, 21 May 2026 17:37:02 +0900 Subject: [PATCH] fix: harden recovery migration flow --- api/user/2fa_test.go | 7 +++- app/src/components/TwoFA/use2FAModal.ts | 31 +++++++++----- app/src/pinia/moudule/user.ts | 12 ++++-- internal/user/otp.go | 54 +++++++++++++++---------- internal/user/otp_test.go | 15 +++++++ 5 files changed, 82 insertions(+), 37 deletions(-) diff --git a/api/user/2fa_test.go b/api/user/2fa_test.go index 24e07142..53de01eb 100644 --- a/api/user/2fa_test.go +++ b/api/user/2fa_test.go @@ -2,6 +2,7 @@ package user import ( "fmt" + "net/http/httptest" "testing" "github.com/0xJacky/Nginx-UI/model" @@ -28,7 +29,8 @@ func setup2FAStatusTestDB(t *testing.T) *gorm.DB { func TestGet2FAStatusRequiresRecoveryCodeMigrationForLegacyOTPUser(t *testing.T) { setup2FAStatusTestDB(t) gin.SetMode(gin.TestMode) - c, _ := gin.CreateTestContext(nil) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) c.Set("user", &model.User{OTPSecret: []byte("encrypted-secret")}) status := get2FAStatus(c) @@ -41,7 +43,8 @@ func TestGet2FAStatusRequiresRecoveryCodeMigrationForLegacyOTPUser(t *testing.T) func TestGet2FAStatusDoesNotRequireMigrationWhenRecoveryCodesExist(t *testing.T) { setup2FAStatusTestDB(t) gin.SetMode(gin.TestMode) - c, _ := gin.CreateTestContext(nil) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) c.Set("user", &model.User{ OTPSecret: []byte("encrypted-secret"), RecoveryCodes: model.RecoveryCodes{ diff --git a/app/src/components/TwoFA/use2FAModal.ts b/app/src/components/TwoFA/use2FAModal.ts index 13147d41..f2c476fa 100644 --- a/app/src/components/TwoFA/use2FAModal.ts +++ b/app/src/components/TwoFA/use2FAModal.ts @@ -64,18 +64,29 @@ function use2FAModal() { appContext: getCurrentInstance()?.appContext, width: '500px', content: () => { - const verifyOTP = (passcode: string, recovery: string) => { - twoFA.start_secure_session_by_otp(passcode, recovery).then(async r => { - modalInstance.destroy() - secureSessionId.value = r.session_id - await userStore.refreshTwoFAStatus() - if (r.used_legacy_recovery_code) - guideLegacyRecoveryMigration() - resolve(r.session_id) - }).catch(async () => { + const verifyOTP = async (passcode: string, recovery: string) => { + let result + try { + result = await twoFA.start_secure_session_by_otp(passcode, recovery) + } + catch { refOTPAuthorization.value?.clearInput() await message.error($gettext('Invalid passcode or recovery code')) - }) + return + } + + modalInstance.destroy() + secureSessionId.value = result.session_id + resolve(result.session_id) + + try { + await userStore.refreshTwoFAStatus() + if (result.used_legacy_recovery_code) + guideLegacyRecoveryMigration() + } + catch (error) { + console.error('Failed to handle post-OTP 2FA refresh:', error) + } } const setSessionId = (sessionId: string) => { diff --git a/app/src/pinia/moudule/user.ts b/app/src/pinia/moudule/user.ts index dccd346a..22e24196 100644 --- a/app/src/pinia/moudule/user.ts +++ b/app/src/pinia/moudule/user.ts @@ -129,9 +129,15 @@ export const useUserStore = defineStore('user', () => { if (!token.value) return twoFAStatus.value - const status = await twoFA.status() - twoFAStatus.value = status - return status + try { + const status = await twoFA.status() + twoFAStatus.value = status + return status + } + catch (error) { + console.error('Failed to refresh 2FA status:', error) + return twoFAStatus.value + } } async function updateCurrentUser(userData: Partial) { diff --git a/internal/user/otp.go b/internal/user/otp.go index 153302b4..9c7770fc 100644 --- a/internal/user/otp.go +++ b/internal/user/otp.go @@ -14,6 +14,8 @@ import ( "github.com/0xJacky/Nginx-UI/query" "github.com/google/uuid" "github.com/pquerna/otp/totp" + "gorm.io/gorm" + "gorm.io/gorm/clause" ) type OTPVerificationResult struct { @@ -40,32 +42,40 @@ func VerifyOTP(user *model.User, otp, recoveryCode string) (result OTPVerificati // legacy recovery code compatibility path if !user.RecoveryCodeGenerated() { - if user.OTPSecret == nil { - return result, ErrTOTPNotEnabled - } + err = model.UseDB().Transaction(func(tx *gorm.DB) error { + var lockedUser model.User + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&lockedUser, user.ID).Error; err != nil { + return err + } - if user.RecoveryCodes.LegacyRecoveryCodeUsedAt != nil { - return result, ErrRecoveryCode - } + if lockedUser.OTPSecret == nil { + return ErrTOTPNotEnabled + } - recoverCode, err := hex.DecodeString(recoveryCode) - if err != nil { - return result, err - } - k := sha1.Sum(user.OTPSecret) - if !bytes.Equal(k[:], recoverCode) { - return result, ErrRecoveryCode - } + if lockedUser.RecoveryCodeGenerated() || lockedUser.RecoveryCodes.LegacyRecoveryCodeUsedAt != nil { + return ErrRecoveryCode + } - t := time.Now().Unix() - user.RecoveryCodes.LegacyRecoveryCodeUsedAt = &t - _, err = u.Where(u.ID.Eq(user.ID)).Updates(user) - if err != nil { - return result, err - } + recoverCode, err := hex.DecodeString(recoveryCode) + if err != nil || len(recoverCode) != sha1.Size { + return ErrRecoveryCode + } - result.UsedLegacyRecoveryCode = true - return result, nil + k := sha1.Sum(lockedUser.OTPSecret) + if !bytes.Equal(k[:], recoverCode) { + return ErrRecoveryCode + } + + t := time.Now().Unix() + lockedUser.RecoveryCodes.LegacyRecoveryCodeUsedAt = &t + if err := tx.Model(&lockedUser).Select("recovery_codes").Updates(&lockedUser).Error; err != nil { + return err + } + + result.UsedLegacyRecoveryCode = true + return nil + }) + return result, err } // check recovery code diff --git a/internal/user/otp_test.go b/internal/user/otp_test.go index fcd7cc24..3364fc67 100644 --- a/internal/user/otp_test.go +++ b/internal/user/otp_test.go @@ -54,3 +54,18 @@ func TestVerifyOTPLegacyRecoveryCodeCanOnlyBeUsedOnce(t *testing.T) { require.ErrorIs(t, err, ErrRecoveryCode) assert.False(t, result.UsedLegacyRecoveryCode) } + +func TestVerifyOTPLegacyRecoveryCodeRejectsMalformedInput(t *testing.T) { + db := setupOTPTestDB(t) + + testUser := &model.User{ + Name: "legacy-user", + Status: true, + OTPSecret: []byte("encrypted-otp-secret"), + } + require.NoError(t, db.Create(testUser).Error) + + result, err := VerifyOTP(testUser, "", "not-hex") + require.ErrorIs(t, err, ErrRecoveryCode) + assert.False(t, result.UsedLegacyRecoveryCode) +}