TOTP google authenticator implemented properly well

src/controllers/account.rs

@@ -363,6 +363,177 @@ async fn change_password(
     password_view(&v, &jar, &user, true, None)
 }
 
+// ---- Two-factor authentication (TOTP / Google Authenticator) -------------
+//
+// Entirely opt-in. The security page has three shapes, all rendered from
+// `security.html`:
+//   * disabled  -> an "enable" button,
+//   * enrolling -> the QR + a confirm-code field (secret staged, not yet on),
+//   * enabled   -> status, remaining backup codes, disable/regenerate forms.
+// Both turning 2FA off and regenerating backup codes require re-entering the
+// account password, so a walk-up attacker on an open session can't weaken it.
+
+#[derive(Debug, Deserialize)]
+struct ConfirmTotpForm {
+    code: String,
+}
+
+#[derive(Debug, Deserialize)]
+struct PasswordConfirmForm {
+    current_password: String,
+}
+
+/// Render the security page. Exactly one of (`enrolling`, plain status) applies;
+/// `backup_codes` is non-empty only on the one render right after enabling or
+/// regenerating, where the plaintext codes are shown once.
+#[allow(clippy::too_many_arguments)]
+fn security_view(
+    v: &TeraView,
+    jar: &CookieJar,
+    user: &users::Model,
+    enrolling: bool,
+    qr: Option<&str>,
+    secret: Option<&str>,
+    backup_codes: &[String],
+    error: Option<&str>,
+) -> Result<Response> {
+    format::view(
+        v,
+        "account/security.html",
+        json!({
+            "logged_in_admin": false,
+            "logged_in_customer": true,
+            "account_nav": true,
+            "customer_name": user.name,
+            "customer_account_type": user.account_type,
+            "totp_enabled": user.totp_enabled(),
+            "enrolling": enrolling,
+            "qr": qr,
+            "secret": secret,
+            "backup_codes": backup_codes,
+            "backup_remaining": user.backup_codes_remaining(),
+            "error": error,
+            "lang": current_lang(jar),
+        }),
+    )
+}
+
+/// Common guard for every security handler: a signed-in, non-admin customer.
+async fn require_customer(ctx: &AppContext, jar: &CookieJar) -> Result<users::Model> {
+    match guard::current_user(ctx, jar).await {
+        Some(user) if guard::is_admin(ctx, &user) => Err(Error::string("admin")),
+        Some(user) => Ok(user),
+        None => Err(Error::Unauthorized("login required".into())),
+    }
+}
+
+#[debug_handler]
+async fn security_page(
+    jar: CookieJar,
+    ViewEngine(v): ViewEngine<TeraView>,
+    State(ctx): State<AppContext>,
+) -> Result<Response> {
+    let Some(user) = guard::current_user(&ctx, &jar).await else {
+        return format::redirect("/login");
+    };
+    if guard::is_admin(&ctx, &user) {
+        return format::redirect("/admin/dashboard");
+    }
+    security_view(&v, &jar, &user, false, None, None, &[], None)
+}
+
+/// Stage a fresh secret and show the QR + confirm-code field.
+#[debug_handler]
+async fn enable_totp(
+    jar: CookieJar,
+    ViewEngine(v): ViewEngine<TeraView>,
+    State(ctx): State<AppContext>,
+) -> Result<Response> {
+    let Ok(user) = require_customer(&ctx, &jar).await else {
+        return format::redirect("/login");
+    };
+    // Already on — nothing to enroll.
+    if user.totp_enabled() {
+        return security_view(&v, &jar, &user, false, None, None, &[], None);
+    }
+    let user = user.into_active_model().begin_totp_enrollment(&ctx.db).await?;
+    let Some((qr, secret)) = user.totp_provisioning() else {
+        return security_view(&v, &jar, &user, false, None, None, &[], Some("enroll"));
+    };
+    security_view(&v, &jar, &user, true, Some(&qr), Some(&secret), &[], None)
+}
+
+/// Verify the first code against the staged secret; on success flip 2FA on and
+/// show the one-time backup codes. On a wrong code, re-show the QR to retry.
+#[debug_handler]
+async fn confirm_totp(
+    jar: CookieJar,
+    ViewEngine(v): ViewEngine<TeraView>,
+    State(ctx): State<AppContext>,
+    Form(form): Form<ConfirmTotpForm>,
+) -> Result<Response> {
+    let Ok(user) = require_customer(&ctx, &jar).await else {
+        return format::redirect("/login");
+    };
+    if user.totp_enabled() {
+        return security_view(&v, &jar, &user, false, None, None, &[], None);
+    }
+    if !user.verify_totp_code(&form.code) {
+        let qr = user.totp_provisioning();
+        let (qr, secret) = match &qr {
+            Some((q, s)) => (Some(q.as_str()), Some(s.as_str())),
+            None => (None, None),
+        };
+        return security_view(&v, &jar, &user, true, qr, secret, &[], Some("code"));
+    }
+    let (user, backup_codes) = user.into_active_model().enable_totp(&ctx.db).await?;
+    security_view(&v, &jar, &user, false, None, None, &backup_codes, None)
+}
+
+/// Turn 2FA off — requires the account password as confirmation.
+#[debug_handler]
+async fn disable_totp(
+    jar: CookieJar,
+    ViewEngine(v): ViewEngine<TeraView>,
+    State(ctx): State<AppContext>,
+    Form(form): Form<PasswordConfirmForm>,
+) -> Result<Response> {
+    let Ok(user) = require_customer(&ctx, &jar).await else {
+        return format::redirect("/login");
+    };
+    if !user.totp_enabled() {
+        return security_view(&v, &jar, &user, false, None, None, &[], None);
+    }
+    if !user.verify_password(&form.current_password) {
+        return security_view(&v, &jar, &user, false, None, None, &[], Some("password"));
+    }
+    let user = user.into_active_model().disable_totp(&ctx.db).await?;
+    security_view(&v, &jar, &user, false, None, None, &[], None)
+}
+
+/// Issue a fresh set of backup codes (invalidating the old ones) — also gated by
+/// the account password.
+#[debug_handler]
+async fn regenerate_backup_codes(
+    jar: CookieJar,
+    ViewEngine(v): ViewEngine<TeraView>,
+    State(ctx): State<AppContext>,
+    Form(form): Form<PasswordConfirmForm>,
+) -> Result<Response> {
+    let Ok(user) = require_customer(&ctx, &jar).await else {
+        return format::redirect("/login");
+    };
+    if !user.totp_enabled() {
+        return security_view(&v, &jar, &user, false, None, None, &[], None);
+    }
+    if !user.verify_password(&form.current_password) {
+        return security_view(&v, &jar, &user, false, None, None, &[], Some("password"));
+    }
+    let (user, backup_codes) =
+        user.into_active_model().regenerate_backup_codes(&ctx.db).await?;
+    security_view(&v, &jar, &user, false, None, None, &backup_codes, None)
+}
+
 pub fn routes() -> Routes {
     Routes::new()
         .add("/account/profile", get(profile_page))
@@ -371,4 +542,9 @@ pub fn routes() -> Routes {
         .add("/account/orders/{order_number}", get(order_detail_page))
         .add("/account/password", get(change_password_page))
         .add("/account/password", post(change_password))
+        .add("/account/security", get(security_page))
+        .add("/account/security/enable", post(enable_totp))
+        .add("/account/security/confirm", post(confirm_totp))
+        .add("/account/security/disable", post(disable_totp))
+        .add("/account/security/backup-codes", post(regenerate_backup_codes))
 }

src/controllers/auth.rs

@@ -13,6 +13,13 @@ use time::Duration as TimeDuration;
 
 pub static EMAIL_DOMAIN_RE: OnceLock<Regex> = OnceLock::new();
 pub(crate) const AUTH_COOKIE: &str = "auth_token";
+/// Short-lived cookie that carries a half-authenticated session between the
+/// password step and the TOTP step. It is a *separate* name from `auth_token`
+/// on purpose: the auth guards only read `auth_token`, so this cookie can never
+/// authenticate a request on its own — it only proves the password step passed.
+pub(crate) const TOTP_PENDING_COOKIE: &str = "totp_pending";
+/// How long the user has to enter their 2FA code after the password step.
+pub(crate) const TOTP_PENDING_TTL_SECS: u64 = 300;
 
 fn get_allow_email_domain_re() -> &'static Regex {
     EMAIL_DOMAIN_RE.get_or_init(|| {
@@ -38,6 +45,24 @@ pub(crate) fn clear_auth_cookie() -> Cookie<'static> {
         .build()
 }
 
+pub(crate) fn totp_pending_cookie(token: &str, max_age_seconds: u64) -> Cookie<'static> {
+    Cookie::build((TOTP_PENDING_COOKIE, token.to_string()))
+        .path("/")
+        .http_only(true)
+        .same_site(SameSite::Lax)
+        .max_age(TimeDuration::seconds(max_age_seconds as i64))
+        .build()
+}
+
+pub(crate) fn clear_totp_pending_cookie() -> Cookie<'static> {
+    Cookie::build((TOTP_PENDING_COOKIE, ""))
+        .path("/")
+        .http_only(true)
+        .same_site(SameSite::Lax)
+        .max_age(TimeDuration::seconds(0))
+        .build()
+}
+
 #[derive(Debug, Deserialize, Serialize)]
 pub struct ForgotParams {
     pub email: String,

src/controllers/auth_pages.rs

@@ -85,6 +85,23 @@ async fn login(
     }
 
     let jwt_secret = ctx.config.get_jwt_config()?;
+
+    // If the user opted into 2FA, the password is only the first factor: don't
+    // issue the real auth cookie yet. Hand out a short-lived, separate "pending"
+    // cookie and send them to the code-entry page. Everyone without 2FA logs in
+    // in a single step exactly as before.
+    if user.totp_enabled() {
+        let pending = user
+            .generate_jwt(&jwt_secret.secret, auth_controller::TOTP_PENDING_TTL_SECS)
+            .or_else(|_| unauthorized("unauthorized!"))?;
+        return format::render()
+            .cookies(&[auth_controller::totp_pending_cookie(
+                &pending,
+                auth_controller::TOTP_PENDING_TTL_SECS,
+            )])?
+            .redirect("/login/totp");
+    }
+
     let token = user
         .generate_jwt(&jwt_secret.secret, jwt_secret.expiration)
         .or_else(|_| unauthorized("unauthorized!"))?;
@@ -94,6 +111,89 @@ async fn login(
         .redirect(home_for(&ctx, &user))
 }
 
+/// Resolve the user behind a valid, unexpired `totp_pending` cookie. Returns
+/// `None` (never errors) when the cookie is missing, malformed, or expired —
+/// the caller bounces such requests back to `/login`.
+async fn user_from_pending(ctx: &AppContext, jar: &CookieJar) -> Option<users::Model> {
+    let cookie = jar.get(auth_controller::TOTP_PENDING_COOKIE)?;
+    let jwt_config = ctx.config.get_jwt_config().ok()?;
+    let claims = loco_rs::auth::jwt::JWT::new(&jwt_config.secret)
+        .validate(cookie.value())
+        .ok()?;
+    let user = users::Model::find_by_pid(&ctx.db, &claims.claims.pid).await.ok()?;
+    // Defend against a stale pending cookie outliving a 2FA disable.
+    user.totp_enabled().then_some(user)
+}
+
+fn login_totp_view(v: &TeraView, jar: &CookieJar, error: Option<&str>) -> Result<Response> {
+    format::view(
+        v,
+        "auth/login_totp.html",
+        json!({
+            "error": error,
+            "logged_in_admin": false,
+            "lang": current_lang(jar),
+        }),
+    )
+}
+
+#[debug_handler]
+async fn login_totp_page(
+    jar: CookieJar,
+    ViewEngine(v): ViewEngine<TeraView>,
+    State(ctx): State<AppContext>,
+) -> Result<Response> {
+    if user_from_pending(&ctx, &jar).await.is_none() {
+        return format::redirect("/login");
+    }
+    login_totp_view(&v, &jar, None)
+}
+
+/// Second login factor. Accepts either a 6-digit authenticator code or one of
+/// the one-time backup codes (auto-detected by length). On success the pending
+/// cookie is cleared and the real `auth_token` is issued.
+#[derive(Debug, serde::Deserialize)]
+struct TotpLoginForm {
+    code: String,
+}
+
+#[debug_handler]
+async fn login_totp(
+    jar: CookieJar,
+    ViewEngine(v): ViewEngine<TeraView>,
+    State(ctx): State<AppContext>,
+    Form(form): Form<TotpLoginForm>,
+) -> Result<Response> {
+    let Some(user) = user_from_pending(&ctx, &jar).await else {
+        return format::redirect("/login");
+    };
+
+    let code = form.code.trim();
+    let via_totp = user.verify_totp_code(code);
+    let via_backup = !via_totp && user.matches_backup_code(code);
+
+    if !via_totp && !via_backup {
+        return login_totp_view(&v, &jar, Some("invalid"));
+    }
+
+    // A used backup code must be burned so it can't be replayed.
+    if via_backup {
+        user.clone().into_active_model().consume_backup_code(&ctx.db, code).await?;
+    }
+
+    let jwt_secret = ctx.config.get_jwt_config()?;
+    let token = user
+        .generate_jwt(&jwt_secret.secret, jwt_secret.expiration)
+        .or_else(|_| unauthorized("unauthorized!"))?;
+
+    format::render()
+        .cookies(&[
+            auth_controller::auth_cookie(&token, jwt_secret.expiration),
+            auth_controller::clear_totp_pending_cookie(),
+        ])?
+        .redirect(home_for(&ctx, &user))
+}
+
 #[debug_handler]
 async fn register_page(
     jar: CookieJar,
@@ -366,6 +466,8 @@ pub fn routes() -> Routes {
     Routes::new()
         .add("/login", get(login_page))
         .add("/login", post(login))
+        .add("/login/totp", get(login_totp_page))
+        .add("/login/totp", post(login_totp))
         .add("/register", get(register_page))
         .add("/register", post(register))
         .add("/verify/{token}", get(verify))