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

allowthem_server/
mfa_page_routes.rs
