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;
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 let totp_row = sqlx::query!(
153 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
154 auth.0.did
155 )
156 .fetch_optional(&state.db)
157 .await;
158
159 let totp_row = match totp_row {
160 Ok(Some(row)) => row,
161 Ok(None) => {
162 return (
163 StatusCode::BAD_REQUEST,
164 Json(json!({
165 "error": "TotpNotSetup",
166 "message": "Please call createTotpSecret first"
167 })),
168 )
169 .into_response();
170 }
171 Err(e) => {
172 error!("DB error fetching TOTP: {:?}", e);
173 return (
174 StatusCode::INTERNAL_SERVER_ERROR,
175 Json(json!({"error": "InternalError"})),
176 )
177 .into_response();
178 }
179 };
180
181 if totp_row.verified {
182 return (
183 StatusCode::CONFLICT,
184 Json(json!({
185 "error": "TotpAlreadyEnabled",
186 "message": "TOTP is already enabled"
187 })),
188 )
189 .into_response();
190 }
191
192 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
193 {
194 Ok(s) => s,
195 Err(e) => {
196 error!("Failed to decrypt TOTP secret: {:?}", e);
197 return (
198 StatusCode::INTERNAL_SERVER_ERROR,
199 Json(json!({"error": "InternalError"})),
200 )
201 .into_response();
202 }
203 };
204
205 let code = input.code.trim();
206 if !verify_totp_code(&secret, code) {
207 return (
208 StatusCode::UNAUTHORIZED,
209 Json(json!({
210 "error": "InvalidCode",
211 "message": "Invalid verification code"
212 })),
213 )
214 .into_response();
215 }
216
217 let backup_codes = generate_backup_codes();
218 let mut tx = match state.db.begin().await {
219 Ok(tx) => tx,
220 Err(e) => {
221 error!("Failed to begin transaction: {:?}", e);
222 return (
223 StatusCode::INTERNAL_SERVER_ERROR,
224 Json(json!({"error": "InternalError"})),
225 )
226 .into_response();
227 }
228 };
229
230 if let Err(e) = sqlx::query!(
231 "UPDATE user_totp SET verified = true, last_used = NOW() WHERE did = $1",
232 auth.0.did
233 )
234 .execute(&mut *tx)
235 .await
236 {
237 error!("Failed to enable TOTP: {:?}", e);
238 return (
239 StatusCode::INTERNAL_SERVER_ERROR,
240 Json(json!({"error": "InternalError"})),
241 )
242 .into_response();
243 }
244
245 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", auth.0.did)
246 .execute(&mut *tx)
247 .await
248 {
249 error!("Failed to clear old backup codes: {:?}", e);
250 return (
251 StatusCode::INTERNAL_SERVER_ERROR,
252 Json(json!({"error": "InternalError"})),
253 )
254 .into_response();
255 }
256
257 for code in &backup_codes {
258 let hash = match hash_backup_code(code) {
259 Ok(h) => h,
260 Err(e) => {
261 error!("Failed to hash backup code: {:?}", e);
262 return (
263 StatusCode::INTERNAL_SERVER_ERROR,
264 Json(json!({"error": "InternalError"})),
265 )
266 .into_response();
267 }
268 };
269
270 if let Err(e) = sqlx::query!(
271 "INSERT INTO backup_codes (did, code_hash, created_at) VALUES ($1, $2, NOW())",
272 auth.0.did,
273 hash
274 )
275 .execute(&mut *tx)
276 .await
277 {
278 error!("Failed to store backup code: {:?}", e);
279 return (
280 StatusCode::INTERNAL_SERVER_ERROR,
281 Json(json!({"error": "InternalError"})),
282 )
283 .into_response();
284 }
285 }
286
287 if let Err(e) = tx.commit().await {
288 error!("Failed to commit transaction: {:?}", e);
289 return (
290 StatusCode::INTERNAL_SERVER_ERROR,
291 Json(json!({"error": "InternalError"})),
292 )
293 .into_response();
294 }
295
296 info!(did = %auth.0.did, "TOTP enabled with {} backup codes", backup_codes.len());
297
298 Json(EnableTotpResponse { backup_codes }).into_response()
299}
300
301#[derive(Deserialize)]
302pub struct DisableTotpInput {
303 pub password: String,
304 pub code: String,
305}
306
307pub async fn disable_totp(
308 State(state): State<AppState>,
309 auth: BearerAuth,
310 Json(input): Json<DisableTotpInput>,
311) -> Response {
312 let user = sqlx::query!("SELECT password_hash FROM users WHERE did = $1", auth.0.did)
313 .fetch_optional(&state.db)
314 .await;
315
316 let password_hash = match user {
317 Ok(Some(row)) => row.password_hash,
318 Ok(None) => {
319 return (
320 StatusCode::NOT_FOUND,
321 Json(json!({"error": "AccountNotFound", "message": "Account not found"})),
322 )
323 .into_response();
324 }
325 Err(e) => {
326 error!("DB error fetching user: {:?}", e);
327 return (
328 StatusCode::INTERNAL_SERVER_ERROR,
329 Json(json!({"error": "InternalError"})),
330 )
331 .into_response();
332 }
333 };
334
335 let password_valid = password_hash
336 .as_ref()
337 .map(|h| bcrypt::verify(&input.password, h).unwrap_or(false))
338 .unwrap_or(false);
339 if !password_valid {
340 return (
341 StatusCode::UNAUTHORIZED,
342 Json(json!({
343 "error": "InvalidPassword",
344 "message": "Password is incorrect"
345 })),
346 )
347 .into_response();
348 }
349
350 let totp_row = sqlx::query!(
351 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
352 auth.0.did
353 )
354 .fetch_optional(&state.db)
355 .await;
356
357 let totp_row = match totp_row {
358 Ok(Some(row)) if row.verified => row,
359 Ok(Some(_)) | Ok(None) => {
360 return (
361 StatusCode::BAD_REQUEST,
362 Json(json!({
363 "error": "TotpNotEnabled",
364 "message": "TOTP is not enabled for this account"
365 })),
366 )
367 .into_response();
368 }
369 Err(e) => {
370 error!("DB error fetching TOTP: {:?}", e);
371 return (
372 StatusCode::INTERNAL_SERVER_ERROR,
373 Json(json!({"error": "InternalError"})),
374 )
375 .into_response();
376 }
377 };
378
379 let code = input.code.trim();
380 let code_valid = if is_backup_code_format(code) {
381 verify_backup_code_for_user(&state, &auth.0.did, code).await
382 } else {
383 let secret =
384 match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version) {
385 Ok(s) => s,
386 Err(e) => {
387 error!("Failed to decrypt TOTP secret: {:?}", e);
388 return (
389 StatusCode::INTERNAL_SERVER_ERROR,
390 Json(json!({"error": "InternalError"})),
391 )
392 .into_response();
393 }
394 };
395 verify_totp_code(&secret, code)
396 };
397
398 if !code_valid {
399 return (
400 StatusCode::UNAUTHORIZED,
401 Json(json!({
402 "error": "InvalidCode",
403 "message": "Invalid verification code"
404 })),
405 )
406 .into_response();
407 }
408
409 let mut tx = match state.db.begin().await {
410 Ok(tx) => tx,
411 Err(e) => {
412 error!("Failed to begin transaction: {:?}", e);
413 return (
414 StatusCode::INTERNAL_SERVER_ERROR,
415 Json(json!({"error": "InternalError"})),
416 )
417 .into_response();
418 }
419 };
420
421 if let Err(e) = sqlx::query!("DELETE FROM user_totp WHERE did = $1", auth.0.did)
422 .execute(&mut *tx)
423 .await
424 {
425 error!("Failed to delete TOTP: {:?}", e);
426 return (
427 StatusCode::INTERNAL_SERVER_ERROR,
428 Json(json!({"error": "InternalError"})),
429 )
430 .into_response();
431 }
432
433 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", auth.0.did)
434 .execute(&mut *tx)
435 .await
436 {
437 error!("Failed to delete backup codes: {:?}", e);
438 return (
439 StatusCode::INTERNAL_SERVER_ERROR,
440 Json(json!({"error": "InternalError"})),
441 )
442 .into_response();
443 }
444
445 if let Err(e) = tx.commit().await {
446 error!("Failed to commit transaction: {:?}", e);
447 return (
448 StatusCode::INTERNAL_SERVER_ERROR,
449 Json(json!({"error": "InternalError"})),
450 )
451 .into_response();
452 }
453
454 info!(did = %auth.0.did, "TOTP disabled");
455
456 (StatusCode::OK, Json(json!({}))).into_response()
457}
458
459#[derive(Serialize)]
460#[serde(rename_all = "camelCase")]
461pub struct GetTotpStatusResponse {
462 pub enabled: bool,
463 pub has_backup_codes: bool,
464 pub backup_codes_remaining: i64,
465}
466
467pub async fn get_totp_status(State(state): State<AppState>, auth: BearerAuth) -> Response {
468 let totp_row = sqlx::query!("SELECT verified FROM user_totp WHERE did = $1", auth.0.did)
469 .fetch_optional(&state.db)
470 .await;
471
472 let enabled = match totp_row {
473 Ok(Some(row)) => row.verified,
474 Ok(None) => false,
475 Err(e) => {
476 error!("DB error fetching TOTP status: {:?}", e);
477 return (
478 StatusCode::INTERNAL_SERVER_ERROR,
479 Json(json!({"error": "InternalError"})),
480 )
481 .into_response();
482 }
483 };
484
485 let backup_count_row = sqlx::query!(
486 "SELECT COUNT(*) as count FROM backup_codes WHERE did = $1 AND used_at IS NULL",
487 auth.0.did
488 )
489 .fetch_one(&state.db)
490 .await;
491
492 let backup_count = backup_count_row.map(|r| r.count.unwrap_or(0)).unwrap_or(0);
493
494 Json(GetTotpStatusResponse {
495 enabled,
496 has_backup_codes: backup_count > 0,
497 backup_codes_remaining: backup_count,
498 })
499 .into_response()
500}
501
502#[derive(Deserialize)]
503pub struct RegenerateBackupCodesInput {
504 pub password: String,
505 pub code: String,
506}
507
508#[derive(Serialize)]
509#[serde(rename_all = "camelCase")]
510pub struct RegenerateBackupCodesResponse {
511 pub backup_codes: Vec<String>,
512}
513
514pub async fn regenerate_backup_codes(
515 State(state): State<AppState>,
516 auth: BearerAuth,
517 Json(input): Json<RegenerateBackupCodesInput>,
518) -> Response {
519 let user = sqlx::query!("SELECT password_hash FROM users WHERE did = $1", auth.0.did)
520 .fetch_optional(&state.db)
521 .await;
522
523 let password_hash = match user {
524 Ok(Some(row)) => row.password_hash,
525 Ok(None) => {
526 return (
527 StatusCode::NOT_FOUND,
528 Json(json!({"error": "AccountNotFound", "message": "Account not found"})),
529 )
530 .into_response();
531 }
532 Err(e) => {
533 error!("DB error fetching user: {:?}", e);
534 return (
535 StatusCode::INTERNAL_SERVER_ERROR,
536 Json(json!({"error": "InternalError"})),
537 )
538 .into_response();
539 }
540 };
541
542 let password_valid = password_hash
543 .as_ref()
544 .map(|h| bcrypt::verify(&input.password, h).unwrap_or(false))
545 .unwrap_or(false);
546 if !password_valid {
547 return (
548 StatusCode::UNAUTHORIZED,
549 Json(json!({
550 "error": "InvalidPassword",
551 "message": "Password is incorrect"
552 })),
553 )
554 .into_response();
555 }
556
557 let totp_row = sqlx::query!(
558 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
559 auth.0.did
560 )
561 .fetch_optional(&state.db)
562 .await;
563
564 let totp_row = match totp_row {
565 Ok(Some(row)) if row.verified => row,
566 Ok(Some(_)) | Ok(None) => {
567 return (
568 StatusCode::BAD_REQUEST,
569 Json(json!({
570 "error": "TotpNotEnabled",
571 "message": "TOTP must be enabled to regenerate backup codes"
572 })),
573 )
574 .into_response();
575 }
576 Err(e) => {
577 error!("DB error fetching TOTP: {:?}", e);
578 return (
579 StatusCode::INTERNAL_SERVER_ERROR,
580 Json(json!({"error": "InternalError"})),
581 )
582 .into_response();
583 }
584 };
585
586 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
587 {
588 Ok(s) => s,
589 Err(e) => {
590 error!("Failed to decrypt TOTP secret: {:?}", e);
591 return (
592 StatusCode::INTERNAL_SERVER_ERROR,
593 Json(json!({"error": "InternalError"})),
594 )
595 .into_response();
596 }
597 };
598
599 let code = input.code.trim();
600 if !verify_totp_code(&secret, code) {
601 return (
602 StatusCode::UNAUTHORIZED,
603 Json(json!({
604 "error": "InvalidCode",
605 "message": "Invalid verification code"
606 })),
607 )
608 .into_response();
609 }
610
611 let backup_codes = generate_backup_codes();
612 let mut tx = match state.db.begin().await {
613 Ok(tx) => tx,
614 Err(e) => {
615 error!("Failed to begin transaction: {:?}", e);
616 return (
617 StatusCode::INTERNAL_SERVER_ERROR,
618 Json(json!({"error": "InternalError"})),
619 )
620 .into_response();
621 }
622 };
623
624 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", auth.0.did)
625 .execute(&mut *tx)
626 .await
627 {
628 error!("Failed to clear old backup codes: {:?}", e);
629 return (
630 StatusCode::INTERNAL_SERVER_ERROR,
631 Json(json!({"error": "InternalError"})),
632 )
633 .into_response();
634 }
635
636 for code in &backup_codes {
637 let hash = match hash_backup_code(code) {
638 Ok(h) => h,
639 Err(e) => {
640 error!("Failed to hash backup code: {:?}", e);
641 return (
642 StatusCode::INTERNAL_SERVER_ERROR,
643 Json(json!({"error": "InternalError"})),
644 )
645 .into_response();
646 }
647 };
648
649 if let Err(e) = sqlx::query!(
650 "INSERT INTO backup_codes (did, code_hash, created_at) VALUES ($1, $2, NOW())",
651 auth.0.did,
652 hash
653 )
654 .execute(&mut *tx)
655 .await
656 {
657 error!("Failed to store backup code: {:?}", e);
658 return (
659 StatusCode::INTERNAL_SERVER_ERROR,
660 Json(json!({"error": "InternalError"})),
661 )
662 .into_response();
663 }
664 }
665
666 if let Err(e) = tx.commit().await {
667 error!("Failed to commit transaction: {:?}", e);
668 return (
669 StatusCode::INTERNAL_SERVER_ERROR,
670 Json(json!({"error": "InternalError"})),
671 )
672 .into_response();
673 }
674
675 info!(did = %auth.0.did, "Backup codes regenerated");
676
677 Json(RegenerateBackupCodesResponse { backup_codes }).into_response()
678}
679
680async fn verify_backup_code_for_user(state: &AppState, did: &str, code: &str) -> bool {
681 let code = code.trim().to_uppercase();
682
683 let backup_codes = sqlx::query!(
684 "SELECT id, code_hash FROM backup_codes WHERE did = $1 AND used_at IS NULL",
685 did
686 )
687 .fetch_all(&state.db)
688 .await;
689
690 let backup_codes = match backup_codes {
691 Ok(codes) => codes,
692 Err(e) => {
693 warn!("Failed to fetch backup codes: {:?}", e);
694 return false;
695 }
696 };
697
698 for row in backup_codes {
699 if verify_backup_code(&code, &row.code_hash) {
700 let _ = sqlx::query!(
701 "UPDATE backup_codes SET used_at = $1 WHERE id = $2",
702 Utc::now(),
703 row.id
704 )
705 .execute(&state.db)
706 .await;
707 return true;
708 }
709 }
710
711 false
712}
713
714pub async fn verify_totp_or_backup_for_user(state: &AppState, did: &str, code: &str) -> bool {
715 let code = code.trim();
716
717 if is_backup_code_format(code) {
718 return verify_backup_code_for_user(state, did, code).await;
719 }
720
721 let totp_row = sqlx::query!(
722 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
723 did
724 )
725 .fetch_optional(&state.db)
726 .await;
727
728 let totp_row = match totp_row {
729 Ok(Some(row)) if row.verified => row,
730 _ => return false,
731 };
732
733 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
734 {
735 Ok(s) => s,
736 Err(_) => return false,
737 };
738
739 if verify_totp_code(&secret, code) {
740 let _ = sqlx::query!("UPDATE user_totp SET last_used = NOW() WHERE did = $1", did)
741 .execute(&state.db)
742 .await;
743 return true;
744 }
745
746 false
747}
748
749pub async fn has_totp_enabled(state: &AppState, did: &str) -> bool {
750 has_totp_enabled_db(&state.db, did).await
751}
752
753pub async fn has_totp_enabled_db(db: &sqlx::PgPool, did: &str) -> bool {
754 let result = sqlx::query_scalar!("SELECT verified FROM user_totp WHERE did = $1", did)
755 .fetch_optional(db)
756 .await;
757
758 matches!(result, Ok(Some(true)))
759}