From f9d629aa5599e519c75b3d2c10acaaf163fc468d Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Sat, 24 Jan 2026 11:00:28 +0100 Subject: [PATCH 01/13] Refactor remote MFA flow to use ClientRemoteMfaFinish RPC --- src/handlers/desktop_client_mfa.rs | 101 +++++++++-------------------- src/http.rs | 10 +-- 2 files changed, 33 insertions(+), 78 deletions(-) diff --git a/src/handlers/desktop_client_mfa.rs b/src/handlers/desktop_client_mfa.rs index 55a62b6..cf81926 100644 --- a/src/handlers/desktop_client_mfa.rs +++ b/src/handlers/desktop_client_mfa.rs @@ -1,5 +1,3 @@ -use std::collections::hash_map::Entry; - use axum::{ extract::{ ws::{Message, WebSocket}, @@ -19,8 +17,10 @@ use crate::{ handlers::get_core_response, http::AppState, proto::{ - core_request, core_response, ClientMfaFinishRequest, ClientMfaFinishResponse, - ClientMfaStartRequest, ClientMfaStartResponse, DeviceInfo, + core_request, + core_response::{self, Payload}, + ClientMfaFinishRequest, ClientMfaFinishResponse, ClientMfaStartRequest, + ClientMfaStartResponse, ClientRemoteMfaFinishRequest, DeviceInfo, }, }; @@ -53,63 +53,49 @@ async fn await_remote_auth( token: token.clone(), }, ), - device_info, + device_info.clone(), )?; let payload = get_core_response(rx).await?; if let core_response::Payload::ClientMfaTokenValidation(response) = payload { if !response.token_valid { return Err(ApiError::Unauthorized(String::new())); } - // check if its already in the map - let contains_key = { - let sessions = state.remote_mfa_sessions.lock().await; - sessions.contains_key(&token) - }; - if contains_key { - return Err(ApiError::Unauthorized(String::new())); - } - Ok(ws.on_upgrade(move |socket| handle_remote_auth_socket(socket, state.clone(), token))) + + let request = ClientRemoteMfaFinishRequest { token }; + let rx = state.grpc_server.send( + core_request::Payload::ClientRemoteMfaFinish(request), + device_info, + )?; + Ok(ws.on_upgrade(move |socket| handle_remote_auth_socket(socket, rx))) } else { Err(ApiError::InvalidResponseType) } } /// Handle axum web socket upgrade for `await_remote_auth`. -async fn handle_remote_auth_socket(socket: WebSocket, state: AppState, token: String) { - let (tx, rx) = oneshot::channel(); - - { - let mut sessions = state.remote_mfa_sessions.lock().await; - match sessions.entry(token.clone()) { - Entry::Occupied(_) => { - return; - } - Entry::Vacant(v) => { - v.insert(tx); - } - } - } - +async fn handle_remote_auth_socket(socket: WebSocket, rx: oneshot::Receiver) { let (mut ws_tx, mut ws_rx) = socket.split(); let mut set = JoinSet::new(); set.spawn(async move { - if let Ok(msg) = rx.await { - let payload = json!({ - "type": "mfa_success", - "preshared_key": &msg, - }); - if let Ok(serialized) = serde_json::to_string(&payload) { - let message = Message::Text(serialized.into()); - if ws_tx.send(message).await.is_err() { - error!("Failed to send preshared key via ws"); + // TODO(jck) unwrap + match rx.await.unwrap() { + Payload::ClientRemoteMfaFinish(response) => { + let ws_response = json!({ + "type": "mfa_success", + "preshared_key": &response.preshared_key, + }); + if let Ok(serialized) = serde_json::to_string(&ws_response) { + let message = Message::Text(serialized.into()); + if let Err(err) = ws_tx.send(message).await { + error!("Failed to send preshared key via ws: {err:?}"); + } } - } else { - error!("Failed to serialize remote mfa ws client response message"); } - } else { - error!("Failed to receive preshared key from receiver"); - } + _ => { + error!("Received wrong response type"); + } + }; let _ = ws_tx.close().await; }); @@ -131,8 +117,6 @@ async fn handle_remote_auth_socket(socket: WebSocket, state: AppState, token: St let _ = set.join_next().await; set.shutdown().await; - // This will remove token, if it's still there. - state.remote_mfa_sessions.lock().await.remove(&token); } #[instrument(level = "debug", skip(state, req))] @@ -186,30 +170,9 @@ async fn finish_remote_mfa( let rx = state .grpc_server .send(core_request::Payload::ClientMfaFinish(req), device_info)?; - let payload = get_core_response(rx).await?; - if let core_response::Payload::ClientMfaFinish(response) = payload { - // Check if this needs to be forwarded. - if let Some(token) = response.token { - let sender_option = { - let mut sessions = state.remote_mfa_sessions.lock().await; - sessions.remove(&token) - }; - if let Some(sender) = sender_option { - let _ = sender.send(response.preshared_key); - } - // If desktop stopped listening for the result, there will be no place to send the - // result. - else { - error!("Remote MFA approve finished but session was not found."); - return Err(ApiError::Unexpected(String::new())); - } - - info!("Finished desktop client authorization via mobile device"); - Ok(Json(json!({}))) - } else { - error!("Remote MFA Unexpected core response, token was not returned"); - Err(ApiError::Unexpected(String::new())) - } + // TODO(jck) can we make the response proto::Empty here? + if let core_response::Payload::ClientMfaFinish(_response) = get_core_response(rx).await? { + Ok(Json(json!({}))) } else { error!("Received invalid gRPC response type"); Err(ApiError::InvalidResponseType) diff --git a/src/http.rs b/src/http.rs index 0f298bb..ad64e25 100644 --- a/src/http.rs +++ b/src/http.rs @@ -1,5 +1,4 @@ use std::{ - collections::HashMap, fs::read_to_string, net::{IpAddr, Ipv4Addr, SocketAddr}, path::Path, @@ -21,11 +20,7 @@ use axum_extra::extract::cookie::Key; use clap::crate_version; use defguard_version::{server::DefguardVersionLayer, Version}; use serde::Serialize; -use tokio::{ - net::TcpListener, - sync::{oneshot, Mutex}, - task::JoinSet, -}; +use tokio::{net::TcpListener, sync::Mutex, task::JoinSet}; use tower_governor::{ governor::GovernorConfigBuilder, key_extractor::SmartIpKeyExtractor, GovernorLayer, }; @@ -62,8 +57,6 @@ pub static GRPC_SERVER_RESTART_CHANNEL: LazyLock> = LazyLock::n #[derive(Clone)] pub(crate) struct AppState { pub(crate) grpc_server: ProxyServer, - pub(crate) remote_mfa_sessions: - Arc>>>, cookie_key: Arc>>, url: Url, } @@ -292,7 +285,6 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> { let shared_state = AppState { cookie_key, grpc_server, - remote_mfa_sessions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), url: config.url.clone(), }; From ccbeb11b6293b419a952afe71d46aa96b7278e77 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Sat, 24 Jan 2026 11:13:32 +0100 Subject: [PATCH 02/13] handle remote mfa core comms in ws thread --- src/handlers/desktop_client_mfa.rs | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/src/handlers/desktop_client_mfa.rs b/src/handlers/desktop_client_mfa.rs index cf81926..0301fb4 100644 --- a/src/handlers/desktop_client_mfa.rs +++ b/src/handlers/desktop_client_mfa.rs @@ -10,7 +10,7 @@ use axum::{ use futures_util::{sink::SinkExt, stream::StreamExt}; use serde::Deserialize; use serde_json::json; -use tokio::{sync::oneshot, task::JoinSet}; +use tokio::{task::JoinSet}; use crate::{ error::ApiError, @@ -61,24 +61,34 @@ async fn await_remote_auth( return Err(ApiError::Unauthorized(String::new())); } - let request = ClientRemoteMfaFinishRequest { token }; - let rx = state.grpc_server.send( - core_request::Payload::ClientRemoteMfaFinish(request), - device_info, - )?; - Ok(ws.on_upgrade(move |socket| handle_remote_auth_socket(socket, rx))) + Ok(ws.on_upgrade(move |socket| { + handle_remote_auth_socket(socket, state.clone(), token, device_info) + })) } else { Err(ApiError::InvalidResponseType) } } /// Handle axum web socket upgrade for `await_remote_auth`. -async fn handle_remote_auth_socket(socket: WebSocket, rx: oneshot::Receiver) { +async fn handle_remote_auth_socket( + socket: WebSocket, + state: AppState, + token: String, + device_info: DeviceInfo, +) { let (mut ws_tx, mut ws_rx) = socket.split(); let mut set = JoinSet::new(); set.spawn(async move { - // TODO(jck) unwrap + let request = ClientRemoteMfaFinishRequest { token }; + let rx = state + .grpc_server + .send( + core_request::Payload::ClientRemoteMfaFinish(request), + device_info, + ) + .unwrap(); // TODO(jck) unwrap + // TODO(jck) unwrap match rx.await.unwrap() { Payload::ClientRemoteMfaFinish(response) => { let ws_response = json!({ From 91459e2476656aabcd564a8684378f6cb2aa6323 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Sun, 25 Jan 2026 08:17:16 +0100 Subject: [PATCH 03/13] longer timeout for remote mfa --- src/enterprise/handlers/desktop_client_mfa.rs | 2 +- src/enterprise/handlers/openid_login.rs | 4 ++-- src/handlers/desktop_client_mfa.rs | 12 ++++++++---- src/handlers/enrollment.rs | 8 ++++---- src/handlers/mobile_client.rs | 2 +- src/handlers/mod.rs | 9 ++++++--- src/handlers/password_reset.rs | 6 +++--- src/handlers/polling.rs | 2 +- src/handlers/register_mfa.rs | 4 ++-- 9 files changed, 28 insertions(+), 21 deletions(-) diff --git a/src/enterprise/handlers/desktop_client_mfa.rs b/src/enterprise/handlers/desktop_client_mfa.rs index fc75614..60b5a67 100644 --- a/src/enterprise/handlers/desktop_client_mfa.rs +++ b/src/enterprise/handlers/desktop_client_mfa.rs @@ -90,7 +90,7 @@ pub(super) async fn mfa_auth_callback( device_info, )?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; if let core_response::Payload::Empty(()) = payload { info!("MFA authentication callback completed successfully"); diff --git a/src/enterprise/handlers/openid_login.rs b/src/enterprise/handlers/openid_login.rs index 49ad838..848c281 100644 --- a/src/enterprise/handlers/openid_login.rs +++ b/src/enterprise/handlers/openid_login.rs @@ -76,7 +76,7 @@ async fn auth_info( let rx = state .grpc_server .send(core_request::Payload::AuthInfo(request), device_info)?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; if let core_response::Payload::AuthInfo(response) = payload { debug!("Received auth info response"); @@ -164,7 +164,7 @@ async fn auth_callback( let rx = state .grpc_server .send(core_request::Payload::AuthCallback(request), device_info)?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; if let core_response::Payload::AuthCallback(AuthCallbackResponse { url, token }) = payload { debug!("Received auth callback response {url:?} {token:?}"); diff --git a/src/handlers/desktop_client_mfa.rs b/src/handlers/desktop_client_mfa.rs index 0301fb4..736794c 100644 --- a/src/handlers/desktop_client_mfa.rs +++ b/src/handlers/desktop_client_mfa.rs @@ -1,3 +1,5 @@ +use std::time::Duration; + use axum::{ extract::{ ws::{Message, WebSocket}, @@ -24,6 +26,8 @@ use crate::{ }, }; +const REMOTE_AUTH_TIMEOUT: Duration = Duration::from_secs(60); + pub(crate) fn router() -> Router { Router::new() .route("/start", post(start_client_mfa)) @@ -55,7 +59,7 @@ async fn await_remote_auth( ), device_info.clone(), )?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, Some(REMOTE_AUTH_TIMEOUT)).await?; if let core_response::Payload::ClientMfaTokenValidation(response) = payload { if !response.token_valid { return Err(ApiError::Unauthorized(String::new())); @@ -140,7 +144,7 @@ async fn start_client_mfa( core_request::Payload::ClientMfaStart(req.clone()), device_info, )?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; if let core_response::Payload::ClientMfaStart(response) = payload { info!("Started desktop client authorization {req:?}"); @@ -161,7 +165,7 @@ async fn finish_client_mfa( let rx = state .grpc_server .send(core_request::Payload::ClientMfaFinish(req), device_info)?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; if let core_response::Payload::ClientMfaFinish(response) = payload { Ok(Json(response)) } else { @@ -181,7 +185,7 @@ async fn finish_remote_mfa( .grpc_server .send(core_request::Payload::ClientMfaFinish(req), device_info)?; // TODO(jck) can we make the response proto::Empty here? - if let core_response::Payload::ClientMfaFinish(_response) = get_core_response(rx).await? { + if let core_response::Payload::ClientMfaFinish(_response) = get_core_response(rx, None).await? { Ok(Json(json!({}))) } else { error!("Received invalid gRPC response type"); diff --git a/src/handlers/enrollment.rs b/src/handlers/enrollment.rs index b8c0e73..4979f35 100644 --- a/src/handlers/enrollment.rs +++ b/src/handlers/enrollment.rs @@ -45,7 +45,7 @@ async fn start_enrollment_process( let rx = state .grpc_server .send(core_request::Payload::EnrollmentStart(req), device_info)?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; debug!("Receving payload from the core service. Try to set private cookie for starting enrollment process."); if let core_response::Payload::EnrollmentStart(response) = payload { info!( @@ -83,7 +83,7 @@ async fn activate_user( let rx = state .grpc_server .send(core_request::Payload::ActivateUser(req), device_info)?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; debug!("Receiving payload from the core service. Trying to remove private cookie..."); if let core_response::Payload::Empty(()) = payload { info!("Activated user - phone number {phone:?}"); @@ -116,7 +116,7 @@ async fn create_device( let rx = state .grpc_server .send(core_request::Payload::NewDevice(req), device_info)?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; if let core_response::Payload::DeviceConfig(response) = payload { info!("Added new device {name} {pubkey}"); Ok(Json(response)) @@ -144,7 +144,7 @@ async fn get_network_info( let rx = state .grpc_server .send(core_request::Payload::ExistingDevice(req), device_info)?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; if let core_response::Payload::DeviceConfig(response) = payload { info!("Got network info for device {pubkey}"); Ok(Json(response)) diff --git a/src/handlers/mobile_client.rs b/src/handlers/mobile_client.rs index b39c2e9..c825d89 100644 --- a/src/handlers/mobile_client.rs +++ b/src/handlers/mobile_client.rs @@ -53,7 +53,7 @@ pub(crate) async fn register_mobile_auth( core_request::Payload::RegisterMobileAuth(send_data), device_info, )?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; if let core_response::Payload::Empty(()) = payload { info!("Registered mobile device for auth"); Ok(()) diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs index 629cbc2..955e2ac 100644 --- a/src/handlers/mod.rs +++ b/src/handlers/mod.rs @@ -3,7 +3,7 @@ use std::time::Duration; use axum::{extract::FromRequestParts, http::request::Parts}; use axum_client_ip::{InsecureClientIp, LeftmostXForwardedFor}; use axum_extra::{headers::UserAgent, TypedHeader}; -use tokio::{sync::oneshot::Receiver, time::timeout}; +use tokio::{sync::oneshot::Receiver, time}; use tonic::Code; use super::proto::DeviceInfo; @@ -69,9 +69,12 @@ where /// Helper which awaits core response /// /// Waits for core response with a given timeout and returns the response payload. -pub(crate) async fn get_core_response(rx: Receiver) -> Result { +pub(crate) async fn get_core_response( + rx: Receiver, + timeout: Option, +) -> Result { debug!("Fetching core response."); - if let Ok(core_response) = timeout(CORE_RESPONSE_TIMEOUT, rx).await { + if let Ok(core_response) = time::timeout(timeout.unwrap_or(CORE_RESPONSE_TIMEOUT), rx).await { debug!("Got gRPC response from Defguard Core"); if let Ok(Payload::CoreError(core_error)) = core_response { if core_error.status_code == Code::FailedPrecondition as i32 diff --git a/src/handlers/password_reset.rs b/src/handlers/password_reset.rs index 855b383..020eb2b 100644 --- a/src/handlers/password_reset.rs +++ b/src/handlers/password_reset.rs @@ -31,7 +31,7 @@ async fn request_password_reset( core_request::Payload::PasswordResetInit(req.clone()), device_info, )?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; if let core_response::Payload::Empty(()) = payload { info!("Started password reset request for {}", req.email); Ok(()) @@ -61,7 +61,7 @@ async fn start_password_reset( let rx = state .grpc_server .send(core_request::Payload::PasswordResetStart(req), device_info)?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; if let core_response::Payload::PasswordResetStart(response) = payload { // set session cookie let cookie = Cookie::build((PASSWORD_RESET_COOKIE_NAME, token)) @@ -92,7 +92,7 @@ async fn reset_password( let rx = state .grpc_server .send(core_request::Payload::PasswordReset(req), device_info)?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; if let core_response::Payload::Empty(()) = payload { if let Some(cookie) = private_cookies.get(PASSWORD_RESET_COOKIE_NAME) { info!("Password reset finished. Removing session cookie"); diff --git a/src/handlers/polling.rs b/src/handlers/polling.rs index 673e097..3749f03 100644 --- a/src/handlers/polling.rs +++ b/src/handlers/polling.rs @@ -18,7 +18,7 @@ pub(crate) async fn info( core_request::Payload::InstanceInfo(req.clone()), device_info, )?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; if let core_response::Payload::InstanceInfo(response) = payload { info!("Retrieved info for polling request"); diff --git a/src/handlers/register_mfa.rs b/src/handlers/register_mfa.rs index a8a2994..0a5266c 100644 --- a/src/handlers/register_mfa.rs +++ b/src/handlers/register_mfa.rs @@ -49,7 +49,7 @@ async fn register_code_mfa_start( }), device_info, )?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; match payload { core_response::Payload::CodeMfaSetupStartResponse(response) => Ok(Json(response)), _ => Err(ApiError::InvalidResponseType), @@ -90,7 +90,7 @@ async fn register_code_mfa_finish( }), device_info, )?; - let payload = get_core_response(rx).await?; + let payload = get_core_response(rx, None).await?; match payload { core_response::Payload::CodeMfaSetupFinishResponse(response) => Ok(Json(response)), _ => Err(ApiError::InvalidResponseType), From 38110cbe124cf3731e4e6daf73b49648053c8bfe Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Sun, 25 Jan 2026 14:43:37 +0100 Subject: [PATCH 04/13] Fix error msg --- src/handlers/mod.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs index 955e2ac..a1557cb 100644 --- a/src/handlers/mod.rs +++ b/src/handlers/mod.rs @@ -95,7 +95,10 @@ pub(crate) async fn get_core_response( core_response .map_err(|err| ApiError::Unexpected(format!("Failed to receive core response: {err}"))) } else { - error!("Did not receive response from Core within {CORE_RESPONSE_TIMEOUT:?}"); + error!( + "Did not receive response from Core within {:?}", + timeout.unwrap_or(CORE_RESPONSE_TIMEOUT) + ); Err(ApiError::CoreTimeout) } } From ae8d4d666db3a87449792b735e3ef45da8d8062e Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Mon, 26 Jan 2026 07:50:11 +0100 Subject: [PATCH 05/13] Use RwLocks instead of Mutexes --- src/grpc.rs | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/grpc.rs b/src/grpc.rs index 4ccbb20..f5da8c7 100644 --- a/src/grpc.rs +++ b/src/grpc.rs @@ -41,8 +41,8 @@ pub(crate) struct Configuration { pub(crate) struct ProxyServer { current_id: Arc, - clients: Arc>, - results: Arc>>>, + clients: Arc>, + results: Arc>>>, pub(crate) connected: Arc, pub(crate) core_version: Arc>>, config: Arc>>, @@ -57,8 +57,8 @@ impl ProxyServer { Self { cookie_key, current_id: Arc::new(AtomicU64::new(1)), - clients: Arc::new(Mutex::new(HashMap::new())), - results: Arc::new(Mutex::new(HashMap::new())), + clients: Arc::new(RwLock::new(HashMap::new())), + results: Arc::new(RwLock::new(HashMap::new())), connected: Arc::new(AtomicBool::new(false)), core_version: Arc::new(Mutex::new(None)), config: Arc::new(Mutex::new(None)), @@ -144,7 +144,7 @@ impl ProxyServer { ) -> Result, ApiError> { if let Some(client_tx) = self .clients - .lock() + .read() .expect("Failed to acquire lock on clients hashmap when sending message to core") .values() .next() @@ -161,7 +161,7 @@ impl ProxyServer { } let (tx, rx) = oneshot::channel(); self.results - .lock() + .write() .expect("Failed to acquire lock on results hashmap when sending CoreRequest") .insert(id, tx); self.connected.store(true, Ordering::Relaxed); @@ -233,7 +233,7 @@ impl proxy_server::Proxy for ProxyServer { info!("Defguard Core gRPC client connected from: {address}"); let (tx, rx) = mpsc::unbounded_channel(); self.clients - .lock() + .write() .expect( "Failed to acquire lock on clients hashmap when registering new core connection", ) @@ -260,7 +260,7 @@ impl proxy_server::Proxy for ProxyServer { *cookie_key.write().unwrap() = Some(key); }, _ => { - let maybe_rx = results.lock().expect("Failed to acquire lock on results hashmap when processing response").remove(&response.id); + let maybe_rx = results.write().expect("Failed to acquire lock on results hashmap when processing response").remove(&response.id); if let Some(rx) = maybe_rx { if let Err(err) = rx.send(payload) { error!("Failed to send message to rx {:?}", err.type_id()); @@ -284,7 +284,7 @@ impl proxy_server::Proxy for ProxyServer { } info!("Defguard core client disconnected: {address}"); connected.store(false, Ordering::Relaxed); - clients.lock().expect("Failed to acquire lock on clients hashmap when removing disconnected client").remove(&address); + clients.write().expect("Failed to acquire lock on clients hashmap when removing disconnected client").remove(&address); } .instrument(tracing::Span::current()), ); From 86f235d6e095897d341783162155bb0fd5cce600 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Mon, 26 Jan 2026 10:42:05 +0100 Subject: [PATCH 06/13] update protos --- proto | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/proto b/proto index ec48aca..246f5f8 160000 --- a/proto +++ b/proto @@ -1 +1 @@ -Subproject commit ec48aca9438e7cdcb4fcdb01ce6dcb5dac7f8dd3 +Subproject commit 246f5f856e3ed3d6549897c8a0a84ed46b6165fc From 11e67fcdacf465cd2d1296b1a14627e115cfceb9 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Mon, 26 Jan 2026 11:08:46 +0100 Subject: [PATCH 07/13] remove unused proto field --- proto | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/proto b/proto index 246f5f8..ba933c4 160000 --- a/proto +++ b/proto @@ -1 +1 @@ -Subproject commit 246f5f856e3ed3d6549897c8a0a84ed46b6165fc +Subproject commit ba933c4c1fa3edb7c0e3dbafa123f9446b3f9533 From 84cebfd65f692a8c05d7d23fba16df2681e836e7 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Mon, 26 Jan 2026 11:09:40 +0100 Subject: [PATCH 08/13] cargo fmt --- src/handlers/desktop_client_mfa.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/handlers/desktop_client_mfa.rs b/src/handlers/desktop_client_mfa.rs index 736794c..98f9c87 100644 --- a/src/handlers/desktop_client_mfa.rs +++ b/src/handlers/desktop_client_mfa.rs @@ -12,7 +12,7 @@ use axum::{ use futures_util::{sink::SinkExt, stream::StreamExt}; use serde::Deserialize; use serde_json::json; -use tokio::{task::JoinSet}; +use tokio::task::JoinSet; use crate::{ error::ApiError, @@ -92,7 +92,7 @@ async fn handle_remote_auth_socket( device_info, ) .unwrap(); // TODO(jck) unwrap - // TODO(jck) unwrap + // TODO(jck) unwrap match rx.await.unwrap() { Payload::ClientRemoteMfaFinish(response) => { let ws_response = json!({ From bc16a15e226bb1d47a7f6afc32861bde21c783c7 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Mon, 26 Jan 2026 11:59:45 +0100 Subject: [PATCH 09/13] fix unwraps --- src/handlers/desktop_client_mfa.rs | 40 ++++++++++++++++++++---------- 1 file changed, 27 insertions(+), 13 deletions(-) diff --git a/src/handlers/desktop_client_mfa.rs b/src/handlers/desktop_client_mfa.rs index 98f9c87..dd90c30 100644 --- a/src/handlers/desktop_client_mfa.rs +++ b/src/handlers/desktop_client_mfa.rs @@ -26,6 +26,7 @@ use crate::{ }, }; +// How much time the user has to approve remote MFA with mobile device const REMOTE_AUTH_TIMEOUT: Duration = Duration::from_secs(60); pub(crate) fn router() -> Router { @@ -83,18 +84,23 @@ async fn handle_remote_auth_socket( let (mut ws_tx, mut ws_rx) = socket.split(); let mut set = JoinSet::new(); + let request = ClientRemoteMfaFinishRequest { token }; + let rx = match state.grpc_server.send( + core_request::Payload::ClientRemoteMfaFinish(request), + device_info, + ) { + Ok(rx) => rx, + Err(err) => { + error!("Failed to send ClientRemoteMfaFinishRequest: {err:?}"); + return; + } + }; + + // Response to ClientRemoteMfaFinishRequest comes once the user concludes MFA with mobile device. + // This task then sends the preshared key to the WebSocket where desktop client awaits for it. set.spawn(async move { - let request = ClientRemoteMfaFinishRequest { token }; - let rx = state - .grpc_server - .send( - core_request::Payload::ClientRemoteMfaFinish(request), - device_info, - ) - .unwrap(); // TODO(jck) unwrap - // TODO(jck) unwrap - match rx.await.unwrap() { - Payload::ClientRemoteMfaFinish(response) => { + match rx.await { + Ok(Payload::ClientRemoteMfaFinish(response)) => { let ws_response = json!({ "type": "mfa_success", "preshared_key": &response.preshared_key, @@ -106,13 +112,20 @@ async fn handle_remote_auth_socket( } } } - _ => { - error!("Received wrong response type"); + Ok(_) => { + error!("Received wrong response type, expected ClientRemoteMfaFinish"); + } + Err(err) => { + error!("Failed to receive preshared key from receiver: {err:?}"); } }; + + // Close the websocket once we're done. let _ = ws_tx.close().await; }); + // Another task to monitor the websocket connection in case desktop client disconnects + // or the connection errors-out. set.spawn(async move { while let Some(msg_result) = ws_rx.next().await { match msg_result { @@ -129,6 +142,7 @@ async fn handle_remote_auth_socket( } }); + // Wait for whichever task finishes first and kill the other one. let _ = set.join_next().await; set.shutdown().await; } From ccc99cb6d389e92d5bada982a9bec421f72a15d1 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Mon, 26 Jan 2026 12:02:32 +0100 Subject: [PATCH 10/13] improve log message --- src/handlers/desktop_client_mfa.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/handlers/desktop_client_mfa.rs b/src/handlers/desktop_client_mfa.rs index dd90c30..063efab 100644 --- a/src/handlers/desktop_client_mfa.rs +++ b/src/handlers/desktop_client_mfa.rs @@ -198,11 +198,10 @@ async fn finish_remote_mfa( let rx = state .grpc_server .send(core_request::Payload::ClientMfaFinish(req), device_info)?; - // TODO(jck) can we make the response proto::Empty here? if let core_response::Payload::ClientMfaFinish(_response) = get_core_response(rx, None).await? { Ok(Json(json!({}))) } else { - error!("Received invalid gRPC response type"); + error!("Received invalid gRPC response type, expected ClientMfaFinish"); Err(ApiError::InvalidResponseType) } } From 10057689d6dfc9b801d1019dc7e4b985d38d75ab Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Mon, 26 Jan 2026 13:33:41 +0100 Subject: [PATCH 11/13] rename rpc method --- proto | 2 +- src/handlers/desktop_client_mfa.rs | 11 ++++------- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/proto b/proto index ba933c4..3a18dbc 160000 --- a/proto +++ b/proto @@ -1 +1 @@ -Subproject commit ba933c4c1fa3edb7c0e3dbafa123f9446b3f9533 +Subproject commit 3a18dbc52f098f37d5d78e88904631faf8db78f3 diff --git a/src/handlers/desktop_client_mfa.rs b/src/handlers/desktop_client_mfa.rs index 063efab..7285af7 100644 --- a/src/handlers/desktop_client_mfa.rs +++ b/src/handlers/desktop_client_mfa.rs @@ -19,10 +19,7 @@ use crate::{ handlers::get_core_response, http::AppState, proto::{ - core_request, - core_response::{self, Payload}, - ClientMfaFinishRequest, ClientMfaFinishResponse, ClientMfaStartRequest, - ClientMfaStartResponse, ClientRemoteMfaFinishRequest, DeviceInfo, + core_request, core_response::{self, Payload}, AwaitRemoteMfaFinishRequest, ClientMfaFinishRequest, ClientMfaFinishResponse, ClientMfaStartRequest, ClientMfaStartResponse, DeviceInfo }, }; @@ -84,9 +81,9 @@ async fn handle_remote_auth_socket( let (mut ws_tx, mut ws_rx) = socket.split(); let mut set = JoinSet::new(); - let request = ClientRemoteMfaFinishRequest { token }; + let request = AwaitRemoteMfaFinishRequest { token }; let rx = match state.grpc_server.send( - core_request::Payload::ClientRemoteMfaFinish(request), + core_request::Payload::AwaitRemoteMfaFinish(request), device_info, ) { Ok(rx) => rx, @@ -100,7 +97,7 @@ async fn handle_remote_auth_socket( // This task then sends the preshared key to the WebSocket where desktop client awaits for it. set.spawn(async move { match rx.await { - Ok(Payload::ClientRemoteMfaFinish(response)) => { + Ok(Payload::AwaitRemoteMfaFinish(response)) => { let ws_response = json!({ "type": "mfa_success", "preshared_key": &response.preshared_key, From e20bfcf9b4e36abfac4277a1f2ead89a55c603c6 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Mon, 26 Jan 2026 13:57:38 +0100 Subject: [PATCH 12/13] update protos --- proto | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/proto b/proto index 3a18dbc..0b98292 160000 --- a/proto +++ b/proto @@ -1 +1 @@ -Subproject commit 3a18dbc52f098f37d5d78e88904631faf8db78f3 +Subproject commit 0b982922c4dab3304a8cb01aed1d8cee806600b7 From 083f0f5430b9432ef2b9f56c92e5f3481cd98041 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Mon, 26 Jan 2026 15:08:28 +0100 Subject: [PATCH 13/13] cargo fmt --- src/handlers/desktop_client_mfa.rs | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/src/handlers/desktop_client_mfa.rs b/src/handlers/desktop_client_mfa.rs index 7285af7..9f278ff 100644 --- a/src/handlers/desktop_client_mfa.rs +++ b/src/handlers/desktop_client_mfa.rs @@ -19,7 +19,10 @@ use crate::{ handlers::get_core_response, http::AppState, proto::{ - core_request, core_response::{self, Payload}, AwaitRemoteMfaFinishRequest, ClientMfaFinishRequest, ClientMfaFinishResponse, ClientMfaStartRequest, ClientMfaStartResponse, DeviceInfo + core_request, + core_response::{self, Payload}, + AwaitRemoteMfaFinishRequest, ClientMfaFinishRequest, ClientMfaFinishResponse, + ClientMfaStartRequest, ClientMfaStartResponse, DeviceInfo, }, }; @@ -93,8 +96,8 @@ async fn handle_remote_auth_socket( } }; - // Response to ClientRemoteMfaFinishRequest comes once the user concludes MFA with mobile device. - // This task then sends the preshared key to the WebSocket where desktop client awaits for it. + // Response to ClientRemoteMfaFinishRequest comes once the user concludes MFA with mobile device. + // This task then sends the preshared key to the WebSocket where desktop client awaits for it. set.spawn(async move { match rx.await { Ok(Payload::AwaitRemoteMfaFinish(response)) => { @@ -117,12 +120,12 @@ async fn handle_remote_auth_socket( } }; - // Close the websocket once we're done. + // Close the websocket once we're done. let _ = ws_tx.close().await; }); - // Another task to monitor the websocket connection in case desktop client disconnects - // or the connection errors-out. + // Another task to monitor the websocket connection in case desktop client disconnects + // or the connection errors-out. set.spawn(async move { while let Some(msg_result) = ws_rx.next().await { match msg_result { @@ -139,7 +142,7 @@ async fn handle_remote_auth_socket( } }); - // Wait for whichever task finishes first and kill the other one. + // Wait for whichever task finishes first and kill the other one. let _ = set.join_next().await; set.shutdown().await; }