this repo has no description
lewis.moe
1use crate::api::EmptyResponse;
2use crate::api::error::ApiError;
3use crate::auth::BearerAuth;
4use crate::auth::{
5 decrypt_totp_secret, encrypt_totp_secret, generate_backup_codes, generate_qr_png_base64,
6 generate_totp_secret, generate_totp_uri, hash_backup_code, is_backup_code_format,
7 verify_backup_code, verify_totp_code,
8};
9use crate::state::{AppState, RateLimitKind};
10use crate::types::PlainPassword;
11use axum::{
12 Json,
13 extract::State,
14 response::{IntoResponse, Response},
15};
16use chrono::Utc;
17use serde::{Deserialize, Serialize};
18use tracing::{error, info, warn};
19
20const ENCRYPTION_VERSION: i32 = 1;
21
22#[derive(Serialize)]
23#[serde(rename_all = "camelCase")]
24pub struct CreateTotpSecretResponse {
25 pub secret: String,
26 pub uri: String,
27 pub qr_base64: String,
28}
29
30pub async fn create_totp_secret(State(state): State<AppState>, auth: BearerAuth) -> Response {
31 let existing = sqlx::query_scalar!(
32 "SELECT verified FROM user_totp WHERE did = $1",
33 &*&auth.0.did
34 )
35 .fetch_optional(&state.db)
36 .await;
37
38 if let Ok(Some(true)) = existing {
39 return ApiError::TotpAlreadyEnabled.into_response();
40 }
41
42 let secret = generate_totp_secret();
43
44 let handle = sqlx::query_scalar!("SELECT handle FROM users WHERE did = $1", &*&auth.0.did)
45 .fetch_optional(&state.db)
46 .await;
47
48 let handle = match handle {
49 Ok(Some(h)) => h,
50 Ok(None) => return ApiError::AccountNotFound.into_response(),
51 Err(e) => {
52 error!("DB error fetching handle: {:?}", e);
53 return ApiError::InternalError(None).into_response();
54 }
55 };
56
57 let hostname = std::env::var("PDS_HOSTNAME").unwrap_or_else(|_| "localhost".to_string());
58 let uri = generate_totp_uri(&secret, &handle, &hostname);
59
60 let qr_code = match generate_qr_png_base64(&secret, &handle, &hostname) {
61 Ok(qr) => qr,
62 Err(e) => {
63 error!("Failed to generate QR code: {:?}", e);
64 return ApiError::InternalError(Some("Failed to generate QR code".into()))
65 .into_response();
66 }
67 };
68
69 let encrypted_secret = match encrypt_totp_secret(&secret) {
70 Ok(enc) => enc,
71 Err(e) => {
72 error!("Failed to encrypt TOTP secret: {:?}", e);
73 return ApiError::InternalError(None).into_response();
74 }
75 };
76
77 let result = sqlx::query!(
78 r#"
79 INSERT INTO user_totp (did, secret_encrypted, encryption_version, verified, created_at)
80 VALUES ($1, $2, $3, false, NOW())
81 ON CONFLICT (did) DO UPDATE SET
82 secret_encrypted = $2,
83 encryption_version = $3,
84 verified = false,
85 created_at = NOW(),
86 last_used = NULL
87 "#,
88 &auth.0.did,
89 encrypted_secret,
90 ENCRYPTION_VERSION
91 )
92 .execute(&state.db)
93 .await;
94
95 if let Err(e) = result {
96 error!("Failed to store TOTP secret: {:?}", e);
97 return ApiError::InternalError(None).into_response();
98 }
99
100 let secret_base32 = base32::encode(base32::Alphabet::Rfc4648 { padding: false }, &secret);
101
102 info!(did = %&auth.0.did, "TOTP secret created (pending verification)");
103
104 Json(CreateTotpSecretResponse {
105 secret: secret_base32,
106 uri,
107 qr_base64: qr_code,
108 })
109 .into_response()
110}
111
112#[derive(Deserialize)]
113pub struct EnableTotpInput {
114 pub code: String,
115}
116
117#[derive(Serialize)]
118#[serde(rename_all = "camelCase")]
119pub struct EnableTotpResponse {
120 pub backup_codes: Vec<String>,
121}
122
123pub async fn enable_totp(
124 State(state): State<AppState>,
125 auth: BearerAuth,
126 Json(input): Json<EnableTotpInput>,
127) -> Response {
128 if !state
129 .check_rate_limit(RateLimitKind::TotpVerify, &auth.0.did)
130 .await
131 {
132 warn!(did = %&auth.0.did, "TOTP verification rate limit exceeded");
133 return ApiError::RateLimitExceeded(None).into_response();
134 }
135
136 let totp_row = sqlx::query!(
137 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
138 &auth.0.did
139 )
140 .fetch_optional(&state.db)
141 .await;
142
143 let totp_row = match totp_row {
144 Ok(Some(row)) => row,
145 Ok(None) => return ApiError::TotpNotEnabled.into_response(),
146 Err(e) => {
147 error!("DB error fetching TOTP: {:?}", e);
148 return ApiError::InternalError(None).into_response();
149 }
150 };
151
152 if totp_row.verified {
153 return ApiError::TotpAlreadyEnabled.into_response();
154 }
155
156 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
157 {
158 Ok(s) => s,
159 Err(e) => {
160 error!("Failed to decrypt TOTP secret: {:?}", e);
161 return ApiError::InternalError(None).into_response();
162 }
163 };
164
165 let code = input.code.trim();
166 if !verify_totp_code(&secret, code) {
167 return ApiError::InvalidCode(Some("Invalid verification code".into())).into_response();
168 }
169
170 let backup_codes = generate_backup_codes();
171 let mut tx = match state.db.begin().await {
172 Ok(tx) => tx,
173 Err(e) => {
174 error!("Failed to begin transaction: {:?}", e);
175 return ApiError::InternalError(None).into_response();
176 }
177 };
178
179 if let Err(e) = sqlx::query!(
180 "UPDATE user_totp SET verified = true, last_used = NOW() WHERE did = $1",
181 &auth.0.did
182 )
183 .execute(&mut *tx)
184 .await
185 {
186 error!("Failed to enable TOTP: {:?}", e);
187 return ApiError::InternalError(None).into_response();
188 }
189
190 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", &*&auth.0.did)
191 .execute(&mut *tx)
192 .await
193 {
194 error!("Failed to clear old backup codes: {:?}", e);
195 return ApiError::InternalError(None).into_response();
196 }
197
198 let backup_hashes: Result<Vec<_>, _> =
199 backup_codes.iter().map(|c| hash_backup_code(c)).collect();
200 let backup_hashes = match backup_hashes {
201 Ok(hashes) => hashes,
202 Err(e) => {
203 error!("Failed to hash backup code: {:?}", e);
204 return ApiError::InternalError(None).into_response();
205 }
206 };
207
208 if let Err(e) = sqlx::query!(
209 r#"
210 INSERT INTO backup_codes (did, code_hash, created_at)
211 SELECT $1, hash, NOW() FROM UNNEST($2::text[]) AS t(hash)
212 "#,
213 &auth.0.did,
214 &backup_hashes[..]
215 )
216 .execute(&mut *tx)
217 .await
218 {
219 error!("Failed to store backup codes: {:?}", e);
220 return ApiError::InternalError(None).into_response();
221 }
222
223 if let Err(e) = tx.commit().await {
224 error!("Failed to commit transaction: {:?}", e);
225 return ApiError::InternalError(None).into_response();
226 }
227
228 info!(did = %&auth.0.did, "TOTP enabled with {} backup codes", backup_codes.len());
229
230 Json(EnableTotpResponse { backup_codes }).into_response()
231}
232
233#[derive(Deserialize)]
234pub struct DisableTotpInput {
235 pub password: PlainPassword,
236 pub code: String,
237}
238
239pub async fn disable_totp(
240 State(state): State<AppState>,
241 auth: BearerAuth,
242 Json(input): Json<DisableTotpInput>,
243) -> Response {
244 if !crate::api::server::reauth::check_legacy_session_mfa(&state.db, &auth.0.did).await {
245 return crate::api::server::reauth::legacy_mfa_required_response(&state.db, &auth.0.did)
246 .await;
247 }
248
249 if !state
250 .check_rate_limit(RateLimitKind::TotpVerify, &auth.0.did)
251 .await
252 {
253 warn!(did = %&auth.0.did, "TOTP verification rate limit exceeded");
254 return ApiError::RateLimitExceeded(None).into_response();
255 }
256
257 let user = sqlx::query!(
258 "SELECT password_hash FROM users WHERE did = $1",
259 &*&auth.0.did
260 )
261 .fetch_optional(&state.db)
262 .await;
263
264 let password_hash = match user {
265 Ok(Some(row)) => row.password_hash,
266 Ok(None) => return ApiError::AccountNotFound.into_response(),
267 Err(e) => {
268 error!("DB error fetching user: {:?}", e);
269 return ApiError::InternalError(None).into_response();
270 }
271 };
272
273 let password_valid = password_hash
274 .as_ref()
275 .map(|h| bcrypt::verify(&input.password, h).unwrap_or(false))
276 .unwrap_or(false);
277 if !password_valid {
278 return ApiError::InvalidPassword("Password is incorrect".into()).into_response();
279 }
280
281 let totp_row = sqlx::query!(
282 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
283 &auth.0.did
284 )
285 .fetch_optional(&state.db)
286 .await;
287
288 let totp_row = match totp_row {
289 Ok(Some(row)) if row.verified => row,
290 Ok(Some(_)) | Ok(None) => return ApiError::TotpNotEnabled.into_response(),
291 Err(e) => {
292 error!("DB error fetching TOTP: {:?}", e);
293 return ApiError::InternalError(None).into_response();
294 }
295 };
296
297 let code = input.code.trim();
298 let code_valid = if is_backup_code_format(code) {
299 verify_backup_code_for_user(&state, &auth.0.did, code).await
300 } else {
301 let secret =
302 match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version) {
303 Ok(s) => s,
304 Err(e) => {
305 error!("Failed to decrypt TOTP secret: {:?}", e);
306 return ApiError::InternalError(None).into_response();
307 }
308 };
309 verify_totp_code(&secret, code)
310 };
311
312 if !code_valid {
313 return ApiError::InvalidCode(Some("Invalid verification code".into())).into_response();
314 }
315
316 let mut tx = match state.db.begin().await {
317 Ok(tx) => tx,
318 Err(e) => {
319 error!("Failed to begin transaction: {:?}", e);
320 return ApiError::InternalError(None).into_response();
321 }
322 };
323
324 if let Err(e) = sqlx::query!("DELETE FROM user_totp WHERE did = $1", &*&auth.0.did)
325 .execute(&mut *tx)
326 .await
327 {
328 error!("Failed to delete TOTP: {:?}", e);
329 return ApiError::InternalError(None).into_response();
330 }
331
332 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", &*&auth.0.did)
333 .execute(&mut *tx)
334 .await
335 {
336 error!("Failed to delete backup codes: {:?}", e);
337 return ApiError::InternalError(None).into_response();
338 }
339
340 if let Err(e) = tx.commit().await {
341 error!("Failed to commit transaction: {:?}", e);
342 return ApiError::InternalError(None).into_response();
343 }
344
345 info!(did = %&auth.0.did, "TOTP disabled");
346
347 EmptyResponse::ok().into_response()
348}
349
350#[derive(Serialize)]
351#[serde(rename_all = "camelCase")]
352pub struct GetTotpStatusResponse {
353 pub enabled: bool,
354 pub has_backup_codes: bool,
355 pub backup_codes_remaining: i64,
356}
357
358pub async fn get_totp_status(State(state): State<AppState>, auth: BearerAuth) -> Response {
359 let totp_row = sqlx::query!(
360 "SELECT verified FROM user_totp WHERE did = $1",
361 &*&auth.0.did
362 )
363 .fetch_optional(&state.db)
364 .await;
365
366 let enabled = match totp_row {
367 Ok(Some(row)) => row.verified,
368 Ok(None) => false,
369 Err(e) => {
370 error!("DB error fetching TOTP status: {:?}", e);
371 return ApiError::InternalError(None).into_response();
372 }
373 };
374
375 let backup_count_row = sqlx::query!(
376 "SELECT COUNT(*) as count FROM backup_codes WHERE did = $1 AND used_at IS NULL",
377 &auth.0.did
378 )
379 .fetch_one(&state.db)
380 .await;
381
382 let backup_count = backup_count_row.map(|r| r.count.unwrap_or(0)).unwrap_or(0);
383
384 Json(GetTotpStatusResponse {
385 enabled,
386 has_backup_codes: backup_count > 0,
387 backup_codes_remaining: backup_count,
388 })
389 .into_response()
390}
391
392#[derive(Deserialize)]
393pub struct RegenerateBackupCodesInput {
394 pub password: PlainPassword,
395 pub code: String,
396}
397
398#[derive(Serialize)]
399#[serde(rename_all = "camelCase")]
400pub struct RegenerateBackupCodesResponse {
401 pub backup_codes: Vec<String>,
402}
403
404pub async fn regenerate_backup_codes(
405 State(state): State<AppState>,
406 auth: BearerAuth,
407 Json(input): Json<RegenerateBackupCodesInput>,
408) -> Response {
409 if !state
410 .check_rate_limit(RateLimitKind::TotpVerify, &auth.0.did)
411 .await
412 {
413 warn!(did = %&auth.0.did, "TOTP verification rate limit exceeded");
414 return ApiError::RateLimitExceeded(None).into_response();
415 }
416
417 let user = sqlx::query!(
418 "SELECT password_hash FROM users WHERE did = $1",
419 &*&auth.0.did
420 )
421 .fetch_optional(&state.db)
422 .await;
423
424 let password_hash = match user {
425 Ok(Some(row)) => row.password_hash,
426 Ok(None) => return ApiError::AccountNotFound.into_response(),
427 Err(e) => {
428 error!("DB error fetching user: {:?}", e);
429 return ApiError::InternalError(None).into_response();
430 }
431 };
432
433 let password_valid = password_hash
434 .as_ref()
435 .map(|h| bcrypt::verify(&input.password, h).unwrap_or(false))
436 .unwrap_or(false);
437 if !password_valid {
438 return ApiError::InvalidPassword("Password is incorrect".into()).into_response();
439 }
440
441 let totp_row = sqlx::query!(
442 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
443 &auth.0.did
444 )
445 .fetch_optional(&state.db)
446 .await;
447
448 let totp_row = match totp_row {
449 Ok(Some(row)) if row.verified => row,
450 Ok(Some(_)) | Ok(None) => return ApiError::TotpNotEnabled.into_response(),
451 Err(e) => {
452 error!("DB error fetching TOTP: {:?}", e);
453 return ApiError::InternalError(None).into_response();
454 }
455 };
456
457 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
458 {
459 Ok(s) => s,
460 Err(e) => {
461 error!("Failed to decrypt TOTP secret: {:?}", e);
462 return ApiError::InternalError(None).into_response();
463 }
464 };
465
466 let code = input.code.trim();
467 if !verify_totp_code(&secret, code) {
468 return ApiError::InvalidCode(Some("Invalid verification code".into())).into_response();
469 }
470
471 let backup_codes = generate_backup_codes();
472 let mut tx = match state.db.begin().await {
473 Ok(tx) => tx,
474 Err(e) => {
475 error!("Failed to begin transaction: {:?}", e);
476 return ApiError::InternalError(None).into_response();
477 }
478 };
479
480 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", &*&auth.0.did)
481 .execute(&mut *tx)
482 .await
483 {
484 error!("Failed to clear old backup codes: {:?}", e);
485 return ApiError::InternalError(None).into_response();
486 }
487
488 let backup_hashes: Result<Vec<_>, _> =
489 backup_codes.iter().map(|c| hash_backup_code(c)).collect();
490 let backup_hashes = match backup_hashes {
491 Ok(hashes) => hashes,
492 Err(e) => {
493 error!("Failed to hash backup code: {:?}", e);
494 return ApiError::InternalError(None).into_response();
495 }
496 };
497
498 if let Err(e) = sqlx::query!(
499 r#"
500 INSERT INTO backup_codes (did, code_hash, created_at)
501 SELECT $1, hash, NOW() FROM UNNEST($2::text[]) AS t(hash)
502 "#,
503 &auth.0.did,
504 &backup_hashes[..]
505 )
506 .execute(&mut *tx)
507 .await
508 {
509 error!("Failed to store backup codes: {:?}", e);
510 return ApiError::InternalError(None).into_response();
511 }
512
513 if let Err(e) = tx.commit().await {
514 error!("Failed to commit transaction: {:?}", e);
515 return ApiError::InternalError(None).into_response();
516 }
517
518 info!(did = %&auth.0.did, "Backup codes regenerated");
519
520 Json(RegenerateBackupCodesResponse { backup_codes }).into_response()
521}
522
523async fn verify_backup_code_for_user(state: &AppState, did: &str, code: &str) -> bool {
524 let code = code.trim().to_uppercase();
525
526 let backup_codes = sqlx::query!(
527 "SELECT id, code_hash FROM backup_codes WHERE did = $1 AND used_at IS NULL",
528 did
529 )
530 .fetch_all(&state.db)
531 .await;
532
533 let backup_codes = match backup_codes {
534 Ok(codes) => codes,
535 Err(e) => {
536 warn!("Failed to fetch backup codes: {:?}", e);
537 return false;
538 }
539 };
540
541 let matched = backup_codes
542 .iter()
543 .find(|row| verify_backup_code(&code, &row.code_hash));
544
545 match matched {
546 Some(row) => {
547 let _ = sqlx::query!(
548 "UPDATE backup_codes SET used_at = $1 WHERE id = $2",
549 Utc::now(),
550 row.id
551 )
552 .execute(&state.db)
553 .await;
554 true
555 }
556 None => false,
557 }
558}
559
560pub async fn verify_totp_or_backup_for_user(state: &AppState, did: &str, code: &str) -> bool {
561 let code = code.trim();
562
563 if is_backup_code_format(code) {
564 return verify_backup_code_for_user(state, did, code).await;
565 }
566
567 let totp_row = sqlx::query!(
568 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
569 did
570 )
571 .fetch_optional(&state.db)
572 .await;
573
574 let totp_row = match totp_row {
575 Ok(Some(row)) if row.verified => row,
576 _ => return false,
577 };
578
579 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
580 {
581 Ok(s) => s,
582 Err(_) => return false,
583 };
584
585 if verify_totp_code(&secret, code) {
586 let _ = sqlx::query!("UPDATE user_totp SET last_used = NOW() WHERE did = $1", did)
587 .execute(&state.db)
588 .await;
589 return true;
590 }
591
592 false
593}
594
595pub async fn has_totp_enabled(state: &AppState, did: &str) -> bool {
596 has_totp_enabled_db(&state.db, did).await
597}
598
599pub async fn has_totp_enabled_db(db: &sqlx::PgPool, did: &str) -> bool {
600 let result = sqlx::query_scalar!("SELECT verified FROM user_totp WHERE did = $1", did)
601 .fetch_optional(db)
602 .await;
603
604 matches!(result, Ok(Some(true)))
605}