this repo has no description
lewis.moe
1use crate::api::EmptyResponse;
2use crate::api::error::ApiError;
3use crate::auth::BearerAuth;
4use crate::auth::totp::{
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 for code in &backup_codes {
199 let hash = match hash_backup_code(code) {
200 Ok(h) => h,
201 Err(e) => {
202 error!("Failed to hash backup code: {:?}", e);
203 return ApiError::InternalError(None).into_response();
204 }
205 };
206
207 if let Err(e) = sqlx::query!(
208 "INSERT INTO backup_codes (did, code_hash, created_at) VALUES ($1, $2, NOW())",
209 &auth.0.did,
210 hash
211 )
212 .execute(&mut *tx)
213 .await
214 {
215 error!("Failed to store backup code: {:?}", e);
216 return ApiError::InternalError(None).into_response();
217 }
218 }
219
220 if let Err(e) = tx.commit().await {
221 error!("Failed to commit transaction: {:?}", e);
222 return ApiError::InternalError(None).into_response();
223 }
224
225 info!(did = %&auth.0.did, "TOTP enabled with {} backup codes", backup_codes.len());
226
227 Json(EnableTotpResponse { backup_codes }).into_response()
228}
229
230#[derive(Deserialize)]
231pub struct DisableTotpInput {
232 pub password: PlainPassword,
233 pub code: String,
234}
235
236pub async fn disable_totp(
237 State(state): State<AppState>,
238 auth: BearerAuth,
239 Json(input): Json<DisableTotpInput>,
240) -> Response {
241 if !crate::api::server::reauth::check_legacy_session_mfa(&state.db, &auth.0.did).await {
242 return crate::api::server::reauth::legacy_mfa_required_response(&state.db, &auth.0.did)
243 .await;
244 }
245
246 if !state
247 .check_rate_limit(RateLimitKind::TotpVerify, &auth.0.did)
248 .await
249 {
250 warn!(did = %&auth.0.did, "TOTP verification rate limit exceeded");
251 return ApiError::RateLimitExceeded(None).into_response();
252 }
253
254 let user = sqlx::query!(
255 "SELECT password_hash FROM users WHERE did = $1",
256 &*&auth.0.did
257 )
258 .fetch_optional(&state.db)
259 .await;
260
261 let password_hash = match user {
262 Ok(Some(row)) => row.password_hash,
263 Ok(None) => return ApiError::AccountNotFound.into_response(),
264 Err(e) => {
265 error!("DB error fetching user: {:?}", e);
266 return ApiError::InternalError(None).into_response();
267 }
268 };
269
270 let password_valid = password_hash
271 .as_ref()
272 .map(|h| bcrypt::verify(&input.password, h).unwrap_or(false))
273 .unwrap_or(false);
274 if !password_valid {
275 return ApiError::InvalidPassword("Password is incorrect".into()).into_response();
276 }
277
278 let totp_row = sqlx::query!(
279 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
280 &auth.0.did
281 )
282 .fetch_optional(&state.db)
283 .await;
284
285 let totp_row = match totp_row {
286 Ok(Some(row)) if row.verified => row,
287 Ok(Some(_)) | Ok(None) => return ApiError::TotpNotEnabled.into_response(),
288 Err(e) => {
289 error!("DB error fetching TOTP: {:?}", e);
290 return ApiError::InternalError(None).into_response();
291 }
292 };
293
294 let code = input.code.trim();
295 let code_valid = if is_backup_code_format(code) {
296 verify_backup_code_for_user(&state, &auth.0.did, code).await
297 } else {
298 let secret =
299 match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version) {
300 Ok(s) => s,
301 Err(e) => {
302 error!("Failed to decrypt TOTP secret: {:?}", e);
303 return ApiError::InternalError(None).into_response();
304 }
305 };
306 verify_totp_code(&secret, code)
307 };
308
309 if !code_valid {
310 return ApiError::InvalidCode(Some("Invalid verification code".into())).into_response();
311 }
312
313 let mut tx = match state.db.begin().await {
314 Ok(tx) => tx,
315 Err(e) => {
316 error!("Failed to begin transaction: {:?}", e);
317 return ApiError::InternalError(None).into_response();
318 }
319 };
320
321 if let Err(e) = sqlx::query!("DELETE FROM user_totp WHERE did = $1", &*&auth.0.did)
322 .execute(&mut *tx)
323 .await
324 {
325 error!("Failed to delete TOTP: {:?}", e);
326 return ApiError::InternalError(None).into_response();
327 }
328
329 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", &*&auth.0.did)
330 .execute(&mut *tx)
331 .await
332 {
333 error!("Failed to delete backup codes: {:?}", e);
334 return ApiError::InternalError(None).into_response();
335 }
336
337 if let Err(e) = tx.commit().await {
338 error!("Failed to commit transaction: {:?}", e);
339 return ApiError::InternalError(None).into_response();
340 }
341
342 info!(did = %&auth.0.did, "TOTP disabled");
343
344 EmptyResponse::ok().into_response()
345}
346
347#[derive(Serialize)]
348#[serde(rename_all = "camelCase")]
349pub struct GetTotpStatusResponse {
350 pub enabled: bool,
351 pub has_backup_codes: bool,
352 pub backup_codes_remaining: i64,
353}
354
355pub async fn get_totp_status(State(state): State<AppState>, auth: BearerAuth) -> Response {
356 let totp_row = sqlx::query!(
357 "SELECT verified FROM user_totp WHERE did = $1",
358 &*&auth.0.did
359 )
360 .fetch_optional(&state.db)
361 .await;
362
363 let enabled = match totp_row {
364 Ok(Some(row)) => row.verified,
365 Ok(None) => false,
366 Err(e) => {
367 error!("DB error fetching TOTP status: {:?}", e);
368 return ApiError::InternalError(None).into_response();
369 }
370 };
371
372 let backup_count_row = sqlx::query!(
373 "SELECT COUNT(*) as count FROM backup_codes WHERE did = $1 AND used_at IS NULL",
374 &auth.0.did
375 )
376 .fetch_one(&state.db)
377 .await;
378
379 let backup_count = backup_count_row.map(|r| r.count.unwrap_or(0)).unwrap_or(0);
380
381 Json(GetTotpStatusResponse {
382 enabled,
383 has_backup_codes: backup_count > 0,
384 backup_codes_remaining: backup_count,
385 })
386 .into_response()
387}
388
389#[derive(Deserialize)]
390pub struct RegenerateBackupCodesInput {
391 pub password: PlainPassword,
392 pub code: String,
393}
394
395#[derive(Serialize)]
396#[serde(rename_all = "camelCase")]
397pub struct RegenerateBackupCodesResponse {
398 pub backup_codes: Vec<String>,
399}
400
401pub async fn regenerate_backup_codes(
402 State(state): State<AppState>,
403 auth: BearerAuth,
404 Json(input): Json<RegenerateBackupCodesInput>,
405) -> Response {
406 if !state
407 .check_rate_limit(RateLimitKind::TotpVerify, &auth.0.did)
408 .await
409 {
410 warn!(did = %&auth.0.did, "TOTP verification rate limit exceeded");
411 return ApiError::RateLimitExceeded(None).into_response();
412 }
413
414 let user = sqlx::query!(
415 "SELECT password_hash FROM users WHERE did = $1",
416 &*&auth.0.did
417 )
418 .fetch_optional(&state.db)
419 .await;
420
421 let password_hash = match user {
422 Ok(Some(row)) => row.password_hash,
423 Ok(None) => return ApiError::AccountNotFound.into_response(),
424 Err(e) => {
425 error!("DB error fetching user: {:?}", e);
426 return ApiError::InternalError(None).into_response();
427 }
428 };
429
430 let password_valid = password_hash
431 .as_ref()
432 .map(|h| bcrypt::verify(&input.password, h).unwrap_or(false))
433 .unwrap_or(false);
434 if !password_valid {
435 return ApiError::InvalidPassword("Password is incorrect".into()).into_response();
436 }
437
438 let totp_row = sqlx::query!(
439 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
440 &auth.0.did
441 )
442 .fetch_optional(&state.db)
443 .await;
444
445 let totp_row = match totp_row {
446 Ok(Some(row)) if row.verified => row,
447 Ok(Some(_)) | Ok(None) => return ApiError::TotpNotEnabled.into_response(),
448 Err(e) => {
449 error!("DB error fetching TOTP: {:?}", e);
450 return ApiError::InternalError(None).into_response();
451 }
452 };
453
454 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
455 {
456 Ok(s) => s,
457 Err(e) => {
458 error!("Failed to decrypt TOTP secret: {:?}", e);
459 return ApiError::InternalError(None).into_response();
460 }
461 };
462
463 let code = input.code.trim();
464 if !verify_totp_code(&secret, code) {
465 return ApiError::InvalidCode(Some("Invalid verification code".into())).into_response();
466 }
467
468 let backup_codes = generate_backup_codes();
469 let mut tx = match state.db.begin().await {
470 Ok(tx) => tx,
471 Err(e) => {
472 error!("Failed to begin transaction: {:?}", e);
473 return ApiError::InternalError(None).into_response();
474 }
475 };
476
477 if let Err(e) = sqlx::query!("DELETE FROM backup_codes WHERE did = $1", &*&auth.0.did)
478 .execute(&mut *tx)
479 .await
480 {
481 error!("Failed to clear old backup codes: {:?}", e);
482 return ApiError::InternalError(None).into_response();
483 }
484
485 for code in &backup_codes {
486 let hash = match hash_backup_code(code) {
487 Ok(h) => h,
488 Err(e) => {
489 error!("Failed to hash backup code: {:?}", e);
490 return ApiError::InternalError(None).into_response();
491 }
492 };
493
494 if let Err(e) = sqlx::query!(
495 "INSERT INTO backup_codes (did, code_hash, created_at) VALUES ($1, $2, NOW())",
496 &auth.0.did,
497 hash
498 )
499 .execute(&mut *tx)
500 .await
501 {
502 error!("Failed to store backup code: {:?}", e);
503 return ApiError::InternalError(None).into_response();
504 }
505 }
506
507 if let Err(e) = tx.commit().await {
508 error!("Failed to commit transaction: {:?}", e);
509 return ApiError::InternalError(None).into_response();
510 }
511
512 info!(did = %&auth.0.did, "Backup codes regenerated");
513
514 Json(RegenerateBackupCodesResponse { backup_codes }).into_response()
515}
516
517async fn verify_backup_code_for_user(state: &AppState, did: &str, code: &str) -> bool {
518 let code = code.trim().to_uppercase();
519
520 let backup_codes = sqlx::query!(
521 "SELECT id, code_hash FROM backup_codes WHERE did = $1 AND used_at IS NULL",
522 did
523 )
524 .fetch_all(&state.db)
525 .await;
526
527 let backup_codes = match backup_codes {
528 Ok(codes) => codes,
529 Err(e) => {
530 warn!("Failed to fetch backup codes: {:?}", e);
531 return false;
532 }
533 };
534
535 for row in backup_codes {
536 if verify_backup_code(&code, &row.code_hash) {
537 let _ = sqlx::query!(
538 "UPDATE backup_codes SET used_at = $1 WHERE id = $2",
539 Utc::now(),
540 row.id
541 )
542 .execute(&state.db)
543 .await;
544 return true;
545 }
546 }
547
548 false
549}
550
551pub async fn verify_totp_or_backup_for_user(state: &AppState, did: &str, code: &str) -> bool {
552 let code = code.trim();
553
554 if is_backup_code_format(code) {
555 return verify_backup_code_for_user(state, did, code).await;
556 }
557
558 let totp_row = sqlx::query!(
559 "SELECT secret_encrypted, encryption_version, verified FROM user_totp WHERE did = $1",
560 did
561 )
562 .fetch_optional(&state.db)
563 .await;
564
565 let totp_row = match totp_row {
566 Ok(Some(row)) if row.verified => row,
567 _ => return false,
568 };
569
570 let secret = match decrypt_totp_secret(&totp_row.secret_encrypted, totp_row.encryption_version)
571 {
572 Ok(s) => s,
573 Err(_) => return false,
574 };
575
576 if verify_totp_code(&secret, code) {
577 let _ = sqlx::query!("UPDATE user_totp SET last_used = NOW() WHERE did = $1", did)
578 .execute(&state.db)
579 .await;
580 return true;
581 }
582
583 false
584}
585
586pub async fn has_totp_enabled(state: &AppState, did: &str) -> bool {
587 has_totp_enabled_db(&state.db, did).await
588}
589
590pub async fn has_totp_enabled_db(db: &sqlx::PgPool, did: &str) -> bool {
591 let result = sqlx::query_scalar!("SELECT verified FROM user_totp WHERE did = $1", did)
592 .fetch_optional(db)
593 .await;
594
595 matches!(result, Ok(Some(true)))
596}