From 9184711d439e05a0a73b81e21137c837b96c8a56 Mon Sep 17 00:00:00 2001
From: Hintay <hintay@me.com>
Date: Mon, 10 Feb 2025 14:44:01 +0900
Subject: [PATCH] feat: validate new recovery code

---
 app/src/api/2fa.ts                         |  2 +-
 app/src/components/TwoFA/Authorization.vue | 69 +++++++++++-----------
 app/src/views/other/Login.vue              |  1 +
 internal/user/errors.go                    |  1 +
 internal/user/otp.go                       | 37 ++++++++++--
 5 files changed, 69 insertions(+), 41 deletions(-)

diff --git a/app/src/api/2fa.ts b/app/src/api/2fa.ts
index e5282a0a..16850ffa 100644
--- a/app/src/api/2fa.ts
+++ b/app/src/api/2fa.ts
@@ -6,7 +6,7 @@ export interface TwoFAStatus {
   otp_status: boolean
   passkey_status: boolean
   recovery_codes_generated: boolean
-  recovery_codes_viewed: boolean
+  recovery_codes_viewed?: boolean
 }
 
 const twoFA = {
diff --git a/app/src/components/TwoFA/Authorization.vue b/app/src/components/TwoFA/Authorization.vue
index 0eac5c7d..e8b8cc9c 100644
--- a/app/src/components/TwoFA/Authorization.vue
+++ b/app/src/components/TwoFA/Authorization.vue
@@ -65,42 +65,30 @@ onMounted(() => {
 
 <template>
   <div>
-    <div v-if="twoFAStatus.otp_status">
-      <div v-if="!useRecoveryCode">
-        <p>{{ $gettext('Please enter the OTP code:') }}</p>
-        <OTPInput
-          ref="refOTP"
-          v-model="passcode"
-          class="justify-center mb-6"
-          @on-complete="onSubmit"
-        />
-      </div>
-      <div
-        v-else
-        class="mt-2 mb-4"
-      >
-        <p>{{ $gettext('Input the recovery code:') }}</p>
-        <AInputGroup compact>
-          <AInput v-model:value="recoveryCode" />
-          <AButton
-            type="primary"
-            @click="onSubmit"
-          >
-            {{ $gettext('Recovery') }}
-          </AButton>
-        </AInputGroup>
-      </div>
+    <div
+      v-if="useRecoveryCode"
+      class="mt-2 mb-4"
+    >
+      <p>{{ $gettext('Input the recovery code:') }}</p>
+      <AInputGroup compact>
+        <AInput v-model:value="recoveryCode" placeholder="xxxxx-xxxxx" />
+        <AButton
+          type="primary"
+          @click="onSubmit"
+        >
+          {{ $gettext('Recovery') }}
+        </AButton>
+      </AInputGroup>
+    </div>
 
-      <div class="flex justify-center">
-        <a
-          v-if="!useRecoveryCode"
-          @click="clickUseRecoveryCode"
-        >{{ $gettext('Use recovery code') }}</a>
-        <a
-          v-else
-          @click="clickUseOTP"
-        >{{ $gettext('Use OTP') }}</a>
-      </div>
+    <div v-if="twoFAStatus.otp_status && !useRecoveryCode">
+      <p>{{ $gettext('Please enter the OTP code:') }}</p>
+      <OTPInput
+        ref="refOTP"
+        v-model="passcode"
+        class="justify-center mb-6"
+        @on-complete="onSubmit"
+      />
     </div>
 
     <div
@@ -121,6 +109,17 @@ onMounted(() => {
         {{ $gettext('Authenticate with a passkey') }}
       </AButton>
     </div>
+
+    <div v-if="twoFAStatus.otp_status || twoFAStatus.recovery_codes_generated" class="flex justify-center">
+      <a
+        v-if="!useRecoveryCode"
+        @click="clickUseRecoveryCode"
+      >{{ $gettext('Use recovery code') }}</a>
+      <a
+        v-else-if="twoFAStatus.otp_status"
+        @click="clickUseOTP"
+      >{{ $gettext('Use OTP') }}</a>
+    </div>
   </div>
 </template>
 
diff --git a/app/src/views/other/Login.vue b/app/src/views/other/Login.vue
index 2407da9d..9295e2bd 100644
--- a/app/src/views/other/Login.vue
+++ b/app/src/views/other/Login.vue
@@ -205,6 +205,7 @@ async function handlePasskeyLogin() {
                   enabled: true,
                   otp_status: true,
                   passkey_status: false,
+                  recovery_codes_generated: true,
                 }"
                 @submit-o-t-p="handleOTPSubmit"
               />
diff --git a/internal/user/errors.go b/internal/user/errors.go
index a39ae32b..6969ef76 100644
--- a/internal/user/errors.go
+++ b/internal/user/errors.go
@@ -8,6 +8,7 @@ var (
 	ErrUserBanned              = e.New(40303, "user banned")
 	ErrOTPCode                 = e.New(40304, "invalid otp code")
 	ErrRecoveryCode            = e.New(40305, "invalid recovery code")
+	ErrTOTPNotEnabled          = e.New(40306, "legacy recovery code not allowed since totp is not enabled")
 	ErrWebAuthnNotConfigured   = e.New(50000, "WebAuthn settings are not configured")
 	ErrUserNotEnabledOTPAs2FA  = e.New(50001, "user not enabled otp as 2fa")
 	ErrOTPOrRecoveryCodeEmpty  = e.New(50002, "otp or recovery code empty")
diff --git a/internal/user/otp.go b/internal/user/otp.go
index 763c33fe..6339800c 100644
--- a/internal/user/otp.go
+++ b/internal/user/otp.go
@@ -5,12 +5,14 @@ import (
 	"crypto/sha1"
 	"encoding/hex"
 	"fmt"
+	"time"
+
 	"github.com/0xJacky/Nginx-UI/internal/cache"
 	"github.com/0xJacky/Nginx-UI/internal/crypto"
 	"github.com/0xJacky/Nginx-UI/model"
+	"github.com/0xJacky/Nginx-UI/query"
 	"github.com/google/uuid"
 	"github.com/pquerna/otp/totp"
-	"time"
 )
 
 func VerifyOTP(user *model.User, otp, recoveryCode string) (err error) {
@@ -24,14 +26,39 @@ func VerifyOTP(user *model.User, otp, recoveryCode string) (err error) {
 			return ErrOTPCode
 		}
 	} else {
-		recoverCode, err := hex.DecodeString(recoveryCode)
+		// get user from db
+		u := query.User
+		user, err = u.Where(u.ID.Eq(user.ID)).First()
 		if err != nil {
 			return err
 		}
-		k := sha1.Sum(user.OTPSecret)
-		if !bytes.Equal(k[:], recoverCode) {
-			return ErrRecoveryCode
+
+		// legacy recovery code
+		if !user.RecoveryCodeGenerated() {
+			if user.OTPSecret == nil {
+				return ErrTOTPNotEnabled
+			}
+
+			recoverCode, err := hex.DecodeString(recoveryCode)
+			if err != nil {
+				return err
+			}
+			k := sha1.Sum(user.OTPSecret)
+			if !bytes.Equal(k[:], recoverCode) {
+				return ErrRecoveryCode
+			}
 		}
+
+		// check recovery code
+		for _, code := range user.RecoveryCodes.Codes {
+			if code.Code == recoveryCode && code.UsedTime == nil {
+				t := time.Now()
+				code.UsedTime = &t
+				_, err = u.Where(u.ID.Eq(user.ID)).Updates(user)
+				return
+			}
+		}
+		return ErrRecoveryCode
 	}
 	return
 }