Skip to main content

max / makenotwork

14.7 KB · 449 lines History Blame Raw
1 //! TOTP 2FA management API: setup, confirm, disable, backup codes, status.
2
3 use axum::{
4 Form,
5 extract::State,
6 response::{Html, IntoResponse, Response},
7 };
8 use serde::Deserialize;
9
10 use crate::config::Config;
11 use sqlx::PgPool;
12
13 use crate::{
14 auth::{AuthUser, verify_password_async},
15 constants::{BACKUP_CODE_COUNT, BACKUP_CODE_LENGTH, TOTP_DIGITS, TOTP_SKEW, TOTP_STEP},
16 db,
17 error::{AppError, Result, ResultExt},
18 helpers::hx_toast,
19 templates::{SaveStatusTemplate, TotpSetupTemplate, TotpStatusTemplate},
20 };
21
22 /// Generate a TOTP secret, QR code, and backup codes (does not enable 2FA yet).
23 #[tracing::instrument(skip_all, name = "totp::setup")]
24 pub(super) async fn setup(
25 State(db): State<PgPool>,
26 State(config): State<Config>,
27 AuthUser(user): AuthUser,
28 ) -> Result<Response> {
29 user.check_not_sandbox()?;
30 // Generate a 20-byte (160-bit) random secret
31 use rand::RngExt;
32 let secret_bytes: Vec<u8> = (0..20).map(|_| rand::rng().random()).collect();
33
34 let totp = totp_rs::TOTP::new(
35 totp_rs::Algorithm::SHA1,
36 TOTP_DIGITS,
37 TOTP_SKEW,
38 TOTP_STEP,
39 secret_bytes,
40 Some("Makenotwork".to_string()),
41 user.email.clone(),
42 )
43 .context("totp generation")?;
44
45 let secret_base32 = totp.get_secret_base32();
46
47 // Store the secret (not yet enabled), encrypted at rest.
48 db::totp::set_totp_secret(&db, user.id, &secret_base32, &config.signing_secret).await?;
49
50 // Generate QR code as base64 PNG
51 let qr_base64 = totp
52 .get_qr_base64()
53 .map_err(|e| AppError::Internal(anyhow::anyhow!("qr code generation: {e}")))?;
54
55 let backup_codes = generate_backup_codes();
56 // Hash the backup codes (Argon2id, one per code) on a blocking thread so the
57 // batch can't occupy a Tokio worker.
58 let code_hashes: Vec<String> = {
59 let codes = backup_codes.clone();
60 let secret = config.signing_secret.clone();
61 tokio::task::spawn_blocking(move || {
62 codes
63 .iter()
64 .map(|code| hash_backup_code(code, &secret))
65 .collect::<Vec<String>>()
66 })
67 .await
68 .map_err(|e| AppError::Internal(anyhow::anyhow!("backup-code hash task join: {e}")))?
69 };
70
71 db::totp::create_backup_codes(&db, user.id, &code_hashes).await?;
72
73 Ok(TotpSetupTemplate {
74 qr_base64,
75 secret_base32,
76 backup_codes,
77 }
78 .into_response())
79 }
80
81 /// Verify the user's first TOTP code and enable 2FA.
82 #[derive(Deserialize)]
83 pub(crate) struct ConfirmForm {
84 code: String,
85 }
86
87 #[tracing::instrument(skip_all, name = "totp::confirm")]
88 pub(super) async fn confirm(
89 State(db): State<PgPool>,
90 State(config): State<Config>,
91 AuthUser(user): AuthUser,
92 Form(form): Form<ConfirmForm>,
93 ) -> Result<Response> {
94 let secret = db::totp::get_totp_secret(&db, user.id, &config.signing_secret)
95 .await?
96 .ok_or_else(|| AppError::BadRequest("2FA setup not started".to_string()))?;
97
98 let totp = build_totp(&secret, &user.email)?;
99
100 // Validate via the matched time step (not `check_current`) and apply the same
101 // step-monotonicity gate the login path uses, so confirm and login share one
102 // replay rule: a code is accepted only if its step is newer than the last
103 // accepted step, and that step is then recorded. `setup` cleared the step to
104 // NULL (defaults to 0), so a genuine first code (step ~= now/30) always wins.
105 let now = chrono::Utc::now().timestamp() as u64;
106 let invalid = || -> Result<Response> {
107 Ok((
108 [
109 ("HX-Retarget", "#totp-confirm-status"),
110 ("HX-Reswap", "innerHTML"),
111 ],
112 Html(
113 SaveStatusTemplate {
114 success: false,
115 message: "Invalid code. Please try again.".to_string(),
116 }
117 .render_string()?,
118 ),
119 )
120 .into_response())
121 };
122
123 let Some(step) = find_matching_step(&totp, &form.code, now) else {
124 return invalid();
125 };
126
127 // Record the matched step atomically; the guarded write is the authoritative
128 // replay gate. A concurrent submission of the same code that already advanced
129 // the step loses here and is rejected, the step monotonicity is enforced by
130 // the DB, not by a separate read-then-write.
131 if !db::totp::set_totp_last_used_step(&db, user.id, step).await? {
132 return invalid();
133 }
134
135 db::totp::enable_totp(&db, user.id).await?;
136
137 Ok((
138 [(
139 "HX-Trigger",
140 hx_toast("Two-factor authentication enabled", "success"),
141 )],
142 TotpStatusTemplate { enabled: true },
143 )
144 .into_response())
145 }
146
147 /// Disable 2FA (requires password confirmation).
148 #[derive(Deserialize)]
149 pub(crate) struct DisableForm {
150 password: String,
151 }
152
153 #[tracing::instrument(skip_all, name = "totp::disable")]
154 pub(super) async fn disable(
155 State(db): State<PgPool>,
156 AuthUser(user): AuthUser,
157 Form(form): Form<DisableForm>,
158 ) -> Result<Response> {
159 // Verify password
160 let db_user = db::users::get_user_by_id(&db, user.id)
161 .await?
162 .ok_or(AppError::Unauthorized)?;
163
164 if !verify_password_async(form.password.clone(), db_user.password_hash.clone()).await? {
165 return Ok((
166 [
167 ("HX-Retarget", "#totp-disable-status"),
168 ("HX-Reswap", "innerHTML"),
169 ],
170 Html(
171 SaveStatusTemplate {
172 success: false,
173 message: "Incorrect password.".to_string(),
174 }
175 .render_string()?,
176 ),
177 )
178 .into_response());
179 }
180
181 db::totp::disable_totp(&db, user.id).await?;
182
183 Ok((
184 [(
185 "HX-Trigger",
186 hx_toast("Two-factor authentication disabled", "success"),
187 )],
188 TotpStatusTemplate { enabled: false },
189 )
190 .into_response())
191 }
192
193 /// Regenerate backup codes (requires password confirmation).
194 #[derive(Deserialize)]
195 pub(crate) struct RegenerateForm {
196 password: String,
197 }
198
199 #[tracing::instrument(skip_all, name = "totp::regenerate_backup_codes")]
200 pub(super) async fn regenerate_backup_codes(
201 State(db): State<PgPool>,
202 State(config): State<Config>,
203 AuthUser(user): AuthUser,
204 Form(form): Form<RegenerateForm>,
205 ) -> Result<Response> {
206 // Verify password
207 let db_user = db::users::get_user_by_id(&db, user.id)
208 .await?
209 .ok_or(AppError::Unauthorized)?;
210
211 if !verify_password_async(form.password.clone(), db_user.password_hash.clone()).await? {
212 return Ok((
213 [
214 ("HX-Retarget", "#backup-regen-status"),
215 ("HX-Reswap", "innerHTML"),
216 ],
217 Html(
218 SaveStatusTemplate {
219 success: false,
220 message: "Incorrect password.".to_string(),
221 }
222 .render_string()?,
223 ),
224 )
225 .into_response());
226 }
227
228 let backup_codes = generate_backup_codes();
229 // Hash the backup codes (Argon2id, one per code) on a blocking thread so the
230 // batch can't occupy a Tokio worker.
231 let code_hashes: Vec<String> = {
232 let codes = backup_codes.clone();
233 let secret = config.signing_secret.clone();
234 tokio::task::spawn_blocking(move || {
235 codes
236 .iter()
237 .map(|code| hash_backup_code(code, &secret))
238 .collect::<Vec<String>>()
239 })
240 .await
241 .map_err(|e| AppError::Internal(anyhow::anyhow!("backup-code hash task join: {e}")))?
242 };
243
244 db::totp::create_backup_codes(&db, user.id, &code_hashes).await?;
245
246 // Return the new codes as an HTML partial
247 let codes_html: String = backup_codes
248 .iter()
249 .map(|c| format!("<code>{c}</code>"))
250 .collect::<Vec<_>>()
251 .join("\n");
252
253 Ok((
254 [("HX-Trigger", hx_toast("Backup codes regenerated", "success"))],
255 Html(format!(
256 "<div class=\"backup-codes-grid\">\n{codes_html}\n</div>\n<p style=\"opacity: 0.7; font-size: 0.85rem; margin-top: 0.75rem;\">Save these codes somewhere safe. Each code can only be used once.</p>"
257 )),
258 )
259 .into_response())
260 }
261
262 /// Return the current 2FA status as an HTMX partial for the dashboard.
263 #[tracing::instrument(skip_all, name = "totp::status")]
264 pub(super) async fn status(State(db): State<PgPool>, AuthUser(user): AuthUser) -> Result<Response> {
265 let enabled = db::totp::is_totp_enabled(&db, user.id).await?;
266
267 Ok(TotpStatusTemplate { enabled }.into_response())
268 }
269
270 // ── Helpers ─────────────────────────────────────────────────────────────────
271
272 /// Build a TOTP instance from a stored base32 secret.
273 pub(crate) fn build_totp(secret_base32: &str, account_name: &str) -> Result<totp_rs::TOTP> {
274 let secret_bytes = totp_rs::Secret::Encoded(secret_base32.to_string())
275 .to_bytes()
276 .context("parse totp secret")?;
277
278 totp_rs::TOTP::new(
279 totp_rs::Algorithm::SHA1,
280 TOTP_DIGITS,
281 TOTP_SKEW,
282 TOTP_STEP,
283 secret_bytes,
284 Some("Makenotwork".to_string()),
285 account_name.to_string(),
286 )
287 .context("totp creation")
288 }
289
290 /// Generate random alphanumeric backup codes.
291 fn generate_backup_codes() -> Vec<String> {
292 use rand::RngExt;
293 let mut rng = rand::rng();
294
295 (0..BACKUP_CODE_COUNT)
296 .map(|_| {
297 (0..BACKUP_CODE_LENGTH)
298 .map(|_| {
299 let idx: u8 = rng.random_range(0..36);
300 if idx < 10 {
301 (b'0' + idx) as char
302 } else {
303 (b'a' + idx - 10) as char
304 }
305 })
306 .collect()
307 })
308 .collect()
309 }
310
311 /// Find which TOTP time step a code matches, returning the step number.
312 ///
313 /// This is used instead of `totp.check_current()` so we can store the
314 /// *matched* step for replay prevention, not just the wall-clock step.
315 /// Without this, a code valid for step N can be replayed at step N+1
316 /// within the skew window.
317 pub(crate) fn find_matching_step(totp: &totp_rs::TOTP, code: &str, time_secs: u64) -> Option<i64> {
318 let base_step = time_secs / TOTP_STEP;
319 let skew = TOTP_SKEW as u64;
320 let start = base_step.saturating_sub(skew);
321 for i in 0..=(skew * 2) {
322 let step = start + i;
323 let step_time = step * TOTP_STEP;
324 let expected = totp.generate(step_time);
325 if crate::crypto::constant_time_compare(&expected, code) {
326 return Some(step as i64);
327 }
328 }
329 None
330 }
331
332 /// Argon2id hash of a backup code.
333 ///
334 /// Backup codes have only ~41 bits of entropy (8 alphanumeric chars), so the
335 /// previous HMAC-SHA256 scheme was brute-forceable in minutes if the DB and
336 /// the server's signing secret both leaked. Argon2id with even modest
337 /// parameters multiplies that work by ~10^5, putting offline attack in the
338 /// "needs a real GPU farm and time" range.
339 ///
340 /// Uses lower-than-password params (8 MiB, 1 iteration), backup codes are
341 /// random tokens, not user-chosen passwords, so the security floor is set by
342 /// the wordlist, not the hash function. Per-code unique salt.
343 pub(crate) fn hash_backup_code(code: &str, _secret: &str) -> String {
344 use argon2::password_hash::{PasswordHasher, SaltString, rand_core::OsRng};
345 use argon2::{Algorithm, Argon2, Params, Version};
346
347 let salt = SaltString::generate(&mut OsRng);
348 let params = Params::new(8 * 1024, 1, 1, None).expect("argon2 backup-code params are valid");
349 let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
350 let hash = argon2
351 .hash_password(code.as_bytes(), &salt)
352 .expect("argon2 backup-code hashing");
353 hash.to_string()
354 }
355
356 /// Legacy HMAC-SHA256 hash kept for verifying pre-migration backup codes.
357 ///
358 /// Used only as a fallback inside `verify_and_consume_backup_code` when the
359 /// stored hash isn't an Argon2 PHC string. Do not call from new code paths,
360 /// any newly issued backup code goes through `hash_backup_code` (Argon2).
361 pub(crate) fn legacy_hmac_backup_code(code: &str, secret: &str) -> String {
362 use hmac::{Hmac, KeyInit, Mac};
363 use sha2::Sha256;
364
365 let mut mac =
366 Hmac::<Sha256>::new_from_slice(secret.as_bytes()).expect("HMAC accepts any key length");
367 mac.update(code.as_bytes());
368 hex::encode(mac.finalize().into_bytes())
369 }
370
371 #[cfg(test)]
372 mod tests {
373 use super::*;
374
375 #[test]
376 fn backup_code_generation_produces_correct_count() {
377 let codes = generate_backup_codes();
378 assert_eq!(codes.len(), BACKUP_CODE_COUNT);
379 }
380
381 #[test]
382 fn backup_codes_are_correct_length() {
383 let codes = generate_backup_codes();
384 for code in &codes {
385 assert_eq!(code.len(), BACKUP_CODE_LENGTH);
386 }
387 }
388
389 #[test]
390 fn backup_codes_are_alphanumeric() {
391 let codes = generate_backup_codes();
392 for code in &codes {
393 assert!(
394 code.chars().all(|c| c.is_ascii_alphanumeric()),
395 "Code should be alphanumeric: {code}"
396 );
397 }
398 }
399
400 #[test]
401 fn backup_codes_are_unique() {
402 let codes = generate_backup_codes();
403 let unique: std::collections::HashSet<&String> = codes.iter().collect();
404 assert_eq!(unique.len(), codes.len());
405 }
406
407 #[test]
408 fn hash_backup_code_is_argon2_phc() {
409 // Argon2 hashes are non-deterministic (unique salt per call) and use
410 // the PHC string format that starts with `$argon2`.
411 let h = hash_backup_code("abc12345", "secret");
412 assert!(h.starts_with("$argon2"), "got {h}");
413 }
414
415 #[test]
416 fn hash_backup_code_non_deterministic() {
417 // Two hashes of the same code must differ, distinct salts.
418 let h1 = hash_backup_code("abc12345", "secret");
419 let h2 = hash_backup_code("abc12345", "secret");
420 assert_ne!(h1, h2);
421 }
422
423 #[test]
424 fn hash_backup_code_verifies_against_itself() {
425 use argon2::{Argon2, PasswordHash, password_hash::PasswordVerifier};
426 let h = hash_backup_code("abc12345", "ignored");
427 let parsed = PasswordHash::new(&h).unwrap();
428 assert!(
429 Argon2::default()
430 .verify_password(b"abc12345", &parsed)
431 .is_ok()
432 );
433 assert!(
434 Argon2::default()
435 .verify_password(b"wrong", &parsed)
436 .is_err()
437 );
438 }
439
440 #[test]
441 fn legacy_hmac_is_deterministic_and_secret_keyed() {
442 let h1 = legacy_hmac_backup_code("abc12345", "secret");
443 let h2 = legacy_hmac_backup_code("abc12345", "secret");
444 assert_eq!(h1, h2);
445 let h3 = legacy_hmac_backup_code("abc12345", "different-secret");
446 assert_ne!(h1, h3);
447 }
448 }
449