From 21e4679b6ce0234daec1c9fc8a9c2884b47dfc36 Mon Sep 17 00:00:00 2001
From: rustdesk <info@rustdesk.com>
Date: Wed, 20 Mar 2024 20:00:29 +0800
Subject: [PATCH] improve tfa

---
 src/auth_2fa.rs          | 3 +--
 src/server/connection.rs | 4 ++--
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/auth_2fa.rs b/src/auth_2fa.rs
index bfbecd4e9..1ea773342 100644
--- a/src/auth_2fa.rs
+++ b/src/auth_2fa.rs
@@ -89,8 +89,7 @@ pub fn generate2fa() -> String {
 
 pub fn verify2fa(code: String) -> bool {
     if let Some((info, totp)) = CURRENT_2FA.lock().unwrap().as_ref() {
-        if let Ok(cur) = totp.generate_current() {
-            let res = code == cur;
+        if let Ok(res) = totp.check_current(&code) {
             if res {
                 if let Ok(v) = info.into_string() {
                     #[cfg(not(any(target_os = "android", target_os = "ios")))]
diff --git a/src/server/connection.rs b/src/server/connection.rs
index d88ab0c3e..158441785 100644
--- a/src/server/connection.rs
+++ b/src/server/connection.rs
@@ -1750,8 +1750,8 @@ impl Connection {
                 return true;
             }
             if let Some(totp) = self.require_2fa.as_ref() {
-                if let Ok(code) = totp.generate_current() {
-                    if tfa.code == code {
+                if let Ok(res) = totp.check_current(&tfa.code) {
+                    if res {
                         self.update_failure(failure, true, 1);
                         self.require_2fa.take();
                         self.send_logon_response().await;