TOTP google authenticator implemented properly well

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))