desktop_client_mfa.rs

use axum::{Json, extract::State};
use axum_extra::extract::{PrivateCookieJar, cookie::Cookie};
use tracing::{debug, error, info, warn};

use crate::{
    enterprise::handlers::openid_login::{
        AuthenticationResponse, CSRF_COOKIE_NAME, FlowType, NONCE_COOKIE_NAME,
    },
    error::ApiError,
    handlers::get_core_response,
    http::AppState,
    proto::{ClientMfaOidcAuthenticateRequest, DeviceInfo, core_request, core_response},
};

#[instrument(level = "debug", skip(state))]
pub(super) async fn mfa_auth_callback(
    State(state): State<AppState>,
    device_info: DeviceInfo,
    mut private_cookies: PrivateCookieJar,
    Json(payload): Json<AuthenticationResponse>,
) -> Result<PrivateCookieJar, ApiError> {
    info!("Processing MFA authentication callback");
    debug!(
        "Received payload: state={}, flow_type={:?}",
        payload.state, payload.flow_type
    );

    match payload.flow_type {
        FlowType::Mfa => (),
        FlowType::Enrollment => {
            warn!(
                "Invalid flow type for MFA callback: {:?}",
                payload.flow_type
            );
            return Err(ApiError::BadRequest(
                "Invalid flow type for MFA callback".into(),
            ));
        }
    }

    debug!("Flow type validation passed: {:?}", payload.flow_type);

    let nonce = private_cookies
        .get(NONCE_COOKIE_NAME)
        .ok_or_else(|| {
            warn!("Nonce cookie not found in request");
            ApiError::Unauthorized("Nonce cookie not found".into())
        })?
        .value_trimmed()
        .to_string();

    let csrf = private_cookies
        .get(CSRF_COOKIE_NAME)
        .ok_or_else(|| {
            warn!("CSRF cookie not found in request");
            ApiError::Unauthorized("CSRF cookie not found".into())
        })?
        .value_trimmed()
        .to_string();

    debug!("Retrieved cookies successfully");

    if payload.state != csrf {
        warn!(
            "CSRF token mismatch: expected={csrf}, received={}",
            payload.state
        );
        return Err(ApiError::Unauthorized("CSRF token mismatch".into()));
    }

    debug!("CSRF token validation passed");

    private_cookies = private_cookies
        .remove(Cookie::from(NONCE_COOKIE_NAME))
        .remove(Cookie::from(CSRF_COOKIE_NAME));

    debug!("Removed security cookies");

    let request = ClientMfaOidcAuthenticateRequest {
        code: payload.code,
        nonce,
        state: payload.state,
    };

    debug!("Sending MFA OIDC authenticate request to core service");

    let rx = state.grpc_server.send(
        core_request::Payload::ClientMfaOidcAuthenticate(request),
        device_info,
    )?;

    let payload = get_core_response(rx, None).await?;

    if let core_response::Payload::Empty(()) = payload {
        info!("MFA authentication callback completed successfully");
        Ok(private_cookies)
    } else {
        error!(
            "Received invalid gRPC response type during handling the MFA OpenID authentication \
            callback"
        );
        Err(ApiError::InvalidResponseType)
    }
}