From 0368ccd01d2cdc87311e7a3b9c7f0810dfa61171 Mon Sep 17 00:00:00 2001
From: FatttSnake <fatttsnake@gmail.com>
Date: Tue, 5 Mar 2024 11:21:32 +0800
Subject: [PATCH] Optimize: login - check two-factor authentication after
 verifying password

---
 .../impl/AuthenticationServiceImpl.kt         | 21 ++++++++-----------
 1 file changed, 9 insertions(+), 12 deletions(-)

diff --git a/src/main/kotlin/top/fatweb/oxygen/api/service/permission/impl/AuthenticationServiceImpl.kt b/src/main/kotlin/top/fatweb/oxygen/api/service/permission/impl/AuthenticationServiceImpl.kt
index 81dd802..81372a3 100644
--- a/src/main/kotlin/top/fatweb/oxygen/api/service/permission/impl/AuthenticationServiceImpl.kt
+++ b/src/main/kotlin/top/fatweb/oxygen/api/service/permission/impl/AuthenticationServiceImpl.kt
@@ -360,18 +360,6 @@ class AuthenticationServiceImpl(
         password: String,
         twoFactorCode: String? = null
     ): LoginVo {
-        val userWithPowerByAccount = userService.getUserWithPowerByAccount(account) ?: throw UserNotFoundException()
-        if (!userWithPowerByAccount.twoFactor.isNullOrBlank()
-            && !userWithPowerByAccount.twoFactor!!.endsWith("?")
-        ) {
-            if (twoFactorCode.isNullOrBlank()) {
-                throw NeedTwoFactorException()
-            }
-            if (!TOTPUtil.validateCode(userWithPowerByAccount.twoFactor!!, twoFactorCode)) {
-                throw TwoFactorVerificationCodeErrorException()
-            }
-        }
-
         val usernamePasswordAuthenticationToken =
             UsernamePasswordAuthenticationToken(account, password)
         val authentication = authenticationManager.authenticate(usernamePasswordAuthenticationToken)
@@ -380,6 +368,15 @@ class AuthenticationServiceImpl(
         val loginUser = authentication.principal as LoginUser
         loginUser.user.password = ""
 
+        if (!loginUser.user.twoFactor.isNullOrBlank() && !loginUser.user.twoFactor!!.endsWith("?")) {
+            if (twoFactorCode.isNullOrBlank()) {
+                throw NeedTwoFactorException()
+            }
+            if (!TOTPUtil.validateCode(loginUser.user.twoFactor!!, twoFactorCode)) {
+                throw TwoFactorVerificationCodeErrorException()
+            }
+        }
+
         logger.info("用户登录 [用户名: '{}', IP: '{}']", loginUser.username, request.remoteAddr)
         userService.update(User().apply {
             currentLoginIp = request.remoteAddr
-- 
2.49.1