track pending 2FA challenge state

Signed-off-by: 21pages <sunboeasy@gmail.com>
This commit is contained in:
21pages
2026-07-27 17:51:22 +08:00
parent 6692bfeefc
commit b0b1fa7ff9

View File

@@ -323,6 +323,7 @@ pub struct Connection {
tx_to_cm: mpsc::UnboundedSender<ipc::Data>, tx_to_cm: mpsc::UnboundedSender<ipc::Data>,
authorized: bool, authorized: bool,
require_2fa: Option<totp_rs::TOTP>, require_2fa: Option<totp_rs::TOTP>,
awaiting_2fa: bool,
keyboard: bool, keyboard: bool,
clipboard: bool, clipboard: bool,
audio: bool, audio: bool,
@@ -509,6 +510,7 @@ impl Connection {
tx_video: Some(tx_video), tx_video: Some(tx_video),
}, },
require_2fa: crate::auth_2fa::get_2fa(None), require_2fa: crate::auth_2fa::get_2fa(None),
awaiting_2fa: false,
// Defer display enumeration until login succeeds. Monitor login replaces this // Defer display enumeration until login succeeds. Monitor login replaces this
// with the primary index returned with the refreshed display snapshot. // with the primary index returned with the refreshed display snapshot.
display_idx: 0, display_idx: 0,
@@ -1678,10 +1680,12 @@ impl Connection {
}); });
} }
}); });
self.awaiting_2fa = true;
self.send_login_error(crate::client::REQUIRE_2FA).await; self.send_login_error(crate::client::REQUIRE_2FA).await;
// Keep the connection alive so the client can continue with 2FA. // Keep the connection alive so the client can continue with 2FA.
return true; return true;
} }
self.awaiting_2fa = false;
if let Some(keep_alive) = self.prepare_terminal_login_for_authorization().await { if let Some(keep_alive) = self.prepare_terminal_login_for_authorization().await {
return keep_alive; return keep_alive;
} }
@@ -2559,6 +2563,7 @@ impl Connection {
} }
// After handling CloseReason messages, proceed to process other message types // After handling CloseReason messages, proceed to process other message types
if let Some(message::Union::LoginRequest(lr)) = msg.union { if let Some(message::Union::LoginRequest(lr)) = msg.union {
self.awaiting_2fa = false;
self.handle_login_request_without_validation(&lr).await; self.handle_login_request_without_validation(&lr).await;
if self.authorized { if self.authorized {
return true; return true;
@@ -2791,6 +2796,11 @@ impl Connection {
} }
} }
} else if let Some(message::Union::Auth2fa(tfa)) = msg.union { } else if let Some(message::Union::Auth2fa(tfa)) = msg.union {
// A 2FA response may arrive after click authorization has completed.
// Ignore it unless this connection is still waiting for the response.
if !self.awaiting_2fa {
return true;
}
let (failure, res) = self.check_failure(1).await; let (failure, res) = self.check_failure(1).await;
if !res { if !res {
return true; return true;