this repo has no description
lewis.moe
1use crate::auth::BearerAuth;
2use crate::auth::totp::{
3 decrypt_totp_secret, encrypt_totp_secret, generate_backup_codes, generate_qr_png_base64,
4 generate_totp_secret, generate_totp_uri, hash_backup_code, is_backup_code_format,
5 verify_backup_code, verify_totp_code,
6};
7use crate::state::{AppState, RateLimitKind};
8use axum::{
9 Json,
10 extract::State,
11 http::StatusCode,
12 response::{IntoResponse, Response},
13};
14use chrono::Utc;
15use serde::{Deserialize, Serialize};
16use serde_json::json;
17use tracing::{error, info, warn};
18
19const ENCRYPTION_VERSION: i32 = 1;
20
21#[derive(Serialize)]
22#[serde(rename_all = "camelCase")]
23pub struct CreateTotpSecretResponse {
24 pub secret: String,
25 pub uri: String,
26 pub qr_base64: String,
27}
28
29pub async fn create_totp_secret(State(state): State<AppState>, auth: BearerAuth) -> Response {
30 let existing = sqlx::query_scalar!("SELECT verified FROM user_totp WHERE did = $1", auth.0.did)
31 .fetch_optional(&state.db)
32 .await;
33
34 if let Ok(Some(true)) = existing {
35 return (
36 StatusCode::CONFLICT,
37 Json(json!({
38 "error": "TotpAlreadyEnabled",
39 "message": "TOTP is already enabled for this account"
40 })),
41 )
42 .into_response();
43 }
44
45 let secret = generate_totp_secret();
46
47 let handle = sqlx::query_scalar!("SELECT handle FROM users WHERE did = $1", auth.0.did)
48 .fetch_optional(&state.db)
49 .await;
50
51 let handle = match handle {
52 Ok(Some(h)) => h,
53 Ok(None) => {
54 return (
55 StatusCode::NOT_FOUND,
56 Json(json!({"error": "AccountNotFound", "message": "Account not found"})),
57 )
58 .into_response();
59 }
60 Err(e) => {
61 error!("DB error fetching handle: {:?}", e);
62 return (
63 StatusCode::INTERNAL_SERVER_ERROR,
64 Json(json!({"error": "InternalError"})),
65 )
66 .into_response();
67 }
68 };
69
70 let hostname = std::env::var("PDS_HOSTNAME").unwrap_or_else(|_| "localhost".to_string());
71 let uri = generate_totp_uri(&secret, &handle, &hostname);
72
73 let qr_code = match generate_qr_png_base64(&secret, &handle, &hostname) {
74 Ok(qr) => qr,
75 Err(e) => {
76 error!("Failed to generate QR code: {:?}", e);
77 return (
78 StatusCode::INTERNAL_SERVER_ERROR,
79 Json(json!({"error": "InternalError", "message": "Failed to generate QR code"})),
80 )
81 .into_response();
82 }
83 };
84
85 let encrypted_secret = match encrypt_totp_secret(&secret) {
86 Ok(enc) => enc,
87 Err(e) => {
88 error!("Failed to encrypt TOTP secret: {:?}", e);
89 return (
90 StatusCode::INTERNAL_SERVER_ERROR,
91 Json(json!({"error": "InternalError"})),
92 )
93 .into_response();
94 }
95 };
96
97 let result = sqlx::query!(
98 r#"
99 INSERT INTO user_totp (did, secret_encrypted, encryption_version, verified, created_at)
100 VALUES ($1, $2, $3, false, NOW())
101 ON CONFLICT (did) DO UPDATE SET
102 secret_encrypted = $2,
103 encryption_version = $3,
104 verified = false,
105 created_at = NOW(),
106 last_used = NULL
107 "#,
108 auth.0.did,
109 encrypted_secret,
110 ENCRYPTION_VERSION
111 )
112 .execute(&state.db)
113 .await;
114
115 if let Err(e) = result {
116 error!("Failed to store TOTP secret: {:?}", e);
117 return (
118 StatusCode::INTERNAL_SERVER_ERROR,
119 Json(json!({"error": "InternalError"})),
120 )
121 .into_response();
122 }
123
124 let secret_base32 = base32::encode(base32::Alphabet::Rfc4648 { padding: false }, &secret);
125
126 info!(did = %auth.0.did, "TOTP secret created (pending verification)");
127
128 Json(CreateTotpSecretResponse {
129 secret: secret_base32,
130 uri,
131 qr_base64: qr_code,
132 })
133 .into_response()
134}
135
136#[derive(Deserialize)]
137pub struct EnableTotpInput {
138 pub code: String,
139}
140
141#[derive(Serialize)]
142#[serde(rename_all = "camelCase")]
143pub struct EnableTotpResponse {
144 pub backup_codes: Vec<String>,
145}
146
147pub async fn enable_totp(
148 State(state): State<AppState>,
149 auth: BearerAuth,
150 Json(input): Json<EnableTotpInput>,
151) -> Response {
152 if !state
153 .check_rate_limit(RateLimitKind::TotpVerify, &auth.0.did)
154 .await
155 {
156 warn!(did = %auth.0.did, "TOTP verification rate limit exceeded");
157 return (
158 StatusCode::TOO_MANY_REQUESTS,
159 Json(json!({
160 "error": "RateLimitExceeded",
161 "message": "Too many verification attempts. Please try again in a few minutes."
162 })),
163 )
164 .into_response();
165 }
166
167 let totp_row = sqlx::query!(
168 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
169 auth.0.did
170 )
171 .fetch_optional(&state.db)
172 .await;
173
174 let totp_row = match totp_row {
175 Ok(Some(row)) => row,
176 Ok(None) => {
177 return (
178 StatusCode::BAD_REQUEST,
179 Json(json!({
180 "error": "TotpNotSetup",
181 "message": "Please call createTotpSecret first"
182 })),
183 )
184 .into_response();
185 }
186 Err(e) => {
187 error!("DB error fetching TOTP: {:?}", e);
188 return (
189 StatusCode::INTERNAL_SERVER_ERROR,
190 Json(json!({"error": "InternalError"})),
191 )
192 .into_response();
193 }
194 };
195
196 if totp_row.verified {
197 return (
198 StatusCode::CONFLICT,
199 Json(json!({
200 "error": "TotpAlreadyEnabled",
201 "message": "TOTP is already enabled"
202 })),
203 )
204 .into_response();
205 }
206
207 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
208 {
209 Ok(s) => s,
210 Err(e) => {
211 error!("Failed to decrypt TOTP secret: {:?}", e);
212 return (
213 StatusCode::INTERNAL_SERVER_ERROR,
214 Json(json!({"error": "InternalError"})),
215 )
216 .into_response();
217 }
218 };
219
220 let code = input.code.trim();
221 if !verify_totp_code(&secret, code) {
222 return (
223 StatusCode::UNAUTHORIZED,
224 Json(json!({
225 "error": "InvalidCode",
226 "message": "Invalid verification code"
227 })),
228 )
229 .into_response();
230 }
231
232 let backup_codes = generate_backup_codes();
233 let mut tx = match state.db.begin().await {
234 Ok(tx) => tx,
235 Err(e) => {
236 error!("Failed to begin transaction: {:?}", e);
237 return (
238 StatusCode::INTERNAL_SERVER_ERROR,
239 Json(json!({"error": "InternalError"})),
240 )
241 .into_response();
242 }
243 };
244
245 if let Err(e) = sqlx::query!(
246 "UPDATE user_totp SET verified = true, last_used = NOW() WHERE did = $1",
247 auth.0.did
248 )
249 .execute(&mut *tx)
250 .await
251 {
252 error!("Failed to enable TOTP: {:?}", e);
253 return (
254 StatusCode::INTERNAL_SERVER_ERROR,
255 Json(json!({"error": "InternalError"})),
256 )
257 .into_response();
258 }
259
260 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", auth.0.did)
261 .execute(&mut *tx)
262 .await
263 {
264 error!("Failed to clear old backup codes: {:?}", e);
265 return (
266 StatusCode::INTERNAL_SERVER_ERROR,
267 Json(json!({"error": "InternalError"})),
268 )
269 .into_response();
270 }
271
272 for code in &backup_codes {
273 let hash = match hash_backup_code(code) {
274 Ok(h) => h,
275 Err(e) => {
276 error!("Failed to hash backup code: {:?}", e);
277 return (
278 StatusCode::INTERNAL_SERVER_ERROR,
279 Json(json!({"error": "InternalError"})),
280 )
281 .into_response();
282 }
283 };
284
285 if let Err(e) = sqlx::query!(
286 "INSERT INTO backup_codes (did, code_hash, created_at) VALUES ($1, $2, NOW())",
287 auth.0.did,
288 hash
289 )
290 .execute(&mut *tx)
291 .await
292 {
293 error!("Failed to store backup code: {:?}", e);
294 return (
295 StatusCode::INTERNAL_SERVER_ERROR,
296 Json(json!({"error": "InternalError"})),
297 )
298 .into_response();
299 }
300 }
301
302 if let Err(e) = tx.commit().await {
303 error!("Failed to commit transaction: {:?}", e);
304 return (
305 StatusCode::INTERNAL_SERVER_ERROR,
306 Json(json!({"error": "InternalError"})),
307 )
308 .into_response();
309 }
310
311 info!(did = %auth.0.did, "TOTP enabled with {} backup codes", backup_codes.len());
312
313 Json(EnableTotpResponse { backup_codes }).into_response()
314}
315
316#[derive(Deserialize)]
317pub struct DisableTotpInput {
318 pub password: String,
319 pub code: String,
320}
321
322pub async fn disable_totp(
323 State(state): State<AppState>,
324 auth: BearerAuth,
325 Json(input): Json<DisableTotpInput>,
326) -> Response {
327 if !crate::api::server::reauth::check_legacy_session_mfa(&state.db, &auth.0.did).await {
328 return crate::api::server::reauth::legacy_mfa_required_response(&state.db, &auth.0.did)
329 .await;
330 }
331
332 if !state
333 .check_rate_limit(RateLimitKind::TotpVerify, &auth.0.did)
334 .await
335 {
336 warn!(did = %auth.0.did, "TOTP verification rate limit exceeded");
337 return (
338 StatusCode::TOO_MANY_REQUESTS,
339 Json(json!({
340 "error": "RateLimitExceeded",
341 "message": "Too many verification attempts. Please try again in a few minutes."
342 })),
343 )
344 .into_response();
345 }
346
347 let user = sqlx::query!("SELECT password_hash FROM users WHERE did = $1", auth.0.did)
348 .fetch_optional(&state.db)
349 .await;
350
351 let password_hash = match user {
352 Ok(Some(row)) => row.password_hash,
353 Ok(None) => {
354 return (
355 StatusCode::NOT_FOUND,
356 Json(json!({"error": "AccountNotFound", "message": "Account not found"})),
357 )
358 .into_response();
359 }
360 Err(e) => {
361 error!("DB error fetching user: {:?}", e);
362 return (
363 StatusCode::INTERNAL_SERVER_ERROR,
364 Json(json!({"error": "InternalError"})),
365 )
366 .into_response();
367 }
368 };
369
370 let password_valid = password_hash
371 .as_ref()
372 .map(|h| bcrypt::verify(&input.password, h).unwrap_or(false))
373 .unwrap_or(false);
374 if !password_valid {
375 return (
376 StatusCode::UNAUTHORIZED,
377 Json(json!({
378 "error": "InvalidPassword",
379 "message": "Password is incorrect"
380 })),
381 )
382 .into_response();
383 }
384
385 let totp_row = sqlx::query!(
386 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
387 auth.0.did
388 )
389 .fetch_optional(&state.db)
390 .await;
391
392 let totp_row = match totp_row {
393 Ok(Some(row)) if row.verified => row,
394 Ok(Some(_)) | Ok(None) => {
395 return (
396 StatusCode::BAD_REQUEST,
397 Json(json!({
398 "error": "TotpNotEnabled",
399 "message": "TOTP is not enabled for this account"
400 })),
401 )
402 .into_response();
403 }
404 Err(e) => {
405 error!("DB error fetching TOTP: {:?}", e);
406 return (
407 StatusCode::INTERNAL_SERVER_ERROR,
408 Json(json!({"error": "InternalError"})),
409 )
410 .into_response();
411 }
412 };
413
414 let code = input.code.trim();
415 let code_valid = if is_backup_code_format(code) {
416 verify_backup_code_for_user(&state, &auth.0.did, code).await
417 } else {
418 let secret =
419 match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version) {
420 Ok(s) => s,
421 Err(e) => {
422 error!("Failed to decrypt TOTP secret: {:?}", e);
423 return (
424 StatusCode::INTERNAL_SERVER_ERROR,
425 Json(json!({"error": "InternalError"})),
426 )
427 .into_response();
428 }
429 };
430 verify_totp_code(&secret, code)
431 };
432
433 if !code_valid {
434 return (
435 StatusCode::UNAUTHORIZED,
436 Json(json!({
437 "error": "InvalidCode",
438 "message": "Invalid verification code"
439 })),
440 )
441 .into_response();
442 }
443
444 let mut tx = match state.db.begin().await {
445 Ok(tx) => tx,
446 Err(e) => {
447 error!("Failed to begin transaction: {:?}", e);
448 return (
449 StatusCode::INTERNAL_SERVER_ERROR,
450 Json(json!({"error": "InternalError"})),
451 )
452 .into_response();
453 }
454 };
455
456 if let Err(e) = sqlx::query!("DELETE FROM user_totp WHERE did = $1", auth.0.did)
457 .execute(&mut *tx)
458 .await
459 {
460 error!("Failed to delete TOTP: {:?}", e);
461 return (
462 StatusCode::INTERNAL_SERVER_ERROR,
463 Json(json!({"error": "InternalError"})),
464 )
465 .into_response();
466 }
467
468 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", auth.0.did)
469 .execute(&mut *tx)
470 .await
471 {
472 error!("Failed to delete backup codes: {:?}", e);
473 return (
474 StatusCode::INTERNAL_SERVER_ERROR,
475 Json(json!({"error": "InternalError"})),
476 )
477 .into_response();
478 }
479
480 if let Err(e) = tx.commit().await {
481 error!("Failed to commit transaction: {:?}", e);
482 return (
483 StatusCode::INTERNAL_SERVER_ERROR,
484 Json(json!({"error": "InternalError"})),
485 )
486 .into_response();
487 }
488
489 info!(did = %auth.0.did, "TOTP disabled");
490
491 (StatusCode::OK, Json(json!({}))).into_response()
492}
493
494#[derive(Serialize)]
495#[serde(rename_all = "camelCase")]
496pub struct GetTotpStatusResponse {
497 pub enabled: bool,
498 pub has_backup_codes: bool,
499 pub backup_codes_remaining: i64,
500}
501
502pub async fn get_totp_status(State(state): State<AppState>, auth: BearerAuth) -> Response {
503 let totp_row = sqlx::query!("SELECT verified FROM user_totp WHERE did = $1", auth.0.did)
504 .fetch_optional(&state.db)
505 .await;
506
507 let enabled = match totp_row {
508 Ok(Some(row)) => row.verified,
509 Ok(None) => false,
510 Err(e) => {
511 error!("DB error fetching TOTP status: {:?}", e);
512 return (
513 StatusCode::INTERNAL_SERVER_ERROR,
514 Json(json!({"error": "InternalError"})),
515 )
516 .into_response();
517 }
518 };
519
520 let backup_count_row = sqlx::query!(
521 "SELECT COUNT(*) as count FROM backup_codes WHERE did = $1 AND used_at IS NULL",
522 auth.0.did
523 )
524 .fetch_one(&state.db)
525 .await;
526
527 let backup_count = backup_count_row.map(|r| r.count.unwrap_or(0)).unwrap_or(0);
528
529 Json(GetTotpStatusResponse {
530 enabled,
531 has_backup_codes: backup_count > 0,
532 backup_codes_remaining: backup_count,
533 })
534 .into_response()
535}
536
537#[derive(Deserialize)]
538pub struct RegenerateBackupCodesInput {
539 pub password: String,
540 pub code: String,
541}
542
543#[derive(Serialize)]
544#[serde(rename_all = "camelCase")]
545pub struct RegenerateBackupCodesResponse {
546 pub backup_codes: Vec<String>,
547}
548
549pub async fn regenerate_backup_codes(
550 State(state): State<AppState>,
551 auth: BearerAuth,
552 Json(input): Json<RegenerateBackupCodesInput>,
553) -> Response {
554 if !state
555 .check_rate_limit(RateLimitKind::TotpVerify, &auth.0.did)
556 .await
557 {
558 warn!(did = %auth.0.did, "TOTP verification rate limit exceeded");
559 return (
560 StatusCode::TOO_MANY_REQUESTS,
561 Json(json!({
562 "error": "RateLimitExceeded",
563 "message": "Too many verification attempts. Please try again in a few minutes."
564 })),
565 )
566 .into_response();
567 }
568
569 let user = sqlx::query!("SELECT password_hash FROM users WHERE did = $1", auth.0.did)
570 .fetch_optional(&state.db)
571 .await;
572
573 let password_hash = match user {
574 Ok(Some(row)) => row.password_hash,
575 Ok(None) => {
576 return (
577 StatusCode::NOT_FOUND,
578 Json(json!({"error": "AccountNotFound", "message": "Account not found"})),
579 )
580 .into_response();
581 }
582 Err(e) => {
583 error!("DB error fetching user: {:?}", e);
584 return (
585 StatusCode::INTERNAL_SERVER_ERROR,
586 Json(json!({"error": "InternalError"})),
587 )
588 .into_response();
589 }
590 };
591
592 let password_valid = password_hash
593 .as_ref()
594 .map(|h| bcrypt::verify(&input.password, h).unwrap_or(false))
595 .unwrap_or(false);
596 if !password_valid {
597 return (
598 StatusCode::UNAUTHORIZED,
599 Json(json!({
600 "error": "InvalidPassword",
601 "message": "Password is incorrect"
602 })),
603 )
604 .into_response();
605 }
606
607 let totp_row = sqlx::query!(
608 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
609 auth.0.did
610 )
611 .fetch_optional(&state.db)
612 .await;
613
614 let totp_row = match totp_row {
615 Ok(Some(row)) if row.verified => row,
616 Ok(Some(_)) | Ok(None) => {
617 return (
618 StatusCode::BAD_REQUEST,
619 Json(json!({
620 "error": "TotpNotEnabled",
621 "message": "TOTP must be enabled to regenerate backup codes"
622 })),
623 )
624 .into_response();
625 }
626 Err(e) => {
627 error!("DB error fetching TOTP: {:?}", e);
628 return (
629 StatusCode::INTERNAL_SERVER_ERROR,
630 Json(json!({"error": "InternalError"})),
631 )
632 .into_response();
633 }
634 };
635
636 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
637 {
638 Ok(s) => s,
639 Err(e) => {
640 error!("Failed to decrypt TOTP secret: {:?}", e);
641 return (
642 StatusCode::INTERNAL_SERVER_ERROR,
643 Json(json!({"error": "InternalError"})),
644 )
645 .into_response();
646 }
647 };
648
649 let code = input.code.trim();
650 if !verify_totp_code(&secret, code) {
651 return (
652 StatusCode::UNAUTHORIZED,
653 Json(json!({
654 "error": "InvalidCode",
655 "message": "Invalid verification code"
656 })),
657 )
658 .into_response();
659 }
660
661 let backup_codes = generate_backup_codes();
662 let mut tx = match state.db.begin().await {
663 Ok(tx) => tx,
664 Err(e) => {
665 error!("Failed to begin transaction: {:?}", e);
666 return (
667 StatusCode::INTERNAL_SERVER_ERROR,
668 Json(json!({"error": "InternalError"})),
669 )
670 .into_response();
671 }
672 };
673
674 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", auth.0.did)
675 .execute(&mut *tx)
676 .await
677 {
678 error!("Failed to clear old backup codes: {:?}", e);
679 return (
680 StatusCode::INTERNAL_SERVER_ERROR,
681 Json(json!({"error": "InternalError"})),
682 )
683 .into_response();
684 }
685
686 for code in &backup_codes {
687 let hash = match hash_backup_code(code) {
688 Ok(h) => h,
689 Err(e) => {
690 error!("Failed to hash backup code: {:?}", e);
691 return (
692 StatusCode::INTERNAL_SERVER_ERROR,
693 Json(json!({"error": "InternalError"})),
694 )
695 .into_response();
696 }
697 };
698
699 if let Err(e) = sqlx::query!(
700 "INSERT INTO backup_codes (did, code_hash, created_at) VALUES ($1, $2, NOW())",
701 auth.0.did,
702 hash
703 )
704 .execute(&mut *tx)
705 .await
706 {
707 error!("Failed to store backup code: {:?}", e);
708 return (
709 StatusCode::INTERNAL_SERVER_ERROR,
710 Json(json!({"error": "InternalError"})),
711 )
712 .into_response();
713 }
714 }
715
716 if let Err(e) = tx.commit().await {
717 error!("Failed to commit transaction: {:?}", e);
718 return (
719 StatusCode::INTERNAL_SERVER_ERROR,
720 Json(json!({"error": "InternalError"})),
721 )
722 .into_response();
723 }
724
725 info!(did = %auth.0.did, "Backup codes regenerated");
726
727 Json(RegenerateBackupCodesResponse { backup_codes }).into_response()
728}
729
730async fn verify_backup_code_for_user(state: &AppState, did: &str, code: &str) -> bool {
731 let code = code.trim().to_uppercase();
732
733 let backup_codes = sqlx::query!(
734 "SELECT id, code_hash FROM backup_codes WHERE did = $1 AND used_at IS NULL",
735 did
736 )
737 .fetch_all(&state.db)
738 .await;
739
740 let backup_codes = match backup_codes {
741 Ok(codes) => codes,
742 Err(e) => {
743 warn!("Failed to fetch backup codes: {:?}", e);
744 return false;
745 }
746 };
747
748 for row in backup_codes {
749 if verify_backup_code(&code, &row.code_hash) {
750 let _ = sqlx::query!(
751 "UPDATE backup_codes SET used_at = $1 WHERE id = $2",
752 Utc::now(),
753 row.id
754 )
755 .execute(&state.db)
756 .await;
757 return true;
758 }
759 }
760
761 false
762}
763
764pub async fn verify_totp_or_backup_for_user(state: &AppState, did: &str, code: &str) -> bool {
765 let code = code.trim();
766
767 if is_backup_code_format(code) {
768 return verify_backup_code_for_user(state, did, code).await;
769 }
770
771 let totp_row = sqlx::query!(
772 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
773 did
774 )
775 .fetch_optional(&state.db)
776 .await;
777
778 let totp_row = match totp_row {
779 Ok(Some(row)) if row.verified => row,
780 _ => return false,
781 };
782
783 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
784 {
785 Ok(s) => s,
786 Err(_) => return false,
787 };
788
789 if verify_totp_code(&secret, code) {
790 let _ = sqlx::query!("UPDATE user_totp SET last_used = NOW() WHERE did = $1", did)
791 .execute(&state.db)
792 .await;
793 return true;
794 }
795
796 false
797}
798
799pub async fn has_totp_enabled(state: &AppState, did: &str) -> bool {
800 has_totp_enabled_db(&state.db, did).await
801}
802
803pub async fn has_totp_enabled_db(db: &sqlx::PgPool, did: &str) -> bool {
804 let result = sqlx::query_scalar!("SELECT verified FROM user_totp WHERE did = $1", did)
805 .fetch_optional(db)
806 .await;
807
808 matches!(result, Ok(Some(true)))
809}