Desktop MFA mobile approve (#138)

* compiles

* update proto

* add mfa token validation

* Update src/handlers/desktop_client_mfa.rs

Co-authored-by: Adam <[email protected]>

* fix cargo lock

* remove dashmap

* review changes

* fmt

* review changes

* Update defguard-ui

* update proto

* upgrade ui module

* Update desktop_client_mfa.rs

* Update desktop_client_mfa.rs

---------

Co-authored-by: Adam <[email protected]>

Author: filipslezaklab

Cargo.lock

@@ -8,7 +8,7 @@ repository = "https://github.com/DefGuard/proxy"
 
 [dependencies]
 # base `axum` deps
-axum = { version = "0.7", features = ["macros", "tracing"] }
+axum = { version = "0.7", features = ["macros", "tracing", "ws"] }
 axum-client-ip = "0.6"
 axum-extra = { version = "0.9", features = [
     "cookie",
@@ -48,6 +48,8 @@ tower_governor = "0.4"
 rust-embed = { version = "8.5", features = ["include-exclude"] }
 mime_guess = "2.0"
 base64 = "0.22.1"
+futures = "0.3.31"
+futures-util = "0.3.31"
 
 [build-dependencies]
 tonic-prost-build = "0.14"

Cargo.toml

@@ -1,4 +1,17 @@
-use axum::{extract::State, routing::post, Json, Router};
+use axum::{
+    extract::{
+        ws::{Message, WebSocket},
+        Query, State, WebSocketUpgrade,
+    },
+    response::{IntoResponse, Response},
+    routing::{get, post},
+    Json, Router,
+};
+use futures_util::{sink::SinkExt, stream::StreamExt};
+use serde::Deserialize;
+use serde_json::json;
+use std::collections::hash_map::Entry;
+use tokio::{sync::oneshot, task::JoinSet};
 
 use crate::{
     error::ApiError,
@@ -14,9 +27,116 @@ pub(crate) fn router() -> Router<AppState> {
     Router::new()
         .route("/start", post(start_client_mfa))
         .route("/finish", post(finish_client_mfa))
+        .route("/remote", get(await_remote_auth))
+        .route("/finish-remote", post(finish_remote_mfa))
 }
 
-#[instrument(level = "debug", skip(state))]
+#[derive(Debug, Clone, Deserialize)]
+pub(crate) struct RemoteMfaRequestQuery {
+    pub token: String,
+}
+
+// Allows desktop client to await for another device to complete MFA for it via mobile client
+#[instrument(level = "debug", skip(state, req))]
+async fn await_remote_auth(
+    ws: WebSocketUpgrade,
+    Query(req): Query<RemoteMfaRequestQuery>,
+    State(state): State<AppState>,
+    device_info: DeviceInfo,
+) -> Result<Response, impl IntoResponse> {
+    let token = req.token;
+    // let core validate token first
+    let rx = state.grpc_server.send(
+        core_request::Payload::ClientMfaTokenValidation(
+            crate::proto::ClientMfaTokenValidationRequest {
+                token: token.clone(),
+            },
+        ),
+        device_info,
+    )?;
+    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)))
+    } else {
+        Err(ApiError::InvalidResponseType)
+    }
+}
+
+// handle axum ws socket upgrade for await_remote_auth
+async fn handle_remote_auth_socket(socket: WebSocket, state: AppState, token: String) {
+    let (tx, rx) = oneshot::channel::<String>();
+    let (mut ws_tx, mut ws_rx) = socket.split();
+
+    let occupied = {
+        let mut sessions = state.remote_mfa_sessions.lock().await;
+        match sessions.entry(token.clone()) {
+            Entry::Occupied(_) => true,
+            Entry::Vacant(v) => {
+                v.insert(tx);
+                false
+            }
+        }
+    };
+    if occupied {
+        let _ = ws_tx.close().await;
+        return;
+    }
+
+    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);
+                if ws_tx.send(message).await.is_err() {
+                    error!("Failed to send preshared key via ws");
+                }
+            } else {
+                error!("Failed to serialize remote mfa ws client response message");
+            }
+        } else {
+            error!("Failed to receive preshared key from receiver")
+        }
+        let _ = ws_tx.close().await;
+    });
+    set.spawn(async move {
+        while let Some(msg_result) = ws_rx.next().await {
+            match msg_result {
+                Ok(msg) => {
+                    if let Message::Close(_) = msg {
+                        break;
+                    }
+                }
+                Err(e) => {
+                    error!("Remote desktop mfa WS client listen error {e}");
+                    break;
+                }
+            }
+        }
+    });
+
+    let _ = set.join_next().await;
+    set.shutdown().await;
+    // will remove token if it's still there
+    state.remote_mfa_sessions.lock().await.remove(&token);
+}
+
+#[instrument(level = "debug", skip(state, req))]
 async fn start_client_mfa(
     State(state): State<AppState>,
     device_info: DeviceInfo,
@@ -38,7 +158,7 @@ async fn start_client_mfa(
     }
 }
 
-#[instrument(level = "debug", skip(state))]
+#[instrument(level = "debug", skip(state, req))]
 async fn finish_client_mfa(
     State(state): State<AppState>,
     device_info: DeviceInfo,
@@ -50,10 +170,52 @@ async fn finish_client_mfa(
         .send(core_request::Payload::ClientMfaFinish(req), device_info)?;
     let payload = get_core_response(rx).await?;
     if let core_response::Payload::ClientMfaFinish(response) = payload {
-        info!("Finished desktop client authorization");
         Ok(Json(response))
     } else {
         error!("Received invalid gRPC response type: {payload:#?}");
         Err(ApiError::InvalidResponseType)
     }
 }
+
+#[instrument(level = "debug", skip(state, req))]
+async fn finish_remote_mfa(
+    State(state): State<AppState>,
+    device_info: DeviceInfo,
+    Json(req): Json<ClientMfaFinishRequest>,
+) -> Result<Json<serde_json::Value>, ApiError> {
+    info!("Finishing desktop client authorization");
+    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
+        match response.token {
+            Some(token) => {
+                let sender_option = {
+                    let mut sessions = state.remote_mfa_sessions.lock().await;
+                    sessions.remove(&token)
+                };
+                match sender_option {
+                    Some(sender) => {
+                        let _ = sender.send(response.preshared_key);
+                    }
+                    // if desktop stopped listening for the result there will be no palce to send the result
+                    None => {
+                        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!({})))
+            }
+            None => {
+                error!("Remote MFA Unexpected core response, token was not returned");
+                Err(ApiError::Unexpected(String::new()))
+            }
+        }
+    } else {
+        error!("Received invalid gRPC response type: {payload:#?}");
+        Err(ApiError::InvalidResponseType)
+    }
+}

proto

@@ -1,7 +1,8 @@
 use std::{
+    collections::HashMap,
     fs::read_to_string,
     net::{IpAddr, Ipv4Addr, SocketAddr},
-    sync::atomic::Ordering,
+    sync::{atomic::Ordering, Arc},
     time::Duration,
 };
 
@@ -16,7 +17,7 @@ use axum::{
 use axum_extra::extract::cookie::Key;
 use clap::crate_version;
 use serde::Serialize;
-use tokio::{net::TcpListener, task::JoinSet};
+use tokio::{net::TcpListener, sync::oneshot, task::JoinSet};
 use tonic::transport::{Identity, Server, ServerTlsConfig};
 use tower_governor::{
     governor::GovernorConfigBuilder, key_extractor::SmartIpKeyExtractor, GovernorLayer,
@@ -42,6 +43,8 @@ const RATE_LIMITER_CLEANUP_PERIOD: Duration = Duration::from_secs(60);
 #[derive(Clone)]
 pub(crate) struct AppState {
     pub(crate) grpc_server: ProxyServer,
+    pub(crate) remote_mfa_sessions:
+        Arc<tokio::sync::Mutex<HashMap<String, oneshot::Sender<String>>>>,
     key: Key,
     url: Url,
 }
@@ -129,6 +132,7 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> {
     debug!("Setting up API server");
     let shared_state = AppState {
         grpc_server: grpc_server.clone(),
+        remote_mfa_sessions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
         // Generate secret key for encrypting cookies.
         key: Key::generate(),
         url: config.url.clone(),

src/handlers/desktop_client_mfa.rs

@@ -1,5 +1,5 @@
 {
-  "$schema": "https://biomejs.dev/schemas/2.1.2/schema.json",
+  "$schema": "https://biomejs.dev/schemas/2.2.0/schema.json",
   "vcs": { "enabled": false, "clientKind": "git", "useIgnoreFile": false },
   "files": {
     "ignoreUnknown": false,
@@ -8,7 +8,7 @@
       "!src/i18n/*.ts",
       "!src/i18n/*.tsx",
       "!src/i18n/i18n-util",
-      "!dist/**"
+      "!dist"
     ]
   },
   "formatter": {
@@ -40,7 +40,8 @@
         "noUnusedVariables": "error",
         "useExhaustiveDependencies": "error",
         "useHookAtTopLevel": "error",
-        "useJsxKeyInIterable": "error"
+        "useJsxKeyInIterable": "error",
+        "useUniqueElementIds": "off"
       },
       "security": { "noDangerouslySetInnerHtmlWithChildren": "error" },
       "style": {