fix: harden recovery migration flow

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{

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) => {

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<User>) {

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

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