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 = bcrypt::verify(&input.password, &password_hash).unwrap_or(false);
336 if !password_valid {
337 return (
338 StatusCode::UNAUTHORIZED,
339 Json(json!({
340 "error": "InvalidPassword",
341 "message": "Password is incorrect"
342 })),
343 )
344 .into_response();
345 }
346
347 let totp_row = sqlx::query!(
348 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
349 auth.0.did
350 )
351 .fetch_optional(&state.db)
352 .await;
353
354 let totp_row = match totp_row {
355 Ok(Some(row)) if row.verified => row,
356 Ok(Some(_)) | Ok(None) => {
357 return (
358 StatusCode::BAD_REQUEST,
359 Json(json!({
360 "error": "TotpNotEnabled",
361 "message": "TOTP is not enabled for this account"
362 })),
363 )
364 .into_response();
365 }
366 Err(e) => {
367 error!("DB error fetching TOTP: {:?}", e);
368 return (
369 StatusCode::INTERNAL_SERVER_ERROR,
370 Json(json!({"error": "InternalError"})),
371 )
372 .into_response();
373 }
374 };
375
376 let code = input.code.trim();
377 let code_valid = if is_backup_code_format(code) {
378 verify_backup_code_for_user(&state, &auth.0.did, code).await
379 } else {
380 let secret =
381 match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version) {
382 Ok(s) => s,
383 Err(e) => {
384 error!("Failed to decrypt TOTP secret: {:?}", e);
385 return (
386 StatusCode::INTERNAL_SERVER_ERROR,
387 Json(json!({"error": "InternalError"})),
388 )
389 .into_response();
390 }
391 };
392 verify_totp_code(&secret, code)
393 };
394
395 if !code_valid {
396 return (
397 StatusCode::UNAUTHORIZED,
398 Json(json!({
399 "error": "InvalidCode",
400 "message": "Invalid verification code"
401 })),
402 )
403 .into_response();
404 }
405
406 let mut tx = match state.db.begin().await {
407 Ok(tx) => tx,
408 Err(e) => {
409 error!("Failed to begin transaction: {:?}", e);
410 return (
411 StatusCode::INTERNAL_SERVER_ERROR,
412 Json(json!({"error": "InternalError"})),
413 )
414 .into_response();
415 }
416 };
417
418 if let Err(e) = sqlx::query!("DELETE FROM user_totp WHERE did = $1", auth.0.did)
419 .execute(&mut *tx)
420 .await
421 {
422 error!("Failed to delete TOTP: {:?}", e);
423 return (
424 StatusCode::INTERNAL_SERVER_ERROR,
425 Json(json!({"error": "InternalError"})),
426 )
427 .into_response();
428 }
429
430 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", auth.0.did)
431 .execute(&mut *tx)
432 .await
433 {
434 error!("Failed to delete backup codes: {:?}", e);
435 return (
436 StatusCode::INTERNAL_SERVER_ERROR,
437 Json(json!({"error": "InternalError"})),
438 )
439 .into_response();
440 }
441
442 if let Err(e) = tx.commit().await {
443 error!("Failed to commit transaction: {:?}", e);
444 return (
445 StatusCode::INTERNAL_SERVER_ERROR,
446 Json(json!({"error": "InternalError"})),
447 )
448 .into_response();
449 }
450
451 info!(did = %auth.0.did, "TOTP disabled");
452
453 (StatusCode::OK, Json(json!({}))).into_response()
454}
455
456#[derive(Serialize)]
457#[serde(rename_all = "camelCase")]
458pub struct GetTotpStatusResponse {
459 pub enabled: bool,
460 pub has_backup_codes: bool,
461 pub backup_codes_remaining: i64,
462}
463
464pub async fn get_totp_status(State(state): State<AppState>, auth: BearerAuth) -> Response {
465 let totp_row = sqlx::query!("SELECT verified FROM user_totp WHERE did = $1", auth.0.did)
466 .fetch_optional(&state.db)
467 .await;
468
469 let enabled = match totp_row {
470 Ok(Some(row)) => row.verified,
471 Ok(None) => false,
472 Err(e) => {
473 error!("DB error fetching TOTP status: {:?}", e);
474 return (
475 StatusCode::INTERNAL_SERVER_ERROR,
476 Json(json!({"error": "InternalError"})),
477 )
478 .into_response();
479 }
480 };
481
482 let backup_count_row = sqlx::query!(
483 "SELECT COUNT(*) as count FROM backup_codes WHERE did = $1 AND used_at IS NULL",
484 auth.0.did
485 )
486 .fetch_one(&state.db)
487 .await;
488
489 let backup_count = backup_count_row.map(|r| r.count.unwrap_or(0)).unwrap_or(0);
490
491 Json(GetTotpStatusResponse {
492 enabled,
493 has_backup_codes: backup_count > 0,
494 backup_codes_remaining: backup_count,
495 })
496 .into_response()
497}
498
499#[derive(Deserialize)]
500pub struct RegenerateBackupCodesInput {
501 pub password: String,
502 pub code: String,
503}
504
505#[derive(Serialize)]
506#[serde(rename_all = "camelCase")]
507pub struct RegenerateBackupCodesResponse {
508 pub backup_codes: Vec<String>,
509}
510
511pub async fn regenerate_backup_codes(
512 State(state): State<AppState>,
513 auth: BearerAuth,
514 Json(input): Json<RegenerateBackupCodesInput>,
515) -> Response {
516 let user = sqlx::query!("SELECT password_hash FROM users WHERE did = $1", auth.0.did)
517 .fetch_optional(&state.db)
518 .await;
519
520 let password_hash = match user {
521 Ok(Some(row)) => row.password_hash,
522 Ok(None) => {
523 return (
524 StatusCode::NOT_FOUND,
525 Json(json!({"error": "AccountNotFound", "message": "Account not found"})),
526 )
527 .into_response();
528 }
529 Err(e) => {
530 error!("DB error fetching user: {:?}", e);
531 return (
532 StatusCode::INTERNAL_SERVER_ERROR,
533 Json(json!({"error": "InternalError"})),
534 )
535 .into_response();
536 }
537 };
538
539 let password_valid = bcrypt::verify(&input.password, &password_hash).unwrap_or(false);
540 if !password_valid {
541 return (
542 StatusCode::UNAUTHORIZED,
543 Json(json!({
544 "error": "InvalidPassword",
545 "message": "Password is incorrect"
546 })),
547 )
548 .into_response();
549 }
550
551 let totp_row = sqlx::query!(
552 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
553 auth.0.did
554 )
555 .fetch_optional(&state.db)
556 .await;
557
558 let totp_row = match totp_row {
559 Ok(Some(row)) if row.verified => row,
560 Ok(Some(_)) | Ok(None) => {
561 return (
562 StatusCode::BAD_REQUEST,
563 Json(json!({
564 "error": "TotpNotEnabled",
565 "message": "TOTP must be enabled to regenerate backup codes"
566 })),
567 )
568 .into_response();
569 }
570 Err(e) => {
571 error!("DB error fetching TOTP: {:?}", e);
572 return (
573 StatusCode::INTERNAL_SERVER_ERROR,
574 Json(json!({"error": "InternalError"})),
575 )
576 .into_response();
577 }
578 };
579
580 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
581 {
582 Ok(s) => s,
583 Err(e) => {
584 error!("Failed to decrypt TOTP secret: {:?}", e);
585 return (
586 StatusCode::INTERNAL_SERVER_ERROR,
587 Json(json!({"error": "InternalError"})),
588 )
589 .into_response();
590 }
591 };
592
593 let code = input.code.trim();
594 if !verify_totp_code(&secret, code) {
595 return (
596 StatusCode::UNAUTHORIZED,
597 Json(json!({
598 "error": "InvalidCode",
599 "message": "Invalid verification code"
600 })),
601 )
602 .into_response();
603 }
604
605 let backup_codes = generate_backup_codes();
606 let mut tx = match state.db.begin().await {
607 Ok(tx) => tx,
608 Err(e) => {
609 error!("Failed to begin transaction: {:?}", e);
610 return (
611 StatusCode::INTERNAL_SERVER_ERROR,
612 Json(json!({"error": "InternalError"})),
613 )
614 .into_response();
615 }
616 };
617
618 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", auth.0.did)
619 .execute(&mut *tx)
620 .await
621 {
622 error!("Failed to clear old backup codes: {:?}", e);
623 return (
624 StatusCode::INTERNAL_SERVER_ERROR,
625 Json(json!({"error": "InternalError"})),
626 )
627 .into_response();
628 }
629
630 for code in &backup_codes {
631 let hash = match hash_backup_code(code) {
632 Ok(h) => h,
633 Err(e) => {
634 error!("Failed to hash backup code: {:?}", e);
635 return (
636 StatusCode::INTERNAL_SERVER_ERROR,
637 Json(json!({"error": "InternalError"})),
638 )
639 .into_response();
640 }
641 };
642
643 if let Err(e) = sqlx::query!(
644 "INSERT INTO backup_codes (did, code_hash, created_at) VALUES ($1, $2, NOW())",
645 auth.0.did,
646 hash
647 )
648 .execute(&mut *tx)
649 .await
650 {
651 error!("Failed to store backup code: {:?}", e);
652 return (
653 StatusCode::INTERNAL_SERVER_ERROR,
654 Json(json!({"error": "InternalError"})),
655 )
656 .into_response();
657 }
658 }
659
660 if let Err(e) = tx.commit().await {
661 error!("Failed to commit transaction: {:?}", e);
662 return (
663 StatusCode::INTERNAL_SERVER_ERROR,
664 Json(json!({"error": "InternalError"})),
665 )
666 .into_response();
667 }
668
669 info!(did = %auth.0.did, "Backup codes regenerated");
670
671 Json(RegenerateBackupCodesResponse { backup_codes }).into_response()
672}
673
674async fn verify_backup_code_for_user(state: &AppState, did: &str, code: &str) -> bool {
675 let code = code.trim().to_uppercase();
676
677 let backup_codes = sqlx::query!(
678 "SELECT id, code_hash FROM backup_codes WHERE did = $1 AND used_at IS NULL",
679 did
680 )
681 .fetch_all(&state.db)
682 .await;
683
684 let backup_codes = match backup_codes {
685 Ok(codes) => codes,
686 Err(e) => {
687 warn!("Failed to fetch backup codes: {:?}", e);
688 return false;
689 }
690 };
691
692 for row in backup_codes {
693 if verify_backup_code(&code, &row.code_hash) {
694 let _ = sqlx::query!(
695 "UPDATE backup_codes SET used_at = $1 WHERE id = $2",
696 Utc::now(),
697 row.id
698 )
699 .execute(&state.db)
700 .await;
701 return true;
702 }
703 }
704
705 false
706}
707
708pub async fn verify_totp_or_backup_for_user(state: &AppState, did: &str, code: &str) -> bool {
709 let code = code.trim();
710
711 if is_backup_code_format(code) {
712 return verify_backup_code_for_user(state, did, code).await;
713 }
714
715 let totp_row = sqlx::query!(
716 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
717 did
718 )
719 .fetch_optional(&state.db)
720 .await;
721
722 let totp_row = match totp_row {
723 Ok(Some(row)) if row.verified => row,
724 _ => return false,
725 };
726
727 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
728 {
729 Ok(s) => s,
730 Err(_) => return false,
731 };
732
733 if verify_totp_code(&secret, code) {
734 let _ = sqlx::query!("UPDATE user_totp SET last_used = NOW() WHERE did = $1", did)
735 .execute(&state.db)
736 .await;
737 return true;
738 }
739
740 false
741}
742
743pub async fn has_totp_enabled(state: &AppState, did: &str) -> bool {
744 let result = sqlx::query_scalar!("SELECT verified FROM user_totp WHERE did = $1", did)
745 .fetch_optional(&state.db)
746 .await;
747
748 matches!(result, Ok(Some(true)))
749}