Skip to main content

max / makenotwork

8.2 KB · 229 lines History Blame Raw
1 //! Two-factor authentication queries: TOTP secrets, backup codes.
2
3 use sqlx::PgPool;
4
5 use super::UserId;
6 use crate::error::Result;
7
8 /// Get the stored TOTP secret for a user, decrypted (None if not set up).
9 ///
10 /// Secrets are encrypted at rest with the global signing secret; this returns
11 /// the plaintext base32 seed. A stored value that isn't a valid `enc:v1:`
12 /// ciphertext (e.g. a pre-encryption plaintext seed, backwards compat was
13 /// cut) errors out, forcing the user to re-enroll their authenticator.
14 #[tracing::instrument(skip_all)]
15 pub(crate) async fn get_totp_secret(
16 pool: &PgPool,
17 user_id: UserId,
18 signing_secret: &str,
19 ) -> Result<Option<String>> {
20 let stored: Option<String> = sqlx::query_scalar("SELECT totp_secret FROM users WHERE id = $1")
21 .bind(user_id)
22 .fetch_one(pool)
23 .await?;
24
25 stored
26 .map(|s| crate::crypto::decrypt_totp_secret(&s, signing_secret))
27 .transpose()
28 }
29
30 /// Store a TOTP secret for a user (does not enable 2FA yet).
31 ///
32 /// The plaintext base32 seed is encrypted at rest with the signing secret
33 /// before it touches the database, so a DB read alone cannot recover a usable
34 /// second factor.
35 #[tracing::instrument(skip_all)]
36 pub(crate) async fn set_totp_secret(
37 pool: &PgPool,
38 user_id: UserId,
39 secret: &str,
40 signing_secret: &str,
41 ) -> Result<()> {
42 let encrypted = crate::crypto::encrypt_totp_secret(secret, signing_secret);
43
44 // Clear the replay step alongside the secret. Without this, a user who
45 // disables and re-enables TOTP (potentially with a new secret) inherits a
46 // stale `totp_last_used_step` and any first-attempt code in a lower step
47 // window is false-rejected as a replay.
48 sqlx::query("UPDATE users SET totp_secret = $2, totp_last_used_step = NULL WHERE id = $1")
49 .bind(user_id)
50 .bind(encrypted)
51 .execute(pool)
52 .await?;
53
54 Ok(())
55 }
56
57 /// Enable TOTP 2FA for a user (called after first successful code verification).
58 #[tracing::instrument(skip_all)]
59 pub(crate) async fn enable_totp(pool: &PgPool, user_id: UserId) -> Result<()> {
60 sqlx::query("UPDATE users SET totp_enabled = true WHERE id = $1")
61 .bind(user_id)
62 .execute(pool)
63 .await?;
64
65 Ok(())
66 }
67
68 /// Disable TOTP 2FA: clear the secret, set enabled to false, delete backup codes.
69 #[tracing::instrument(skip_all)]
70 pub(crate) async fn disable_totp(pool: &PgPool, user_id: UserId) -> Result<()> {
71 sqlx::query("UPDATE users SET totp_secret = NULL, totp_enabled = false, totp_last_used_step = NULL WHERE id = $1")
72 .bind(user_id)
73 .execute(pool)
74 .await?;
75
76 sqlx::query("DELETE FROM backup_codes WHERE user_id = $1")
77 .bind(user_id)
78 .execute(pool)
79 .await?;
80
81 Ok(())
82 }
83
84 /// Atomically advance the last accepted TOTP time step. Returns `true` if this
85 /// call recorded the step (it was strictly newer than the stored one), `false`
86 /// if a concurrent verification already consumed this step or a later one.
87 ///
88 /// The `AND $2 > totp_last_used_step` guard makes the read-then-write in the
89 /// caller race-free: two concurrent submissions of the same code both read the
90 /// old step and both match, but only the winner's UPDATE affects a row, so a
91 /// code cannot be accepted twice within its 30s window. Callers must gate
92 /// acceptance on the returned bool, not on their own prior step read.
93 #[tracing::instrument(skip_all)]
94 pub(crate) async fn set_totp_last_used_step(
95 pool: &PgPool,
96 user_id: UserId,
97 step: i64,
98 ) -> Result<bool> {
99 let updated = sqlx::query_scalar::<_, UserId>(
100 "UPDATE users SET totp_last_used_step = $2 \
101 WHERE id = $1 AND $2 > COALESCE(totp_last_used_step, 0) RETURNING id",
102 )
103 .bind(user_id)
104 .bind(step)
105 .fetch_optional(pool)
106 .await?;
107
108 Ok(updated.is_some())
109 }
110
111 /// Check if a user has TOTP 2FA enabled.
112 #[tracing::instrument(skip_all)]
113 pub(crate) async fn is_totp_enabled(pool: &PgPool, user_id: UserId) -> Result<bool> {
114 let enabled: bool = sqlx::query_scalar("SELECT totp_enabled FROM users WHERE id = $1")
115 .bind(user_id)
116 .fetch_one(pool)
117 .await?;
118
119 Ok(enabled)
120 }
121
122 /// Delete existing backup codes and insert new ones (atomic replacement).
123 #[tracing::instrument(skip_all)]
124 pub(crate) async fn create_backup_codes(
125 pool: &PgPool,
126 user_id: UserId,
127 code_hashes: &[String],
128 ) -> Result<()> {
129 let mut tx = pool.begin().await?;
130
131 // Delete any existing codes
132 sqlx::query("DELETE FROM backup_codes WHERE user_id = $1")
133 .bind(user_id)
134 .execute(&mut *tx)
135 .await?;
136
137 // Batch insert all codes in a single query
138 sqlx::query("INSERT INTO backup_codes (user_id, code_hash) SELECT $1, UNNEST($2::text[])")
139 .bind(user_id)
140 .bind(code_hashes)
141 .execute(&mut *tx)
142 .await?;
143
144 tx.commit().await?;
145 Ok(())
146 }
147
148 /// Verify a backup code and mark it as used if found.
149 ///
150 /// `code` is the raw 8-char token the user typed; `legacy_hmac` is the
151 /// HMAC-SHA256 of the same code (passed in pre-computed by the caller so the
152 /// secret stays in route-layer scope). Returns `Ok(true)` when a matching
153 /// unused code is consumed.
154 ///
155 /// Dual-read window: rows hashed under the old HMAC scheme remain valid
156 /// until the user regenerates their backup codes (each regeneration writes
157 /// fresh Argon2 hashes). Argon2 PHC strings begin with `$argon2`; anything
158 /// else is treated as a legacy 64-char hex HMAC.
159 #[tracing::instrument(skip_all)]
160 pub(crate) async fn verify_and_consume_backup_code(
161 pool: &PgPool,
162 user_id: UserId,
163 code: &str,
164 legacy_hmac: &str,
165 ) -> Result<bool> {
166 use argon2::{Argon2, PasswordHash, password_hash::PasswordVerifier};
167
168 let rows: Vec<(uuid::Uuid, String)> = sqlx::query_as(
169 "SELECT id, code_hash FROM backup_codes WHERE user_id = $1 AND used_at IS NULL",
170 )
171 .bind(user_id)
172 .fetch_all(pool)
173 .await?;
174
175 // Timing note: the loop `break`s on the first match, which is NOT a usable
176 // timing oracle. A wrong guess (the attacker's case) matches nothing, so the
177 // loop always runs to completion and scans every row in constant time,
178 // independent of code ordering. The early exit fires only on a *successful*
179 // verify, by which point the caller already supplied a valid code and has
180 // nothing left to learn. Retaining the break also avoids forcing N Argon2
181 // verifications (each ~46 MiB) on every attempt, which would hand an attacker
182 // a memory-amplification lever on the 2FA endpoint. Brute force is bounded
183 // separately by the shared failed-attempt lockout.
184 // Run the per-code Argon2 verifies on a blocking thread so a 2FA attempt
185 // (up to N verifies) can't occupy a Tokio worker.
186 let matched_id: Option<uuid::Uuid> = {
187 let code = code.to_string();
188 let legacy_hmac = legacy_hmac.to_string();
189 tokio::task::spawn_blocking(move || {
190 for (id, stored) in &rows {
191 let is_match = if stored.starts_with("$argon2") {
192 match PasswordHash::new(stored) {
193 Ok(parsed) => Argon2::default()
194 .verify_password(code.as_bytes(), &parsed)
195 .is_ok(),
196 Err(e) => {
197 tracing::warn!(error = %e, "malformed argon2 backup code hash in DB; skipping");
198 false
199 }
200 }
201 } else {
202 // Legacy HMAC-SHA256 hex. Length-equality short-circuits
203 // before the constant-time compare, matching the existing
204 // behavior of `crypto::constant_time_compare`.
205 crate::crypto::constant_time_compare(stored, &legacy_hmac)
206 };
207 if is_match {
208 return Some(*id);
209 }
210 }
211 None
212 })
213 .await
214 .map_err(|e| anyhow::anyhow!("backup-code verify task join: {e}"))?
215 };
216
217 let Some(id) = matched_id else {
218 return Ok(false);
219 };
220
221 let result =
222 sqlx::query("UPDATE backup_codes SET used_at = NOW() WHERE id = $1 AND used_at IS NULL")
223 .bind(id)
224 .execute(pool)
225 .await?;
226
227 Ok(result.rows_affected() > 0)
228 }
229