allowthem_server/
mfa_page_routes.rs

1use std::sync::Arc;
2
3use axum::Form;
4use axum::Router;
5use axum::extract::{Extension, Query};
6use axum::http::HeaderMap;
7use axum::http::StatusCode;
8use axum::http::Uri;
9use axum::http::header::{LOCATION, SET_COOKIE, USER_AGENT};
10use axum::response::{IntoResponse, Response};
11use axum::routing::{get, post};
12use axum_htmx::{HxBoosted, HxRequest};
13use chrono::Utc;
14use minijinja::{Environment, context};
15use serde::Deserialize;
16
17use allowthem_core::applications::BrandingConfig;
18use allowthem_core::totp::totp_uri;
19use allowthem_core::{AllowThem, AuditEvent, AuthError, sessions};
20use qrcode::QrCode;
21use qrcode::render::svg;
22
23use crate::branding::{DefaultBranding, branding_context, default_branding_ref, resolve_branding};
24use crate::browser_error::BrowserError;
25use crate::csrf::CsrfToken;
26use crate::error::BrowserAuthRedirect;
27
28/// Error shown when a wrong TOTP code is entered during MFA setup confirmation.
29const SETUP_INVALID_CODE: &str = "Invalid TOTP code";
30
31/// Error shown when a wrong TOTP code is entered on the MFA challenge page.
32const CHALLENGE_INVALID_TOTP: &str = "Invalid TOTP or recovery code";
33
34/// Error shown when a wrong recovery code is entered on the MFA challenge page.
35const CHALLENGE_INVALID_RECOVERY: &str = "Invalid recovery code";
36
37#[derive(Clone)]
38struct MfaPageConfig {
39    templates: Arc<Environment<'static>>,
40    is_production: bool,
41    base_url: String,
42}
43
44// ---------------------------------------------------------------------------
45// Helpers
46// ---------------------------------------------------------------------------
47
48fn client_ip(headers: &HeaderMap) -> Option<String> {
49    headers
50        .get("x-forwarded-for")
51        .and_then(|v| v.to_str().ok())
52        .and_then(|s| s.split(',').next())
53        .map(|s| s.trim().to_string())
54}
55
56/// Generate a QR code SVG data URI from the given text.
57///
58/// Returns a `data:image/svg+xml,...` URI suitable for an `<img src>` attribute.
59/// Falls back to an empty string if encoding fails (the template will show
60/// the manual-entry secret as an alternative).
61fn qr_data_uri(text: &str) -> String {
62    let code = match QrCode::new(text.as_bytes()) {
63        Ok(c) => c,
64        Err(_) => return String::new(),
65    };
66    let svg_str = code
67        .render()
68        .min_dimensions(200, 200)
69        .dark_color(svg::Color("#000000"))
70        .light_color(svg::Color("#ffffff"))
71        .build();
72    // URI-encode the SVG for use in a data URI (minimal encoding).
73    let encoded = svg_str
74        .replace('#', "%23")
75        .replace('<', "%3C")
76        .replace('>', "%3E")
77        .replace('"', "'");
78    format!("data:image/svg+xml,{encoded}")
79}
80
81/// Extract the host from a base URL for use as the TOTP issuer.
82///
83/// Strips the scheme and path, and also strips the port (the totp-rs
84/// library rejects issuer strings containing colons).
85fn derive_issuer(base_url: &str) -> String {
86    base_url
87        .trim_start_matches("https://")
88        .trim_start_matches("http://")
89        .split('/')
90        .next()
91        .unwrap_or("allowthem")
92        .split(':')
93        .next()
94        .unwrap_or("allowthem")
95        .to_string()
96}
97
98/// Validate session cookie and return the authenticated user.
99///
100/// On failure, returns a 303 redirect to `/login?next={path}` — matching
101/// `BrowserAuthUser` rejection semantics without requiring `Arc<dyn AuthClient>`
102/// in the router state.
103async fn require_browser_user(
104    ath: &AllowThem,
105    headers: &HeaderMap,
106    path: &str,
107) -> Result<allowthem_core::types::User, Response> {
108    let cookie_header = headers
109        .get(axum::http::header::COOKIE)
110        .and_then(|v| v.to_str().ok())
111        .ok_or_else(|| BrowserAuthRedirect::new(path).into_response())?;
112
113    let token = ath
114        .parse_session_cookie(cookie_header)
115        .ok_or_else(|| BrowserAuthRedirect::new(path).into_response())?;
116
117    let ttl = ath.session_config().ttl;
118    let session = ath
119        .db()
120        .validate_session(&token, ttl)
121        .await
122        .map_err(|err| {
123            tracing::error!("session validation error: {err}");
124            BrowserAuthRedirect::new(path).into_response()
125        })?
126        .ok_or_else(|| BrowserAuthRedirect::new(path).into_response())?;
127
128    match ath.db().get_user(session.user_id).await {
129        Ok(user) if user.is_active => Ok(user),
130        Ok(_) => Err(BrowserAuthRedirect::new(path).into_response()),
131        Err(AuthError::NotFound) => Err(BrowserAuthRedirect::new(path).into_response()),
132        Err(err) => {
133            tracing::error!("user lookup error: {err}");
134            Err(BrowserAuthRedirect::new(path).into_response())
135        }
136    }
137}
138
139// ---------------------------------------------------------------------------
140// Setup-side routes (authenticated, CSRF-protected)
141// ---------------------------------------------------------------------------
142
143/// Render just the `_auth_main_mfa_setup.html` partial plus the
144/// `_auth_oob_head.html` OOB head swap, for HTMX fragment responses.
145///
146/// The `.wf-note` style used by the TOTP info box lives in the shell's
147/// `<head>` `<style>` block. Fragment responses don't update `<head>`,
148/// but mfa_setup is always reached from an authenticated /settings page,
149/// so the full page (and its head styles) loads before any HX swap.
150fn render_mfa_setup_fragment(
151    config: &MfaPageConfig,
152    csrf_token: &str,
153    totp_uri: &str,
154    qr_data_uri: &str,
155    secret: &str,
156    error: &str,
157    branding: Option<&BrandingConfig>,
158) -> Result<axum::response::Html<String>, BrowserError> {
159    let ctx = context! {
160        csrf_token,
161        totp_uri,
162        qr_data_uri,
163        secret,
164        error,
165        is_production => config.is_production,
166        page_title => "Set up two-factor authentication — allowthem",
167        status_hint => "ENABLE 2FA",
168        ..branding_context(branding),
169    };
170
171    let main = crate::browser_templates::render(
172        &config.templates,
173        "_partials/_auth_main_mfa_setup.html",
174        ctx.clone(),
175    )?;
176    let oob =
177        crate::browser_templates::render(&config.templates, "_partials/_auth_oob_head.html", ctx)?;
178    Ok(axum::response::Html(format!("{}{}", main.0, oob.0)))
179}
180
181/// Render just the `_auth_main_mfa_recovery.html` partial plus the
182/// `_auth_oob_head.html` OOB head swap, for HTMX fragment responses.
183fn render_mfa_recovery_fragment(
184    config: &MfaPageConfig,
185    recovery_codes: &[String],
186    branding: Option<&BrandingConfig>,
187) -> Result<axum::response::Html<String>, BrowserError> {
188    let ctx = context! {
189        recovery_codes,
190        is_production => config.is_production,
191        page_title => "Recovery codes — allowthem",
192        status_hint => "RECOVERY CODES",
193        ..branding_context(branding),
194    };
195
196    let main = crate::browser_templates::render(
197        &config.templates,
198        "_partials/_auth_main_mfa_recovery.html",
199        ctx.clone(),
200    )?;
201    let oob =
202        crate::browser_templates::render(&config.templates, "_partials/_auth_oob_head.html", ctx)?;
203    Ok(axum::response::Html(format!("{}{}", main.0, oob.0)))
204}
205
206/// GET /settings/mfa/setup — show QR URI, base32 secret, and TOTP code input.
207///
208/// Idempotent: if a pending (non-enabled) secret exists, reuses it.
209/// Only creates a new secret on first visit.
210#[allow(clippy::too_many_arguments)]
211async fn get_mfa_setup(
212    Extension(ath): Extension<AllowThem>,
213    Extension(config): Extension<MfaPageConfig>,
214    default_branding: Option<Extension<Arc<DefaultBranding>>>,
215    uri: Uri,
216    csrf: CsrfToken,
217    headers: HeaderMap,
218    HxBoosted(boosted): HxBoosted,
219    HxRequest(request): HxRequest,
220) -> Result<Response, BrowserError> {
221    let user = match require_browser_user(&ath, &headers, uri.path()).await {
222        Ok(u) => u,
223        Err(redirect) => return Ok(redirect),
224    };
225
226    let default = default_branding_ref(&default_branding);
227    let branding = resolve_branding(&ath, None, default).await;
228
229    // Reuse pending secret if one exists; create only on first visit
230    let secret = match ath.get_pending_mfa_secret(user.id).await? {
231        Some(s) => s,
232        None => ath.create_mfa_secret(user.id).await?,
233    };
234
235    let issuer = derive_issuer(&config.base_url);
236    let uri = totp_uri(&secret, user.email.as_str(), &issuer);
237    let qr = qr_data_uri(&uri);
238
239    if request && !boosted {
240        let html = render_mfa_setup_fragment(
241            &config,
242            csrf.as_str(),
243            &uri,
244            &qr,
245            &secret,
246            "",
247            branding.as_ref(),
248        )?;
249        return Ok(html.into_response());
250    }
251
252    let html = crate::browser_templates::render(
253        &config.templates,
254        "mfa_setup.html",
255        context! {
256            csrf_token => csrf.as_str(),
257            secret => &secret,
258            totp_uri => &uri,
259            qr_data_uri => &qr,
260            error => "",
261            is_production => config.is_production,
262            ..branding_context(branding.as_ref()),
263        },
264    )?;
265    Ok(html.into_response())
266}
267
268#[derive(Deserialize)]
269pub struct MfaConfirmForm {
270    code: String,
271    #[allow(dead_code)]
272    csrf_token: String,
273}
274
275/// POST /settings/mfa/confirm — verify TOTP code and enable MFA.
276///
277/// On success, renders recovery codes page directly (no redirect).
278/// On failure, re-renders setup page with error.
279#[allow(clippy::too_many_arguments)]
280async fn post_mfa_confirm(
281    Extension(ath): Extension<AllowThem>,
282    Extension(config): Extension<MfaPageConfig>,
283    default_branding: Option<Extension<Arc<DefaultBranding>>>,
284    uri: Uri,
285    csrf: CsrfToken,
286    headers: HeaderMap,
287    HxBoosted(boosted): HxBoosted,
288    HxRequest(request): HxRequest,
289    Form(form): Form<MfaConfirmForm>,
290) -> Result<Response, BrowserError> {
291    let user = match require_browser_user(&ath, &headers, uri.path()).await {
292        Ok(u) => u,
293        Err(redirect) => return Ok(redirect),
294    };
295
296    let default = default_branding_ref(&default_branding);
297    let branding = resolve_branding(&ath, None, default).await;
298
299    let ip = client_ip(&headers);
300    let ua = headers.get(USER_AGENT).and_then(|v| v.to_str().ok());
301
302    match ath.enable_mfa(user.id, &form.code).await {
303        Ok(recovery_codes) => {
304            let _ = ath
305                .db()
306                .log_audit(
307                    AuditEvent::MfaEnabled,
308                    Some(&user.id),
309                    None,
310                    ip.as_deref(),
311                    ua,
312                    None,
313                )
314                .await;
315
316            if request && !boosted {
317                let html =
318                    render_mfa_recovery_fragment(&config, &recovery_codes, branding.as_ref())?;
319                return Ok(html.into_response());
320            }
321
322            let html = crate::browser_templates::render(
323                &config.templates,
324                "mfa_recovery.html",
325                context! {
326                    recovery_codes => &recovery_codes,
327                    is_production => config.is_production,
328                    ..branding_context(branding.as_ref()),
329                },
330            )?;
331            Ok(html.into_response())
332        }
333        Err(allowthem_core::AuthError::InvalidTotpCode) => {
334            // Re-render setup page with error
335            let secret = ath
336                .get_pending_mfa_secret(user.id)
337                .await?
338                .unwrap_or_default();
339            let issuer = derive_issuer(&config.base_url);
340            let uri = totp_uri(&secret, user.email.as_str(), &issuer);
341            let qr = qr_data_uri(&uri);
342
343            let html = crate::browser_templates::render(
344                &config.templates,
345                "mfa_setup.html",
346                context! {
347                    csrf_token => csrf.as_str(),
348                    secret => &secret,
349                    totp_uri => &uri,
350                    qr_data_uri => &qr,
351                    error => SETUP_INVALID_CODE,
352                    is_production => config.is_production,
353                    ..branding_context(branding.as_ref()),
354                },
355            )?;
356            Ok(html.into_response())
357        }
358        Err(e) => Err(BrowserError::Auth(e)),
359    }
360}
361
362#[derive(Deserialize)]
363pub struct MfaDisableForm {
364    #[allow(dead_code)]
365    csrf_token: String,
366}
367
368/// POST /settings/mfa/disable — disable MFA and redirect to settings.
369async fn post_mfa_disable(
370    Extension(ath): Extension<AllowThem>,
371    uri: Uri,
372    headers: HeaderMap,
373    Form(_form): Form<MfaDisableForm>,
374) -> Result<Response, BrowserError> {
375    let user = match require_browser_user(&ath, &headers, uri.path()).await {
376        Ok(u) => u,
377        Err(redirect) => return Ok(redirect),
378    };
379
380    let ip = client_ip(&headers);
381    let ua = headers.get(USER_AGENT).and_then(|v| v.to_str().ok());
382
383    ath.disable_mfa(user.id).await?;
384
385    let _ = ath
386        .db()
387        .log_audit(
388            AuditEvent::MfaDisabled,
389            Some(&user.id),
390            None,
391            ip.as_deref(),
392            ua,
393            None,
394        )
395        .await;
396
397    Ok((StatusCode::SEE_OTHER, [(LOCATION, "/settings".to_string())]).into_response())
398}
399
400#[derive(Deserialize)]
401struct RegenerateCodesForm {
402    #[allow(dead_code)]
403    csrf_token: String,
404}
405
406/// POST /settings/mfa/recovery-codes/regenerate — regenerate and show new codes.
407#[allow(clippy::too_many_arguments)]
408async fn post_regenerate_recovery_codes(
409    Extension(ath): Extension<AllowThem>,
410    Extension(config): Extension<MfaPageConfig>,
411    default_branding: Option<Extension<Arc<DefaultBranding>>>,
412    uri: Uri,
413    headers: HeaderMap,
414    HxBoosted(boosted): HxBoosted,
415    HxRequest(request): HxRequest,
416    Form(_form): Form<RegenerateCodesForm>,
417) -> Result<Response, BrowserError> {
418    let user = match require_browser_user(&ath, &headers, uri.path()).await {
419        Ok(u) => u,
420        Err(redirect) => return Ok(redirect),
421    };
422
423    let has_mfa = ath.db().has_mfa_enabled(user.id).await?;
424    if !has_mfa {
425        return Ok((StatusCode::SEE_OTHER, [(LOCATION, "/settings".to_string())]).into_response());
426    }
427
428    let recovery_codes = ath.regenerate_recovery_codes(user.id).await?;
429
430    let default = default_branding_ref(&default_branding);
431    let branding = resolve_branding(&ath, None, default).await;
432
433    if request && !boosted {
434        let html = render_mfa_recovery_fragment(&config, &recovery_codes, branding.as_ref())?;
435        return Ok(html.into_response());
436    }
437
438    let html = crate::browser_templates::render(
439        &config.templates,
440        "mfa_recovery.html",
441        context! {
442            recovery_codes => &recovery_codes,
443            is_production => config.is_production,
444            ..branding_context(branding.as_ref()),
445        },
446    )?;
447    Ok(html.into_response())
448}
449
450// ---------------------------------------------------------------------------
451// Challenge routes (mid-login, no session — outside CSRF layer)
452// ---------------------------------------------------------------------------
453
454#[derive(Deserialize)]
455pub struct ChallengeQuery {
456    token: String,
457}
458
459/// Render just the `_auth_main_mfa_challenge.html` partial plus the
460/// `_auth_oob_head.html` OOB head swap, for HTMX fragment responses.
461///
462/// The CSS-only recovery-code toggle's `<style>` lives in the shell's
463/// `<head>`. Fragment responses don't update `<head>`, but mfa_challenge
464/// is always reached mid-flow from /login, so the full page (and its
465/// head styles) loads before any HX swap; re-rendering via fragment
466/// on errors is safe because the styles are already in the document.
467fn render_mfa_challenge_fragment(
468    config: &MfaPageConfig,
469    mfa_token: &str,
470    error: &str,
471    branding: Option<&BrandingConfig>,
472) -> Result<axum::response::Html<String>, BrowserError> {
473    let ctx = context! {
474        mfa_token,
475        error,
476        is_production => config.is_production,
477        page_title => "Two-factor authentication — allowthem",
478        status_hint => "TWO-FACTOR",
479        ..branding_context(branding),
480    };
481
482    let main = crate::browser_templates::render(
483        &config.templates,
484        "_partials/_auth_main_mfa_challenge.html",
485        ctx.clone(),
486    )?;
487    let oob =
488        crate::browser_templates::render(&config.templates, "_partials/_auth_oob_head.html", ctx)?;
489    Ok(axum::response::Html(format!("{}{}", main.0, oob.0)))
490}
491
492/// GET /mfa/challenge — render TOTP code input form.
493async fn get_mfa_challenge(
494    Extension(ath): Extension<AllowThem>,
495    Extension(config): Extension<MfaPageConfig>,
496    default_branding: Option<Extension<Arc<DefaultBranding>>>,
497    Query(query): Query<ChallengeQuery>,
498    HxBoosted(boosted): HxBoosted,
499    HxRequest(request): HxRequest,
500) -> Result<Response, BrowserError> {
501    // Validate token is still alive (don't consume it)
502    let user_id = ath.db().validate_mfa_challenge(&query.token).await?;
503    if user_id.is_none() {
504        // Invalid or expired token — redirect to login
505        return Ok((StatusCode::SEE_OTHER, [(LOCATION, "/login".to_string())]).into_response());
506    }
507
508    let default = default_branding_ref(&default_branding);
509    let branding = resolve_branding(&ath, None, default).await;
510
511    if request && !boosted {
512        let html = render_mfa_challenge_fragment(&config, &query.token, "", branding.as_ref())?;
513        return Ok(html.into_response());
514    }
515
516    let html = crate::browser_templates::render(
517        &config.templates,
518        "mfa_challenge.html",
519        context! {
520            mfa_token => &query.token,
521            error => "",
522            is_production => config.is_production,
523            ..branding_context(branding.as_ref()),
524        },
525    )?;
526    Ok(html.into_response())
527}
528
529#[derive(Deserialize)]
530pub struct MfaChallengeForm {
531    mfa_token: String,
532    #[serde(default)]
533    code: Option<String>,
534    #[serde(default)]
535    recovery_code: Option<String>,
536    #[serde(default)]
537    use_recovery: Option<String>,
538}
539
540/// POST /mfa/challenge — verify TOTP code or recovery code, create session.
541async fn post_mfa_challenge(
542    Extension(ath): Extension<AllowThem>,
543    Extension(config): Extension<MfaPageConfig>,
544    default_branding: Option<Extension<Arc<DefaultBranding>>>,
545    headers: HeaderMap,
546    Form(form): Form<MfaChallengeForm>,
547) -> Result<Response, BrowserError> {
548    let default = default_branding_ref(&default_branding);
549    let branding = resolve_branding(&ath, None, default).await;
550    let ip = headers
551        .get("x-forwarded-for")
552        .and_then(|v| v.to_str().ok())
553        .and_then(|s| s.split(',').next())
554        .map(|s| s.trim().to_string());
555    let ua = headers.get(USER_AGENT).and_then(|v| v.to_str().ok());
556
557    // 1. Validate challenge token
558    let user_id = match ath.db().validate_mfa_challenge(&form.mfa_token).await? {
559        Some(uid) => uid,
560        None => {
561            return Ok((StatusCode::SEE_OTHER, [(LOCATION, "/login".to_string())]).into_response());
562        }
563    };
564
565    // 2. Branch: recovery code vs TOTP
566    let use_recovery = form.use_recovery.is_some();
567    let verified = if use_recovery {
568        let code = form.recovery_code.as_deref().unwrap_or("");
569        ath.verify_recovery_code(user_id, code).await?
570    } else {
571        let code = form.code.as_deref().unwrap_or("");
572        ath.verify_totp(user_id, code).await?
573    };
574
575    if !verified {
576        // Log failure
577        let _ = ath
578            .db()
579            .log_audit(
580                AuditEvent::MfaChallengeFailed,
581                Some(&user_id),
582                None,
583                ip.as_deref(),
584                ua,
585                None,
586            )
587            .await;
588
589        let error_msg = if use_recovery {
590            CHALLENGE_INVALID_RECOVERY
591        } else {
592            CHALLENGE_INVALID_TOTP
593        };
594
595        let html = crate::browser_templates::render(
596            &config.templates,
597            "mfa_challenge.html",
598            context! {
599                mfa_token => &form.mfa_token,
600                error => error_msg,
601                is_production => config.is_production,
602                ..branding_context(branding.as_ref()),
603            },
604        )?;
605        return Ok(html.into_response());
606    }
607
608    // 3. Success: consume challenge, create session
609    ath.db().consume_mfa_challenge(&form.mfa_token).await?;
610
611    let _ = ath
612        .db()
613        .log_audit(
614            AuditEvent::MfaChallengeSuccess,
615            Some(&user_id),
616            None,
617            ip.as_deref(),
618            ua,
619            None,
620        )
621        .await;
622
623    // Emit Login to maintain the invariant that every session creation
624    // produces a Login audit event, consistent with the non-MFA login path.
625    let _ = ath
626        .db()
627        .log_audit(
628            AuditEvent::Login,
629            Some(&user_id),
630            None,
631            ip.as_deref(),
632            ua,
633            None,
634        )
635        .await;
636
637    let token = sessions::generate_token();
638    let token_hash = sessions::hash_token(&token);
639    let ttl = ath.session_config().ttl;
640    let expires_at = Utc::now() + ttl;
641    ath.db()
642        .create_session(user_id, token_hash, ip.as_deref(), ua, expires_at)
643        .await?;
644
645    ath.notify_user_active(user_id);
646    ath.emit_event(allowthem_core::AuthEvent::new(
647        "session.created",
648        Some(user_id),
649        serde_json::json!({ "user_id": user_id }),
650    ))
651    .await;
652
653    let cookie = ath.session_cookie(&token);
654
655    Ok((
656        StatusCode::SEE_OTHER,
657        [(SET_COOKIE, cookie), (LOCATION, "/".to_string())],
658    )
659        .into_response())
660}
661
662// ---------------------------------------------------------------------------
663// Public router constructors
664// ---------------------------------------------------------------------------
665
666/// Build a router for MFA setup routes (authenticated, CSRF-protected).
667///
668/// Mounts:
669/// - GET  /settings/mfa/setup
670/// - POST /settings/mfa/confirm
671/// - POST /settings/mfa/disable
672pub fn mfa_setup_routes(
673    templates: Arc<Environment<'static>>,
674    is_production: bool,
675    base_url: String,
676) -> Router<()> {
677    let cfg = MfaPageConfig {
678        templates,
679        is_production,
680        base_url,
681    };
682    Router::new()
683        .route("/settings/mfa/setup", get(get_mfa_setup))
684        .route("/settings/mfa/confirm", post(post_mfa_confirm))
685        .route("/settings/mfa/disable", post(post_mfa_disable))
686        .route(
687            "/settings/mfa/recovery-codes/regenerate",
688            post(post_regenerate_recovery_codes),
689        )
690        .layer(Extension(cfg))
691}
692
693/// Build a router for the MFA challenge route (mid-login, no session).
694///
695/// Mounts:
696/// - GET  /mfa/challenge
697/// - POST /mfa/challenge
698pub fn mfa_challenge_routes(
699    templates: Arc<Environment<'static>>,
700    is_production: bool,
701) -> Router<()> {
702    let cfg = MfaPageConfig {
703        templates,
704        is_production,
705        base_url: String::new(),
706    };
707    Router::new()
708        .route(
709            "/mfa/challenge",
710            get(get_mfa_challenge).post(post_mfa_challenge),
711        )
712        .layer(Extension(cfg))
713}
714
715#[cfg(test)]
716mod tests {
717    use super::*;
718
719    use axum::body::Body;
720    use axum::http::{Request, StatusCode, header};
721    use chrono::{Duration, Utc};
722    use totp_rs::{Algorithm, Secret, TOTP};
723    use tower::ServiceExt;
724
725    use allowthem_core::{AllowThemBuilder, Email, generate_token, hash_token};
726
727    const TEST_MFA_KEY: [u8; 32] = [0x42; 32];
728
729    // ---------------------------------------------------------------------------
730    // Helpers
731    // ---------------------------------------------------------------------------
732
733    async fn setup() -> AllowThem {
734        AllowThemBuilder::new("sqlite::memory:")
735            .cookie_secure(false)
736            .mfa_key(TEST_MFA_KEY)
737            .csrf_key(*b"test-csrf-key-for-binary-tests!!")
738            .build()
739            .await
740            .unwrap()
741    }
742
743    /// Build a router that exercises only the MFA routes (no login).
744    /// Setup-side routes are CSRF-protected; challenge routes are not.
745    fn test_app(ath: AllowThem) -> Router {
746        let templates = crate::browser_templates::build_default_browser_env();
747        Router::new()
748            .merge(mfa_setup_routes(
749                templates.clone(),
750                false,
751                "http://127.0.0.1:3100".into(),
752            ))
753            .layer(axum::middleware::from_fn(crate::csrf::csrf_middleware))
754            .merge(mfa_challenge_routes(templates, false))
755            .layer(axum::middleware::from_fn_with_state(
756                ath.clone(),
757                crate::cors::inject_ath_into_extensions,
758            ))
759    }
760
761    async fn create_session(ath: &AllowThem) -> (allowthem_core::types::UserId, String) {
762        let email = Email::new("mfa-test@example.com".into()).unwrap();
763        let user = ath
764            .db()
765            .create_user(email, "pass", None, None)
766            .await
767            .unwrap();
768        let token = generate_token();
769        let token_hash = hash_token(&token);
770        let expires = Utc::now() + Duration::hours(24);
771        ath.db()
772            .create_session(user.id, token_hash, None, None, expires)
773            .await
774            .unwrap();
775        let cookie = ath.session_cookie(&token);
776        let cookie_val = cookie.split(';').next().unwrap().to_string();
777        (user.id, cookie_val)
778    }
779
780    /// Acquire a CSRF token by hitting the setup GET endpoint and parsing it from HTML.
781    async fn get_csrf(app: &Router, session_cookie: &str) -> String {
782        let req = Request::builder()
783            .uri("/settings/mfa/setup")
784            .header(header::COOKIE, session_cookie)
785            .body(Body::empty())
786            .unwrap();
787        let resp = app.clone().oneshot(req).await.unwrap();
788        let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
789            .await
790            .unwrap();
791        let html = String::from_utf8(bytes.to_vec()).unwrap();
792        let marker = "name=\"csrf_token\" value=\"";
793        let start = html.find(marker).expect("csrf_token not found in HTML") + marker.len();
794        let end = html[start..].find('"').unwrap() + start;
795        html[start..end].to_string()
796    }
797
798    /// Create a user with MFA enabled. Returns (totp, recovery_codes).
799    async fn enable_mfa_for_user(
800        ath: &AllowThem,
801        user_id: allowthem_core::types::UserId,
802    ) -> (TOTP, Vec<String>) {
803        let secret_b32 = ath.create_mfa_secret(user_id).await.unwrap();
804        let totp = TOTP::new(
805            Algorithm::SHA1,
806            6,
807            1,
808            30,
809            Secret::Encoded(secret_b32).to_bytes().unwrap(),
810            None,
811            String::new(),
812        )
813        .unwrap();
814        let code = totp.generate_current().unwrap();
815        let recovery_codes = ath.enable_mfa(user_id, &code).await.unwrap();
816        (totp, recovery_codes)
817    }
818
819    // ---------------------------------------------------------------------------
820    // qr_data_uri — pure function, no I/O
821    // ---------------------------------------------------------------------------
822
823    #[test]
824    fn qr_data_uri_produces_svg_data_uri() {
825        let uri = qr_data_uri("otpauth://totp/test?secret=ABC&issuer=test");
826        assert!(
827            uri.starts_with("data:image/svg+xml,"),
828            "must produce an SVG data URI"
829        );
830        assert!(uri.contains("svg"), "must contain SVG content");
831        // The URI is injected via |safe in the template; ensure no raw & that
832        // would break HTML attribute parsing.
833        assert!(
834            !uri.contains('&'),
835            "data URI must not contain raw '&' characters"
836        );
837    }
838
839    #[test]
840    fn qr_data_uri_empty_input_still_works() {
841        let uri = qr_data_uri("");
842        // Empty string is valid QR content
843        assert!(uri.starts_with("data:image/svg+xml,"));
844    }
845
846    // ---------------------------------------------------------------------------
847    // derive_issuer — pure function, no I/O
848    // ---------------------------------------------------------------------------
849
850    #[test]
851    fn derive_issuer_strips_http_scheme() {
852        assert_eq!(derive_issuer("http://example.com"), "example.com");
853    }
854
855    #[test]
856    fn derive_issuer_strips_https_scheme() {
857        assert_eq!(
858            derive_issuer("https://auth.example.com"),
859            "auth.example.com"
860        );
861    }
862
863    #[test]
864    fn derive_issuer_strips_port() {
865        // totp-rs rejects issuer strings containing colons; port must be removed.
866        assert_eq!(derive_issuer("http://127.0.0.1:3100"), "127.0.0.1");
867    }
868
869    #[test]
870    fn derive_issuer_strips_path() {
871        assert_eq!(
872            derive_issuer("https://auth.example.com/some/path"),
873            "auth.example.com"
874        );
875    }
876
877    // ---------------------------------------------------------------------------
878    // GET /settings/mfa/setup — idempotency
879    // ---------------------------------------------------------------------------
880
881    #[tokio::test]
882    async fn get_mfa_setup_renders_secret() {
883        let ath = setup().await;
884        let app = test_app(ath.clone());
885        let (_, cookie) = create_session(&ath).await;
886
887        let csrf = get_csrf(&app, &cookie).await;
888        let req = Request::builder()
889            .uri("/settings/mfa/setup")
890            .header(header::COOKIE, format!("{cookie}; csrf_token={csrf}"))
891            .body(Body::empty())
892            .unwrap();
893        let resp = app.oneshot(req).await.unwrap();
894
895        assert_eq!(resp.status(), StatusCode::OK);
896        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
897            .await
898            .unwrap();
899        let html = String::from_utf8(body.to_vec()).unwrap();
900        assert!(
901            html.contains("totp-secret"),
902            "setup page must show secret element"
903        );
904        // The totp_uri value is HTML-escaped by MiniJinja; check the testid container exists.
905        assert!(
906            html.contains("totp-uri"),
907            "setup page must show QR URI container"
908        );
909        assert!(
910            html.contains("data:image/svg+xml,"),
911            "setup page must include a QR code data URI"
912        );
913    }
914
915    #[tokio::test]
916    async fn get_mfa_setup_is_idempotent() {
917        // Two GETs must return the same secret so wrong-code-then-retry works.
918        let ath = setup().await;
919        let app = test_app(ath.clone());
920        let (_, cookie) = create_session(&ath).await;
921        let csrf = get_csrf(&app, &cookie).await;
922
923        let secret_of = |html: String| -> String {
924            // Extract the text content of the <code data-testid="totp-secret"> element.
925            // The template renders the element with additional class attributes before >,
926            // so split on the data-testid attribute value then find the closing > to skip
927            // all attributes, then read up to </code>.
928            let after_attr = html
929                .split("data-testid=\"totp-secret\"")
930                .nth(1)
931                .expect("totp-secret element not found in HTML");
932            let after_tag_close = after_attr
933                .splitn(2, '>')
934                .nth(1)
935                .expect("closing > of totp-secret element not found");
936            after_tag_close.split('<').next().unwrap_or("").to_string()
937        };
938
939        let req1 = Request::builder()
940            .uri("/settings/mfa/setup")
941            .header(header::COOKIE, format!("{cookie}; csrf_token={csrf}"))
942            .body(Body::empty())
943            .unwrap();
944        let resp1 = app.clone().oneshot(req1).await.unwrap();
945        let html1 = String::from_utf8(
946            axum::body::to_bytes(resp1.into_body(), usize::MAX)
947                .await
948                .unwrap()
949                .to_vec(),
950        )
951        .unwrap();
952
953        let req2 = Request::builder()
954            .uri("/settings/mfa/setup")
955            .header(header::COOKIE, format!("{cookie}; csrf_token={csrf}"))
956            .body(Body::empty())
957            .unwrap();
958        let resp2 = app.clone().oneshot(req2).await.unwrap();
959        let html2 = String::from_utf8(
960            axum::body::to_bytes(resp2.into_body(), usize::MAX)
961                .await
962                .unwrap()
963                .to_vec(),
964        )
965        .unwrap();
966
967        assert_eq!(
968            secret_of(html1),
969            secret_of(html2),
970            "repeated GET /settings/mfa/setup must return the same pending secret"
971        );
972    }
973
974    // ---------------------------------------------------------------------------
975    // POST /settings/mfa/confirm
976    // ---------------------------------------------------------------------------
977
978    #[tokio::test]
979    async fn post_mfa_confirm_invalid_code_shows_error_and_does_not_enable() {
980        let ath = setup().await;
981        let app = test_app(ath.clone());
982        let (user_id, cookie) = create_session(&ath).await;
983
984        // Trigger secret creation via GET (idempotency path)
985        let csrf = get_csrf(&app, &cookie).await;
986
987        let body_str = format!("code=000000&csrf_token={csrf}");
988        let req = Request::builder()
989            .method("POST")
990            .uri("/settings/mfa/confirm")
991            .header(header::COOKIE, format!("{cookie}; csrf_token={csrf}"))
992            .header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
993            .body(Body::from(body_str))
994            .unwrap();
995        let resp = app.oneshot(req).await.unwrap();
996
997        assert_eq!(resp.status(), StatusCode::OK);
998        let html = String::from_utf8(
999            axum::body::to_bytes(resp.into_body(), usize::MAX)
1000                .await
1001                .unwrap()
1002                .to_vec(),
1003        )
1004        .unwrap();
1005        assert!(
1006            html.contains(SETUP_INVALID_CODE),
1007            "wrong code must show setup error"
1008        );
1009        assert!(
1010            !ath.has_mfa_enabled(user_id).await.unwrap(),
1011            "MFA must not be enabled after wrong code"
1012        );
1013    }
1014
1015    #[tokio::test]
1016    async fn post_mfa_confirm_valid_code_enables_mfa_and_renders_recovery_codes() {
1017        let ath = setup().await;
1018        let app = test_app(ath.clone());
1019        let (user_id, cookie) = create_session(&ath).await;
1020
1021        let csrf = get_csrf(&app, &cookie).await;
1022
1023        // Create and retrieve the pending secret
1024        let secret = ath.create_mfa_secret(user_id).await.unwrap();
1025        let totp = TOTP::new(
1026            Algorithm::SHA1,
1027            6,
1028            1,
1029            30,
1030            Secret::Encoded(secret).to_bytes().unwrap(),
1031            None,
1032            String::new(),
1033        )
1034        .unwrap();
1035        let code = totp.generate_current().unwrap();
1036
1037        let body_str = format!("code={code}&csrf_token={csrf}");
1038        let req = Request::builder()
1039            .method("POST")
1040            .uri("/settings/mfa/confirm")
1041            .header(header::COOKIE, format!("{cookie}; csrf_token={csrf}"))
1042            .header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
1043            .body(Body::from(body_str))
1044            .unwrap();
1045        let resp = app.oneshot(req).await.unwrap();
1046
1047        assert_eq!(resp.status(), StatusCode::OK);
1048        let html = String::from_utf8(
1049            axum::body::to_bytes(resp.into_body(), usize::MAX)
1050                .await
1051                .unwrap()
1052                .to_vec(),
1053        )
1054        .unwrap();
1055        assert!(
1056            html.contains("recovery-code"),
1057            "success must render recovery codes"
1058        );
1059        assert!(
1060            ath.has_mfa_enabled(user_id).await.unwrap(),
1061            "MFA must be enabled after valid confirm"
1062        );
1063    }
1064
1065    // ---------------------------------------------------------------------------
1066    // POST /settings/mfa/disable
1067    // ---------------------------------------------------------------------------
1068
1069    #[tokio::test]
1070    async fn post_mfa_disable_removes_mfa_and_redirects() {
1071        let ath = setup().await;
1072        let app = test_app(ath.clone());
1073        let (user_id, cookie) = create_session(&ath).await;
1074        enable_mfa_for_user(&ath, user_id).await;
1075
1076        // Derive CSRF token from the session token (HMAC path — no Set-Cookie on GET).
1077        let session_token_val = cookie.split('=').nth(1).unwrap().to_string();
1078        let session_token = allowthem_core::types::SessionToken::from_encoded(session_token_val);
1079        let csrf =
1080            allowthem_core::derive_csrf_token(&session_token, b"test-csrf-key-for-binary-tests!!");
1081
1082        let body_str = format!("csrf_token={csrf}");
1083        let req = Request::builder()
1084            .method("POST")
1085            .uri("/settings/mfa/disable")
1086            .header(header::COOKIE, &cookie)
1087            .header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
1088            .body(Body::from(body_str))
1089            .unwrap();
1090        let resp = app.oneshot(req).await.unwrap();
1091
1092        assert_eq!(resp.status(), StatusCode::SEE_OTHER);
1093        assert_eq!(resp.headers().get("location").unwrap(), "/settings");
1094        assert!(
1095            !ath.has_mfa_enabled(user_id).await.unwrap(),
1096            "MFA must be disabled after disable POST"
1097        );
1098    }
1099
1100    // ---------------------------------------------------------------------------
1101    // POST /settings/mfa/recovery-codes/regenerate
1102    // ---------------------------------------------------------------------------
1103
1104    #[tokio::test]
1105    async fn post_regenerate_recovery_codes_renders_new_codes() {
1106        let ath = setup().await;
1107        let app = test_app(ath.clone());
1108        let (user_id, cookie) = create_session(&ath).await;
1109        let (_, old_codes) = enable_mfa_for_user(&ath, user_id).await;
1110
1111        let session_token_val = cookie.split('=').nth(1).unwrap().to_string();
1112        let session_token = allowthem_core::types::SessionToken::from_encoded(session_token_val);
1113        let csrf =
1114            allowthem_core::derive_csrf_token(&session_token, b"test-csrf-key-for-binary-tests!!");
1115
1116        let body_str = format!("csrf_token={csrf}");
1117        let req = Request::builder()
1118            .method("POST")
1119            .uri("/settings/mfa/recovery-codes/regenerate")
1120            .header(header::COOKIE, &cookie)
1121            .header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
1122            .body(Body::from(body_str))
1123            .unwrap();
1124        let resp = app.oneshot(req).await.unwrap();
1125
1126        assert_eq!(resp.status(), StatusCode::OK);
1127        let html = String::from_utf8(
1128            axum::body::to_bytes(resp.into_body(), usize::MAX)
1129                .await
1130                .unwrap()
1131                .to_vec(),
1132        )
1133        .unwrap();
1134        assert!(
1135            html.contains("recovery-code"),
1136            "regeneration must render recovery codes"
1137        );
1138        // Verify old codes are no longer valid
1139        for old_code in &old_codes {
1140            let valid = ath.verify_recovery_code(user_id, old_code).await.unwrap();
1141            assert!(
1142                !valid,
1143                "old recovery code must be invalidated after regeneration"
1144            );
1145        }
1146    }
1147
1148    #[tokio::test]
1149    async fn post_regenerate_recovery_codes_without_mfa_redirects() {
1150        let ath = setup().await;
1151        let app = test_app(ath.clone());
1152        let (_, cookie) = create_session(&ath).await;
1153
1154        let session_token_val = cookie.split('=').nth(1).unwrap().to_string();
1155        let session_token = allowthem_core::types::SessionToken::from_encoded(session_token_val);
1156        let csrf =
1157            allowthem_core::derive_csrf_token(&session_token, b"test-csrf-key-for-binary-tests!!");
1158
1159        let body_str = format!("csrf_token={csrf}");
1160        let req = Request::builder()
1161            .method("POST")
1162            .uri("/settings/mfa/recovery-codes/regenerate")
1163            .header(header::COOKIE, &cookie)
1164            .header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
1165            .body(Body::from(body_str))
1166            .unwrap();
1167        let resp = app.oneshot(req).await.unwrap();
1168
1169        assert_eq!(
1170            resp.status(),
1171            StatusCode::SEE_OTHER,
1172            "must redirect when MFA is not enabled"
1173        );
1174        assert_eq!(resp.headers().get("location").unwrap(), "/settings");
1175    }
1176
1177    // ---------------------------------------------------------------------------
1178    // GET /mfa/challenge
1179    // ---------------------------------------------------------------------------
1180
1181    #[tokio::test]
1182    async fn get_mfa_challenge_with_invalid_token_redirects_to_login() {
1183        let ath = setup().await;
1184        let app = test_app(ath);
1185
1186        let req = Request::builder()
1187            .uri("/mfa/challenge?token=not-a-real-token")
1188            .body(Body::empty())
1189            .unwrap();
1190        let resp = app.oneshot(req).await.unwrap();
1191
1192        assert_eq!(resp.status(), StatusCode::SEE_OTHER);
1193        assert_eq!(resp.headers().get("location").unwrap(), "/login");
1194    }
1195
1196    #[tokio::test]
1197    async fn get_mfa_challenge_with_valid_token_renders_form() {
1198        let ath = setup().await;
1199        let app = test_app(ath.clone());
1200        let (user_id, _) = create_session(&ath).await;
1201        enable_mfa_for_user(&ath, user_id).await;
1202
1203        let token = ath.db().create_mfa_challenge(user_id).await.unwrap();
1204        let req = Request::builder()
1205            .uri(format!("/mfa/challenge?token={token}"))
1206            .body(Body::empty())
1207            .unwrap();
1208        let resp = app.oneshot(req).await.unwrap();
1209
1210        assert_eq!(resp.status(), StatusCode::OK);
1211        let html = String::from_utf8(
1212            axum::body::to_bytes(resp.into_body(), usize::MAX)
1213                .await
1214                .unwrap()
1215                .to_vec(),
1216        )
1217        .unwrap();
1218        assert!(
1219            html.contains("name=\"code\""),
1220            "challenge form must have code input"
1221        );
1222        assert!(
1223            html.contains("mfa_token"),
1224            "challenge form must embed mfa_token hidden field"
1225        );
1226    }
1227
1228    #[tokio::test]
1229    async fn get_mfa_challenge_hx_request_returns_fragment() {
1230        let ath = setup().await;
1231        let app = test_app(ath.clone());
1232        let (user_id, _) = create_session(&ath).await;
1233        enable_mfa_for_user(&ath, user_id).await;
1234
1235        let token = ath.db().create_mfa_challenge(user_id).await.unwrap();
1236        let req = Request::builder()
1237            .uri(format!("/mfa/challenge?token={token}"))
1238            .header("HX-Request", "true")
1239            .body(Body::empty())
1240            .unwrap();
1241        let resp = app.oneshot(req).await.unwrap();
1242
1243        assert_eq!(resp.status(), StatusCode::OK);
1244        let html = String::from_utf8(
1245            axum::body::to_bytes(resp.into_body(), usize::MAX)
1246                .await
1247                .unwrap()
1248                .to_vec(),
1249        )
1250        .unwrap();
1251        assert!(
1252            html.contains("<main class=\"wf-auth-form\">"),
1253            "HX response must be a fragment starting at <main>"
1254        );
1255        assert!(
1256            !html.contains("<html"),
1257            "HX response must not render the full shell"
1258        );
1259    }
1260
1261    #[test]
1262    fn render_mfa_setup_fragment_composes_main_and_oob_head() {
1263        let templates = crate::browser_templates::build_default_browser_env();
1264        let config = MfaPageConfig {
1265            templates,
1266            is_production: false,
1267            base_url: "http://127.0.0.1:3100".into(),
1268        };
1269        let totp =
1270            "otpauth://totp/allowthem:user@example.com?secret=JBSWY3DPEHPK3PXP&issuer=allowthem";
1271        let html = render_mfa_setup_fragment(
1272            &config,
1273            "csrf-tok",
1274            totp,
1275            &qr_data_uri(totp),
1276            "JBSWY3DPEHPK3PXP",
1277            "",
1278            None,
1279        )
1280        .unwrap()
1281        .0;
1282        assert!(
1283            html.contains("<main class=\"wf-auth-form\">"),
1284            "fragment must include the <main> root"
1285        );
1286        assert!(
1287            html.contains("<title hx-swap-oob=\"true\">"),
1288            "fragment must include the OOB <title> tag"
1289        );
1290        assert!(
1291            html.contains("id=\"wf-screen-label\""),
1292            "fragment must include the OOB #wf-screen-label span"
1293        );
1294        assert!(
1295            html.contains("ENABLE 2FA"),
1296            "fragment must include the ENABLE 2FA status hint"
1297        );
1298        assert!(
1299            html.contains("JBSWY3DPEHPK3PXP"),
1300            "fragment must include the base32 secret"
1301        );
1302    }
1303
1304    #[tokio::test]
1305    async fn get_mfa_setup_hx_request_returns_fragment() {
1306        let ath = setup().await;
1307        let app = test_app(ath.clone());
1308        let (_, cookie) = create_session(&ath).await;
1309        let csrf = get_csrf(&app, &cookie).await;
1310
1311        let req = Request::builder()
1312            .uri("/settings/mfa/setup")
1313            .header(header::COOKIE, format!("{cookie}; csrf_token={csrf}"))
1314            .header("HX-Request", "true")
1315            .body(Body::empty())
1316            .unwrap();
1317        let resp = app.oneshot(req).await.unwrap();
1318
1319        assert_eq!(resp.status(), StatusCode::OK);
1320        let html = String::from_utf8(
1321            axum::body::to_bytes(resp.into_body(), usize::MAX)
1322                .await
1323                .unwrap()
1324                .to_vec(),
1325        )
1326        .unwrap();
1327        assert!(
1328            html.contains("<main class=\"wf-auth-form\">"),
1329            "HX response must be a fragment starting at <main>"
1330        );
1331        assert!(
1332            !html.contains("<html"),
1333            "HX response must not render the full shell"
1334        );
1335    }
1336
1337    #[test]
1338    fn render_mfa_recovery_fragment_composes_main_and_oob_head() {
1339        let templates = crate::browser_templates::build_default_browser_env();
1340        let config = MfaPageConfig {
1341            templates,
1342            is_production: false,
1343            base_url: "http://127.0.0.1:3100".into(),
1344        };
1345        let codes = vec!["AAAA-BBBB".to_string(), "CCCC-DDDD".to_string()];
1346        let html = render_mfa_recovery_fragment(&config, &codes, None)
1347            .unwrap()
1348            .0;
1349        assert!(
1350            html.contains("<main class=\"wf-auth-form\">"),
1351            "fragment must include the <main> root"
1352        );
1353        assert!(
1354            html.contains("<title hx-swap-oob=\"true\">"),
1355            "fragment must include the OOB <title> tag"
1356        );
1357        assert!(
1358            html.contains("id=\"wf-screen-label\""),
1359            "fragment must include the OOB #wf-screen-label span"
1360        );
1361        assert!(
1362            html.contains("RECOVERY CODES"),
1363            "fragment must include the RECOVERY CODES status hint"
1364        );
1365        assert!(
1366            html.contains("AAAA-BBBB"),
1367            "fragment must include the rendered recovery codes"
1368        );
1369        assert!(
1370            html.contains("wf-grid"),
1371            "fragment must include the recovery code grid"
1372        );
1373    }
1374
1375    #[tokio::test]
1376    async fn post_mfa_confirm_hx_request_returns_recovery_fragment() {
1377        let ath = setup().await;
1378        let app = test_app(ath.clone());
1379        let (user_id, cookie) = create_session(&ath).await;
1380        let csrf = get_csrf(&app, &cookie).await;
1381
1382        let secret = ath.create_mfa_secret(user_id).await.unwrap();
1383        let totp = TOTP::new(
1384            Algorithm::SHA1,
1385            6,
1386            1,
1387            30,
1388            Secret::Encoded(secret).to_bytes().unwrap(),
1389            None,
1390            String::new(),
1391        )
1392        .unwrap();
1393        let code = totp.generate_current().unwrap();
1394
1395        let body_str = format!("code={code}&csrf_token={csrf}");
1396        let req = Request::builder()
1397            .method("POST")
1398            .uri("/settings/mfa/confirm")
1399            .header(header::COOKIE, format!("{cookie}; csrf_token={csrf}"))
1400            .header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
1401            .header("HX-Request", "true")
1402            .body(Body::from(body_str))
1403            .unwrap();
1404        let resp = app.oneshot(req).await.unwrap();
1405
1406        assert_eq!(resp.status(), StatusCode::OK);
1407        let html = String::from_utf8(
1408            axum::body::to_bytes(resp.into_body(), usize::MAX)
1409                .await
1410                .unwrap()
1411                .to_vec(),
1412        )
1413        .unwrap();
1414        assert!(
1415            html.contains("<main class=\"wf-auth-form\">"),
1416            "HX response must be a fragment starting at <main>"
1417        );
1418        assert!(
1419            !html.contains("<html"),
1420            "HX response must not render the full shell"
1421        );
1422        assert!(
1423            html.contains("recovery-code"),
1424            "HX response must render the recovery codes"
1425        );
1426    }
1427
1428    #[test]
1429    fn render_mfa_challenge_fragment_composes_main_and_oob_head() {
1430        let templates = crate::browser_templates::build_default_browser_env();
1431        let config = MfaPageConfig {
1432            templates,
1433            is_production: false,
1434            base_url: String::new(),
1435        };
1436        let html = render_mfa_challenge_fragment(&config, "mfa-token-abc", "", None)
1437            .unwrap()
1438            .0;
1439        assert!(
1440            html.contains("<main class=\"wf-auth-form\">"),
1441            "fragment must include the <main> root"
1442        );
1443        assert!(
1444            html.contains("<title hx-swap-oob=\"true\">"),
1445            "fragment must include the OOB <title> tag"
1446        );
1447        assert!(
1448            html.contains("id=\"wf-screen-label\""),
1449            "fragment must include the OOB #wf-screen-label span"
1450        );
1451        assert!(
1452            html.contains("TWO-FACTOR"),
1453            "fragment must include the TWO-FACTOR status hint"
1454        );
1455    }
1456
1457    // ---------------------------------------------------------------------------
1458    // POST /mfa/challenge
1459    // ---------------------------------------------------------------------------
1460
1461    #[tokio::test]
1462    async fn post_mfa_challenge_invalid_token_redirects_to_login() {
1463        let ath = setup().await;
1464        let app = test_app(ath);
1465
1466        let body_str = "mfa_token=garbage&code=123456";
1467        let req = Request::builder()
1468            .method("POST")
1469            .uri("/mfa/challenge")
1470            .header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
1471            .body(Body::from(body_str))
1472            .unwrap();
1473        let resp = app.oneshot(req).await.unwrap();
1474
1475        assert_eq!(resp.status(), StatusCode::SEE_OTHER);
1476        assert_eq!(resp.headers().get("location").unwrap(), "/login");
1477    }
1478
1479    #[tokio::test]
1480    async fn post_mfa_challenge_wrong_totp_does_not_consume_challenge() {
1481        // Retry must be possible after a wrong code.
1482        let ath = setup().await;
1483        let app = test_app(ath.clone());
1484        let (user_id, _) = create_session(&ath).await;
1485        enable_mfa_for_user(&ath, user_id).await;
1486
1487        let token = ath.db().create_mfa_challenge(user_id).await.unwrap();
1488
1489        let body_str = format!("mfa_token={token}&code=000000");
1490        let req = Request::builder()
1491            .method("POST")
1492            .uri("/mfa/challenge")
1493            .header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
1494            .body(Body::from(body_str))
1495            .unwrap();
1496        let resp = app.oneshot(req).await.unwrap();
1497
1498        assert_eq!(resp.status(), StatusCode::OK);
1499        let html = String::from_utf8(
1500            axum::body::to_bytes(resp.into_body(), usize::MAX)
1501                .await
1502                .unwrap()
1503                .to_vec(),
1504        )
1505        .unwrap();
1506        assert!(
1507            html.contains(CHALLENGE_INVALID_TOTP),
1508            "wrong code must show TOTP error"
1509        );
1510
1511        // Challenge must still be valid (not consumed) so the user can retry
1512        let still_valid = ath.db().validate_mfa_challenge(&token).await.unwrap();
1513        assert!(
1514            still_valid.is_some(),
1515            "challenge must survive a failed attempt"
1516        );
1517    }
1518
1519    #[tokio::test]
1520    async fn post_mfa_challenge_valid_totp_creates_session_and_emits_login() {
1521        let ath = setup().await;
1522        let app = test_app(ath.clone());
1523        let (user_id, _) = create_session(&ath).await;
1524        let (totp, _) = enable_mfa_for_user(&ath, user_id).await;
1525
1526        let token = ath.db().create_mfa_challenge(user_id).await.unwrap();
1527        let code = totp.generate_current().unwrap();
1528
1529        let body_str = format!("mfa_token={token}&code={code}");
1530        let req = Request::builder()
1531            .method("POST")
1532            .uri("/mfa/challenge")
1533            .header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
1534            .body(Body::from(body_str))
1535            .unwrap();
1536        let resp = app.oneshot(req).await.unwrap();
1537
1538        assert_eq!(resp.status(), StatusCode::SEE_OTHER);
1539        assert_eq!(resp.headers().get("location").unwrap(), "/");
1540        assert!(
1541            resp.headers().get(header::SET_COOKIE).is_some(),
1542            "session cookie must be set on success"
1543        );
1544
1545        // Challenge must be consumed
1546        let consumed = ath.db().validate_mfa_challenge(&token).await.unwrap();
1547        assert!(
1548            consumed.is_none(),
1549            "challenge must be consumed after success"
1550        );
1551
1552        // Both MfaChallengeSuccess and Login must be in the audit log
1553        let entries = ath.db().get_audit_log(Some(&user_id), 50, 0).await.unwrap();
1554        let event_types: Vec<&allowthem_core::AuditEvent> =
1555            entries.iter().map(|e| &e.event_type).collect();
1556        assert!(
1557            event_types.contains(&&allowthem_core::AuditEvent::MfaChallengeSuccess),
1558            "MfaChallengeSuccess must be in audit log"
1559        );
1560        assert!(
1561            event_types.contains(&&allowthem_core::AuditEvent::Login),
1562            "Login must be in audit log after MFA challenge success"
1563        );
1564    }
1565
1566    #[tokio::test]
1567    async fn post_mfa_challenge_wrong_recovery_code_shows_error() {
1568        let ath = setup().await;
1569        let app = test_app(ath.clone());
1570        let (user_id, _) = create_session(&ath).await;
1571        enable_mfa_for_user(&ath, user_id).await;
1572
1573        let token = ath.db().create_mfa_challenge(user_id).await.unwrap();
1574
1575        let body_str = format!("mfa_token={token}&recovery_code=AAAAAAAA&use_recovery=on");
1576        let req = Request::builder()
1577            .method("POST")
1578            .uri("/mfa/challenge")
1579            .header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
1580            .body(Body::from(body_str))
1581            .unwrap();
1582        let resp = app.oneshot(req).await.unwrap();
1583
1584        assert_eq!(resp.status(), StatusCode::OK);
1585        let html = String::from_utf8(
1586            axum::body::to_bytes(resp.into_body(), usize::MAX)
1587                .await
1588                .unwrap()
1589                .to_vec(),
1590        )
1591        .unwrap();
1592        assert!(
1593            html.contains(CHALLENGE_INVALID_RECOVERY),
1594            "wrong recovery code must show recovery error"
1595        );
1596    }
1597}
allowthem_server/mfa_page_routes.rs

allowthem_server/
mfa_page_routes.rs
