Skip to main content

haste_server/auth_n/mfa/routes/
totp_verification.rs

1use crate::{
2    auth_n::{
3        mfa::utilities::user_mfa_to_totp,
4        oidc::routes::route_string::tenant_route_string,
5        session::{self, user::SessionAuthorizationState},
6    },
7    extract::{csrf_token::CSRFToken, path_tenant::TenantIdentifier},
8    services::ServerState,
9    ui::pages::mfa,
10};
11use axum::{
12    Form,
13    extract::{OriginalUri, Query, State},
14    response::{IntoResponse, Redirect, Response},
15};
16use axum_extra::{extract::Cached, routing::TypedPath};
17use haste_fhir_model::r4::generated::terminology::IssueType;
18use haste_fhir_operation_error::OperationOutcomeError;
19use haste_fhir_search::SearchEngine;
20use haste_fhir_terminology::FHIRTerminology;
21use haste_jwt::TenantId;
22use haste_repository::{
23    Repository,
24    admin::TenantModelAdmin,
25    types::{
26        mfa::{UserMFACredential, UserMFACredentialCreate, UserMFASearchClaims},
27        scope::UserId,
28        user::User,
29    },
30};
31use serde::Deserialize;
32use std::sync::Arc;
33use tower_sessions::Session;
34use url::form_urlencoded;
35
36fn is_safe_local_redirect_path(path: &str) -> bool {
37    path.starts_with('/') && !path.starts_with("//")
38}
39
40pub fn totp_verification_route(tenant: &TenantId, redirect_to: &str) -> String {
41    let route = tenant_route_string(tenant)
42        .join("mfa")
43        .join("totp-verification");
44
45    let query = form_urlencoded::Serializer::new(String::new())
46        .append_pair("redirect_to", redirect_to)
47        .finish();
48
49    format!("{}?{}", route.to_string_lossy(), query)
50}
51
52async fn get_required_totp_credentials<
53    Repo: Repository + Send + Sync,
54    Search: SearchEngine + Send + Sync,
55    Terminology: FHIRTerminology + Send + Sync,
56>(
57    state: &ServerState<Repo, Search, Terminology>,
58    tenant: &TenantId,
59    current_session: &Session,
60) -> Result<(User, Vec<UserMFACredential>), OperationOutcomeError> {
61    let user = match session::user::get_authorization_state(current_session).await? {
62        Some(SessionAuthorizationState::MFARequired { user }) => user,
63        _ => {
64            return Err(OperationOutcomeError::error(
65                IssueType::security(),
66                "MFA verification is not required.".to_string(),
67            ));
68        }
69    };
70
71    let credentials = TenantModelAdmin::<UserMFACredentialCreate, _, _, _, _>::search(
72        state.repo.as_ref(),
73        tenant,
74        &UserMFASearchClaims {
75            tenant: tenant.clone(),
76            user_id: UserId::new(user.id.clone()),
77            is_active: Some(true),
78        },
79    )
80    .await?
81    .into_iter()
82    .filter(|credential| credential.credential_type == "totp")
83    .collect::<Vec<_>>();
84
85    if credentials.is_empty() {
86        return Err(OperationOutcomeError::error(
87            IssueType::not_found(),
88            "No active TOTP credential found for this user.".to_string(),
89        ));
90    }
91
92    Ok((user, credentials))
93}
94
95#[derive(Deserialize)]
96pub struct TOTPVerificationPOSTBody {
97    pub csrf_token: String,
98    pub otp_code: String,
99}
100
101#[derive(Deserialize)]
102pub struct TOTPVerificationQuery {
103    pub redirect_to: String,
104}
105
106#[derive(TypedPath, Deserialize)]
107#[typed_path("/totp-verification")]
108pub struct TOTPVerificationGET;
109
110#[derive(TypedPath, Deserialize)]
111#[typed_path("/totp-verification")]
112pub struct TOTPVerificationPOST;
113
114pub async fn totp_verification_get<
115    Repo: Repository + Send + Sync,
116    Search: SearchEngine + Send + Sync,
117    Terminology: FHIRTerminology + Send + Sync,
118>(
119    _: TOTPVerificationGET,
120    Query(query): Query<TOTPVerificationQuery>,
121    uri: OriginalUri,
122    CSRFToken(csrf_token): CSRFToken,
123    Cached(TenantIdentifier { tenant }): Cached<TenantIdentifier>,
124    State(state): State<Arc<ServerState<Repo, Search, Terminology>>>,
125    Cached(current_session): Cached<Session>,
126) -> Result<Response, OperationOutcomeError> {
127    if !is_safe_local_redirect_path(&query.redirect_to) {
128        return Err(OperationOutcomeError::error(
129            IssueType::security(),
130            "Invalid MFA redirect target.".to_string(),
131        ));
132    }
133
134    let (_, credentials) =
135        get_required_totp_credentials(state.as_ref(), &tenant, &current_session).await?;
136
137    Ok(mfa::totp_verification::totp_entry_html(
138        &tenant,
139        &csrf_token,
140        credentials[0].totp_digits as usize,
141        &uri.to_string(),
142        None,
143    )
144    .into_response())
145}
146
147pub async fn totp_verification_post<
148    Repo: Repository + Send + Sync,
149    Search: SearchEngine + Send + Sync,
150    Terminology: FHIRTerminology + Send + Sync,
151>(
152    _: TOTPVerificationPOST,
153    Query(query): Query<TOTPVerificationQuery>,
154    uri: OriginalUri,
155    CSRFToken(csrf_token): CSRFToken,
156    Cached(TenantIdentifier { tenant }): Cached<TenantIdentifier>,
157    State(state): State<Arc<ServerState<Repo, Search, Terminology>>>,
158    Cached(current_session): Cached<Session>,
159    Form(form_data): Form<TOTPVerificationPOSTBody>,
160) -> Result<Response, OperationOutcomeError> {
161    if form_data.csrf_token != csrf_token {
162        return Err(OperationOutcomeError::error(
163            IssueType::security(),
164            "Invalid CSRF token.".to_string(),
165        ));
166    }
167
168    if !is_safe_local_redirect_path(&query.redirect_to) {
169        return Err(OperationOutcomeError::error(
170            IssueType::security(),
171            "Invalid MFA redirect target.".to_string(),
172        ));
173    }
174
175    let (user, credentials) =
176        get_required_totp_credentials(state.as_ref(), &tenant, &current_session).await?;
177    let digits = credentials[0].totp_digits as usize;
178
179    let mut is_otp_valid = false;
180
181    for credential in credentials {
182        let totp = user_mfa_to_totp(
183            state.secret_provider.as_ref(),
184            &state.config,
185            &user,
186            credential,
187        )
188        .await?;
189
190        if totp.check_current(&form_data.otp_code).map_err(|_e| {
191            OperationOutcomeError::error(
192                IssueType::security(),
193                "Invalid verification code.".to_string(),
194            )
195        })? {
196            is_otp_valid = true;
197            break;
198        }
199    }
200
201    if !is_otp_valid {
202        return Ok(mfa::totp_verification::totp_entry_html(
203            &tenant,
204            &csrf_token,
205            digits,
206            &uri.to_string(),
207            Some(vec![
208                "Invalid verification code. Please try again.".to_string(),
209            ]),
210        )
211        .into_response());
212    }
213
214    session::user::set_completed_authorization_state(&current_session, user).await?;
215
216    Ok(Redirect::to(&query.redirect_to).into_response())
217}