diff --git a/.gitignore b/.gitignore index 54bf0c2..f2379c7 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ .direnv/ .envrc /node_modules +.env diff --git a/proto b/proto index dbb08b7..4688d4d 160000 --- a/proto +++ b/proto @@ -1 +1 @@ -Subproject commit dbb08b75e1b009c48a73bfc07cb25bea1838243a +Subproject commit 4688d4d587246e09d800c03963721b437f3a3eca diff --git a/src/handlers/enrollment.rs b/src/handlers/enrollment.rs index eed4c98..7d8b99b 100644 --- a/src/handlers/enrollment.rs +++ b/src/handlers/enrollment.rs @@ -2,6 +2,8 @@ use axum::{extract::State, routing::post, Json, Router}; use axum_extra::extract::{cookie::Cookie, PrivateCookieJar}; use time::OffsetDateTime; +use super::register_mfa::router as register_mfa_router; + use crate::{ error::ApiError, handlers::{get_core_response, mobile_client::register_mobile_auth}, @@ -14,6 +16,7 @@ use crate::{ pub(crate) fn router() -> Router { Router::new() + .nest("/register-mfa", register_mfa_router()) .route("/start", post(start_enrollment_process)) .route("/activate_user", post(activate_user)) .route("/create_device", post(create_device)) @@ -82,7 +85,7 @@ async fn activate_user( .grpc_server .send(core_request::Payload::ActivateUser(req), device_info)?; let payload = get_core_response(rx).await?; - debug!("Receving payload from the core service. Trying to remove private cookie..."); + 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:?}"); if let Some(cookie) = private_cookies.get(ENROLLMENT_COOKIE_NAME) { diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs index cb74e02..bb49ab7 100644 --- a/src/handlers/mod.rs +++ b/src/handlers/mod.rs @@ -14,6 +14,7 @@ pub(crate) mod enrollment; pub(crate) mod mobile_client; pub(crate) mod password_reset; pub(crate) mod polling; +pub(crate) mod register_mfa; // Timeout for awaiting response from Defguard Core. const CORE_RESPONSE_TIMEOUT: Duration = Duration::from_secs(5); diff --git a/src/handlers/register_mfa.rs b/src/handlers/register_mfa.rs new file mode 100644 index 0000000..caef630 --- /dev/null +++ b/src/handlers/register_mfa.rs @@ -0,0 +1,99 @@ +use serde::Deserialize; + +use axum::{extract::State, response::IntoResponse, routing::post, Json, Router}; +use axum_extra::extract::PrivateCookieJar; + +use crate::{ + error::ApiError, + handlers::get_core_response, + http::{AppState, ENROLLMENT_COOKIE_NAME}, + proto::{ + core_request, core_response, CodeMfaSetupFinishRequest, CodeMfaSetupFinishResponse, + CodeMfaSetupStartRequest, CodeMfaSetupStartResponse, DeviceInfo, MfaMethod, + }, +}; + +pub(crate) fn router() -> Router { + Router::new() + .route("/code/start", post(register_code_mfa_start)) + .route("/code/finish", post(register_code_mfa_finish)) +} + +#[derive(Debug, Clone, Deserialize)] +struct RegisterMfaCodeStartRequest { + pub method: MfaMethod, +} + +#[instrument(level = "debug", skip(state, req))] +async fn register_code_mfa_start( + State(state): State, + device_info: DeviceInfo, + cookie_jar: PrivateCookieJar, + Json(req): Json, +) -> Result, impl IntoResponse> { + debug!("Register code MFA started"); + let token = cookie_jar + .get(ENROLLMENT_COOKIE_NAME) + .ok_or_else(|| ApiError::Unauthorized(String::new()))? + .value() + .to_string(); + + if req.method != MfaMethod::Email && req.method != MfaMethod::Totp { + error!("Requested method not supported"); + return Err(ApiError::BadRequest("Method not supported.".to_string())); + } + + let rx = state.grpc_server.send( + core_request::Payload::CodeMfaSetupStart(CodeMfaSetupStartRequest { + token, + method: req.method.into(), + }), + device_info, + )?; + let payload = get_core_response(rx).await?; + match payload { + core_response::Payload::CodeMfaSetupStartResponse(response) => Ok(Json(response)), + _ => Err(ApiError::InvalidResponseType), + } +} + +#[derive(Debug, Clone, Deserialize)] +struct RegisterMfaCodeFinishRequest { + pub code: String, + pub method: MfaMethod, +} + +#[instrument(level = "debug", skip(state, req))] +async fn register_code_mfa_finish( + State(state): State, + device_info: DeviceInfo, + cookie_jar: PrivateCookieJar, + Json(req): Json, +) -> Result, impl IntoResponse> { + let token = cookie_jar + .get(ENROLLMENT_COOKIE_NAME) + .ok_or_else(|| ApiError::Unauthorized(String::new()))? + .value() + .to_string(); + + let code = req.code; + let method = req.method; + + if method != MfaMethod::Totp && method != MfaMethod::Email { + return Err(ApiError::BadRequest("Method not supported".to_string())); + } + + let rx = state.grpc_server.send( + core_request::Payload::CodeMfaSetupFinish(CodeMfaSetupFinishRequest { + token, + code, + method: method as i32, + }), + device_info, + )?; + let payload = get_core_response(rx).await?; + match payload { + core_response::Payload::CodeMfaSetupFinishResponse(response) => Ok(Json(response)), + _ => Err(ApiError::InvalidResponseType), + } +}