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