add code based mfa setup

.gitignore

@@ -4,3 +4,4 @@
 .direnv/
 .envrc
 /node_modules
+.env

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<AppState> {
     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) {

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);

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<AppState> {
+    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<AppState>,
+    device_info: DeviceInfo,
+    cookie_jar: PrivateCookieJar,
+    Json(req): Json<RegisterMfaCodeStartRequest>,
+) -> Result<Json<CodeMfaSetupStartResponse>, 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<AppState>,
+    device_info: DeviceInfo,
+    cookie_jar: PrivateCookieJar,
+    Json(req): Json<RegisterMfaCodeFinishRequest>,
+) -> Result<Json<CodeMfaSetupFinishResponse>, 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),
+    }
+}