Skip to main content

max / makenotwork

14.2 KB · 431 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::{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 = || {
107 (
108 [
109 ("HX-Retarget", "#totp-confirm-status"),
110 ("HX-Reswap", "innerHTML"),
111 ],
112 Html("<span class=\"save-error\">Invalid code. Please try again.</span>"),
113 )
114 .into_response()
115 };
116
117 let Some(step) = find_matching_step(&totp, &form.code, now) else {
118 return Ok(invalid());
119 };
120
121 // Record the matched step atomically; the guarded write is the authoritative
122 // replay gate. A concurrent submission of the same code that already advanced
123 // the step loses here and is rejected, the step monotonicity is enforced by
124 // the DB, not by a separate read-then-write.
125 if !db::totp::set_totp_last_used_step(&db, user.id, step).await? {
126 return Ok(invalid());
127 }
128
129 db::totp::enable_totp(&db, user.id).await?;
130
131 Ok((
132 [(
133 "HX-Trigger",
134 hx_toast("Two-factor authentication enabled", "success"),
135 )],
136 TotpStatusTemplate { enabled: true },
137 )
138 .into_response())
139 }
140
141 /// Disable 2FA (requires password confirmation).
142 #[derive(Deserialize)]
143 pub(crate) struct DisableForm {
144 password: String,
145 }
146
147 #[tracing::instrument(skip_all, name = "totp::disable")]
148 pub(super) async fn disable(
149 State(db): State<PgPool>,
150 AuthUser(user): AuthUser,
151 Form(form): Form<DisableForm>,
152 ) -> Result<Response> {
153 // Verify password
154 let db_user = db::users::get_user_by_id(&db, user.id)
155 .await?
156 .ok_or(AppError::Unauthorized)?;
157
158 if !verify_password_async(form.password.clone(), db_user.password_hash.clone()).await? {
159 return Ok((
160 [
161 ("HX-Retarget", "#totp-disable-status"),
162 ("HX-Reswap", "innerHTML"),
163 ],
164 Html("<span class=\"save-error\">Incorrect password.</span>"),
165 )
166 .into_response());
167 }
168
169 db::totp::disable_totp(&db, user.id).await?;
170
171 Ok((
172 [(
173 "HX-Trigger",
174 hx_toast("Two-factor authentication disabled", "success"),
175 )],
176 TotpStatusTemplate { enabled: false },
177 )
178 .into_response())
179 }
180
181 /// Regenerate backup codes (requires password confirmation).
182 #[derive(Deserialize)]
183 pub(crate) struct RegenerateForm {
184 password: String,
185 }
186
187 #[tracing::instrument(skip_all, name = "totp::regenerate_backup_codes")]
188 pub(super) async fn regenerate_backup_codes(
189 State(db): State<PgPool>,
190 State(config): State<Config>,
191 AuthUser(user): AuthUser,
192 Form(form): Form<RegenerateForm>,
193 ) -> Result<Response> {
194 // Verify password
195 let db_user = db::users::get_user_by_id(&db, user.id)
196 .await?
197 .ok_or(AppError::Unauthorized)?;
198
199 if !verify_password_async(form.password.clone(), db_user.password_hash.clone()).await? {
200 return Ok((
201 [
202 ("HX-Retarget", "#backup-regen-status"),
203 ("HX-Reswap", "innerHTML"),
204 ],
205 Html("<span class=\"save-error\">Incorrect password.</span>"),
206 )
207 .into_response());
208 }
209
210 let backup_codes = generate_backup_codes();
211 // Hash the backup codes (Argon2id, one per code) on a blocking thread so the
212 // batch can't occupy a Tokio worker.
213 let code_hashes: Vec<String> = {
214 let codes = backup_codes.clone();
215 let secret = config.signing_secret.clone();
216 tokio::task::spawn_blocking(move || {
217 codes
218 .iter()
219 .map(|code| hash_backup_code(code, &secret))
220 .collect::<Vec<String>>()
221 })
222 .await
223 .map_err(|e| AppError::Internal(anyhow::anyhow!("backup-code hash task join: {e}")))?
224 };
225
226 db::totp::create_backup_codes(&db, user.id, &code_hashes).await?;
227
228 // Return the new codes as an HTML partial
229 let codes_html: String = backup_codes
230 .iter()
231 .map(|c| format!("<code>{c}</code>"))
232 .collect::<Vec<_>>()
233 .join("\n");
234
235 Ok((
236 [("HX-Trigger", hx_toast("Backup codes regenerated", "success"))],
237 Html(format!(
238 "<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>"
239 )),
240 )
241 .into_response())
242 }
243
244 /// Return the current 2FA status as an HTMX partial for the dashboard.
245 #[tracing::instrument(skip_all, name = "totp::status")]
246 pub(super) async fn status(State(db): State<PgPool>, AuthUser(user): AuthUser) -> Result<Response> {
247 let enabled = db::totp::is_totp_enabled(&db, user.id).await?;
248
249 Ok(TotpStatusTemplate { enabled }.into_response())
250 }
251
252 // ── Helpers ─────────────────────────────────────────────────────────────────
253
254 /// Build a TOTP instance from a stored base32 secret.
255 pub(crate) fn build_totp(secret_base32: &str, account_name: &str) -> Result<totp_rs::TOTP> {
256 let secret_bytes = totp_rs::Secret::Encoded(secret_base32.to_string())
257 .to_bytes()
258 .context("parse totp secret")?;
259
260 totp_rs::TOTP::new(
261 totp_rs::Algorithm::SHA1,
262 TOTP_DIGITS,
263 TOTP_SKEW,
264 TOTP_STEP,
265 secret_bytes,
266 Some("Makenotwork".to_string()),
267 account_name.to_string(),
268 )
269 .context("totp creation")
270 }
271
272 /// Generate random alphanumeric backup codes.
273 fn generate_backup_codes() -> Vec<String> {
274 use rand::RngExt;
275 let mut rng = rand::rng();
276
277 (0..BACKUP_CODE_COUNT)
278 .map(|_| {
279 (0..BACKUP_CODE_LENGTH)
280 .map(|_| {
281 let idx: u8 = rng.random_range(0..36);
282 if idx < 10 {
283 (b'0' + idx) as char
284 } else {
285 (b'a' + idx - 10) as char
286 }
287 })
288 .collect()
289 })
290 .collect()
291 }
292
293 /// Find which TOTP time step a code matches, returning the step number.
294 ///
295 /// This is used instead of `totp.check_current()` so we can store the
296 /// *matched* step for replay prevention, not just the wall-clock step.
297 /// Without this, a code valid for step N can be replayed at step N+1
298 /// within the skew window.
299 pub(crate) fn find_matching_step(totp: &totp_rs::TOTP, code: &str, time_secs: u64) -> Option<i64> {
300 let base_step = time_secs / TOTP_STEP;
301 let skew = TOTP_SKEW as u64;
302 let start = base_step.saturating_sub(skew);
303 for i in 0..=(skew * 2) {
304 let step = start + i;
305 let step_time = step * TOTP_STEP;
306 let expected = totp.generate(step_time);
307 if crate::crypto::constant_time_compare(&expected, code) {
308 return Some(step as i64);
309 }
310 }
311 None
312 }
313
314 /// Argon2id hash of a backup code.
315 ///
316 /// Backup codes have only ~41 bits of entropy (8 alphanumeric chars), so the
317 /// previous HMAC-SHA256 scheme was brute-forceable in minutes if the DB and
318 /// the server's signing secret both leaked. Argon2id with even modest
319 /// parameters multiplies that work by ~10^5, putting offline attack in the
320 /// "needs a real GPU farm and time" range.
321 ///
322 /// Uses lower-than-password params (8 MiB, 1 iteration), backup codes are
323 /// random tokens, not user-chosen passwords, so the security floor is set by
324 /// the wordlist, not the hash function. Per-code unique salt.
325 pub(crate) fn hash_backup_code(code: &str, _secret: &str) -> String {
326 use argon2::password_hash::{PasswordHasher, SaltString, rand_core::OsRng};
327 use argon2::{Algorithm, Argon2, Params, Version};
328
329 let salt = SaltString::generate(&mut OsRng);
330 let params = Params::new(8 * 1024, 1, 1, None).expect("argon2 backup-code params are valid");
331 let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
332 let hash = argon2
333 .hash_password(code.as_bytes(), &salt)
334 .expect("argon2 backup-code hashing");
335 hash.to_string()
336 }
337
338 /// Legacy HMAC-SHA256 hash kept for verifying pre-migration backup codes.
339 ///
340 /// Used only as a fallback inside `verify_and_consume_backup_code` when the
341 /// stored hash isn't an Argon2 PHC string. Do not call from new code paths,
342 /// any newly issued backup code goes through `hash_backup_code` (Argon2).
343 pub(crate) fn legacy_hmac_backup_code(code: &str, secret: &str) -> String {
344 use hmac::{Hmac, KeyInit, Mac};
345 use sha2::Sha256;
346
347 let mut mac =
348 Hmac::<Sha256>::new_from_slice(secret.as_bytes()).expect("HMAC accepts any key length");
349 mac.update(code.as_bytes());
350 hex::encode(mac.finalize().into_bytes())
351 }
352
353 #[cfg(test)]
354 mod tests {
355 use super::*;
356
357 #[test]
358 fn backup_code_generation_produces_correct_count() {
359 let codes = generate_backup_codes();
360 assert_eq!(codes.len(), BACKUP_CODE_COUNT);
361 }
362
363 #[test]
364 fn backup_codes_are_correct_length() {
365 let codes = generate_backup_codes();
366 for code in &codes {
367 assert_eq!(code.len(), BACKUP_CODE_LENGTH);
368 }
369 }
370
371 #[test]
372 fn backup_codes_are_alphanumeric() {
373 let codes = generate_backup_codes();
374 for code in &codes {
375 assert!(
376 code.chars().all(|c| c.is_ascii_alphanumeric()),
377 "Code should be alphanumeric: {code}"
378 );
379 }
380 }
381
382 #[test]
383 fn backup_codes_are_unique() {
384 let codes = generate_backup_codes();
385 let unique: std::collections::HashSet<&String> = codes.iter().collect();
386 assert_eq!(unique.len(), codes.len());
387 }
388
389 #[test]
390 fn hash_backup_code_is_argon2_phc() {
391 // Argon2 hashes are non-deterministic (unique salt per call) and use
392 // the PHC string format that starts with `$argon2`.
393 let h = hash_backup_code("abc12345", "secret");
394 assert!(h.starts_with("$argon2"), "got {h}");
395 }
396
397 #[test]
398 fn hash_backup_code_non_deterministic() {
399 // Two hashes of the same code must differ, distinct salts.
400 let h1 = hash_backup_code("abc12345", "secret");
401 let h2 = hash_backup_code("abc12345", "secret");
402 assert_ne!(h1, h2);
403 }
404
405 #[test]
406 fn hash_backup_code_verifies_against_itself() {
407 use argon2::{Argon2, PasswordHash, password_hash::PasswordVerifier};
408 let h = hash_backup_code("abc12345", "ignored");
409 let parsed = PasswordHash::new(&h).unwrap();
410 assert!(
411 Argon2::default()
412 .verify_password(b"abc12345", &parsed)
413 .is_ok()
414 );
415 assert!(
416 Argon2::default()
417 .verify_password(b"wrong", &parsed)
418 .is_err()
419 );
420 }
421
422 #[test]
423 fn legacy_hmac_is_deterministic_and_secret_keyed() {
424 let h1 = legacy_hmac_backup_code("abc12345", "secret");
425 let h2 = legacy_hmac_backup_code("abc12345", "secret");
426 assert_eq!(h1, h2);
427 let h3 = legacy_hmac_backup_code("abc12345", "different-secret");
428 assert_ne!(h1, h3);
429 }
430 }
431