Skip to main content

haste_server/auth_n/mfa/routes/
totp_verification.rs

1use crate::auth_n::constant_time;
2use crate::auth_n::error::AuthError;
3use crate::{
4    auth_n::{
5        mfa::utilities::user_mfa_to_totp,
6        oidc::routes::route_string::tenant_route_string,
7        session::{self, user::SessionAuthorizationState},
8    },
9    extract::csrf_token::CSRFToken,
10    services::ServerState,
11    ui::{components::TenantContext, pages::mfa},
12};
13use axum::{
14    Form,
15    extract::{FromRequestParts, OriginalUri, Query, State},
16    response::{IntoResponse, Redirect, Response},
17};
18use axum_extra::{extract::Cached, routing::TypedPath};
19use haste_fhir_model::r4::generated::terminology::IssueType;
20use haste_fhir_operation_error::OperationOutcomeError;
21use haste_fhir_search::SearchEngine;
22use haste_fhir_terminology::FHIRTerminology;
23use haste_jwt::TenantId;
24use haste_repository::{
25    Repository,
26    admin::TenantModelAdmin,
27    types::{
28        mfa::{UserMFACredential, UserMFACredentialCreate, UserMFASearchClaims},
29        scope::UserId,
30        user::User,
31    },
32};
33use serde::Deserialize;
34use std::sync::Arc;
35use tower_sessions::Session;
36use url::{Url, form_urlencoded};
37
38/// Whether `path` stays on this origin when a browser follows it.
39///
40/// Browsers parse a `Location` with the WHATWG URL algorithm, which treats a
41/// backslash like a slash for http(s), so `/\evil.example` is navigated to
42/// as `//evil.example`: a second origin, which the plain `//` check missed.
43/// The path is therefore resolved against a stand-in origin the way a browser
44/// would, and accepted only when it still points at that origin.
45fn is_safe_local_redirect_path(path: &str) -> bool {
46    if !path.starts_with('/') || path.chars().any(|c| c.is_control() || c == '\\') {
47        return false;
48    }
49
50    let Ok(base) = Url::parse("http://redirect.invalid") else {
51        return false;
52    };
53
54    base.join(path).is_ok_and(|resolved| {
55        resolved.scheme() == "http" && resolved.host_str() == Some("redirect.invalid")
56    })
57}
58
59pub fn totp_verification_route(tenant: &TenantId, redirect_to: &str) -> String {
60    let route = tenant_route_string(tenant)
61        .join("mfa")
62        .join("totp-verification");
63
64    let query = form_urlencoded::Serializer::new(String::new())
65        .append_pair("redirect_to", redirect_to)
66        .finish();
67
68    format!("{}?{}", route.to_string_lossy(), query)
69}
70
71async fn get_required_totp_credentials<
72    Repo: Repository + Send + Sync,
73    Search: SearchEngine + Send + Sync,
74    Terminology: FHIRTerminology + Send + Sync,
75>(
76    state: &ServerState<Repo, Search, Terminology>,
77    tenant: &TenantId,
78    current_session: &Session,
79) -> Result<(User, Vec<UserMFACredential>), OperationOutcomeError> {
80    let user = match session::user::get_authorization_state(current_session).await? {
81        Some(SessionAuthorizationState::MFARequired { user }) => user,
82        _ => {
83            return Err(OperationOutcomeError::error(
84                IssueType::security(),
85                "MFA verification is not required.".to_string(),
86            ));
87        }
88    };
89
90    let credentials = TenantModelAdmin::<UserMFACredentialCreate, _, _, _, _>::search(
91        state.repo.as_ref(),
92        tenant,
93        &UserMFASearchClaims {
94            tenant: tenant.clone(),
95            user_id: UserId::new(user.id.clone()),
96            is_active: Some(true),
97        },
98    )
99    .await?
100    .into_iter()
101    .filter(|credential| credential.credential_type == "totp")
102    .collect::<Vec<_>>();
103
104    if credentials.is_empty() {
105        return Err(OperationOutcomeError::error(
106            IssueType::not_found(),
107            "No active TOTP credential found for this user.".to_string(),
108        ));
109    }
110
111    Ok((user, credentials))
112}
113
114#[derive(Deserialize)]
115pub struct TOTPVerificationPOSTBody {
116    pub csrf_token: String,
117    pub otp_code: String,
118}
119
120#[derive(Deserialize)]
121pub struct TOTPVerificationQuery {
122    pub redirect_to: String,
123}
124
125/// Where a verification request sits in the login flow: the URI it was made
126/// to, which the form posts back to, and where to go once the code is verified.
127#[derive(FromRequestParts)]
128pub struct TOTPVerificationLocation {
129    #[from_request(via(Query))]
130    query: TOTPVerificationQuery,
131    uri: OriginalUri,
132}
133
134#[derive(TypedPath, Deserialize)]
135#[typed_path("/totp-verification")]
136pub struct TOTPVerificationGET;
137
138#[derive(TypedPath, Deserialize)]
139#[typed_path("/totp-verification")]
140pub struct TOTPVerificationPOST;
141
142pub async fn totp_verification_get<
143    Repo: Repository + Send + Sync,
144    Search: SearchEngine + Send + Sync,
145    Terminology: FHIRTerminology + Send + Sync,
146>(
147    _: TOTPVerificationGET,
148    TOTPVerificationLocation { query, uri }: TOTPVerificationLocation,
149    CSRFToken(csrf_token): CSRFToken,
150    Cached(TenantContext { tenant, branding }): Cached<TenantContext>,
151    State(state): State<Arc<ServerState<Repo, Search, Terminology>>>,
152    Cached(current_session): Cached<Session>,
153) -> Result<Response, OperationOutcomeError> {
154    if !is_safe_local_redirect_path(&query.redirect_to) {
155        return Err(AuthError::InvalidMfaRedirectTarget.into());
156    }
157
158    let (_, credentials) =
159        get_required_totp_credentials(state.as_ref(), &tenant, &current_session).await?;
160
161    Ok(mfa::totp_verification::totp_entry_html(
162        &tenant,
163        &csrf_token,
164        credentials[0].totp_digits as usize,
165        &uri.to_string(),
166        None,
167        Some(&branding),
168    )
169    .into_response())
170}
171
172pub async fn totp_verification_post<
173    Repo: Repository + Send + Sync,
174    Search: SearchEngine + Send + Sync,
175    Terminology: FHIRTerminology + Send + Sync,
176>(
177    _: TOTPVerificationPOST,
178    TOTPVerificationLocation { query, uri }: TOTPVerificationLocation,
179    CSRFToken(csrf_token): CSRFToken,
180    Cached(TenantContext { tenant, branding }): Cached<TenantContext>,
181    State(state): State<Arc<ServerState<Repo, Search, Terminology>>>,
182    Cached(current_session): Cached<Session>,
183    Form(form_data): Form<TOTPVerificationPOSTBody>,
184) -> Result<Response, OperationOutcomeError> {
185    if !constant_time::str_eq(&form_data.csrf_token, &csrf_token) {
186        return Err(AuthError::InvalidCsrfToken.into());
187    }
188
189    if !is_safe_local_redirect_path(&query.redirect_to) {
190        return Err(AuthError::InvalidMfaRedirectTarget.into());
191    }
192
193    let (user, credentials) =
194        get_required_totp_credentials(state.as_ref(), &tenant, &current_session).await?;
195    let digits = credentials[0].totp_digits as usize;
196
197    let mut is_otp_valid = false;
198
199    for credential in credentials {
200        let totp = user_mfa_to_totp(
201            state.secret_provider.as_ref(),
202            &state.config,
203            &user,
204            credential,
205        )
206        .await?;
207
208        if totp.check_current(&form_data.otp_code).map_err(|_e| {
209            OperationOutcomeError::error(
210                IssueType::security(),
211                "Invalid verification code.".to_string(),
212            )
213        })? {
214            is_otp_valid = true;
215            break;
216        }
217    }
218
219    if !is_otp_valid {
220        return Ok(mfa::totp_verification::totp_entry_html(
221            &tenant,
222            &csrf_token,
223            digits,
224            &uri.to_string(),
225            Some(vec![
226                "Invalid verification code. Please try again.".to_string(),
227            ]),
228            Some(&branding),
229        )
230        .into_response());
231    }
232
233    session::user::rotate_session_id(&current_session).await?;
234    session::user::set_completed_authorization_state(&current_session, user).await?;
235
236    Ok(Redirect::to(&query.redirect_to).into_response())
237}