use crate::api::auth::{Claims, Session, TokenScope}; use crate::error::{ApiResult, AppError}; use crate::model::auth::AuthResponse; use crate::svc::auth_svc::AuthService; use rocket::serde::json::Json; use rocket::serde::{Deserialize, Serialize}; use rocket::State; use totp_rs::{Algorithm, TOTP}; #[derive(Debug, Deserialize)] pub struct TOTPSixDigitCode { code: String, } #[derive(Debug, sqlx::Type, Clone, Copy, Serialize, Deserialize, PartialEq)] #[serde(rename_all = "snake_case")] #[sqlx(type_name = "totp_status", rename_all = "lowercase")] pub enum TotpStatus { Enabled, Pending, Disabled, } #[derive(Serialize)] pub struct QrResponse { qr_code: String, } #[derive(Deserialize)] pub struct TotpVerifyRequest { pub code: String, } #[derive(Deserialize)] pub struct PasswordConfirmation { password: String, } #[derive(Deserialize)] pub struct PasswordAnd2fa { pub password: String, pub totp_code: String, } pub fn totp_gen(user_id: i64, secret: &[u8]) -> ApiResult { TOTP::new( Algorithm::SHA1, 6, 1, 30, secret.to_owned(), Some("chat.zxq5.dev".to_string()), format!("{}", user_id), ) .map_err(|_| AppError::internal("failed to generate totp")) } #[post("/totp", data = "

")] pub async fn confirm_totp( user: Session, form: Json, svc: &State, ) -> ApiResult<()> { svc.confirm_totp(user.uid, &form.code).await } #[post("/totp.jpg", data = "")] pub async fn get_totp( user: Session, form: Json, svc: &State, ) -> ApiResult> { let secret = svc.get_or_create_totp_secret(user.uid, &form.password).await?; let qr_b64 = totp_gen(user.uid, secret.as_bytes()) .map_err(|_| AppError::internal("invalid totp secret"))? .get_qr_base64() .map_err(|_| AppError::internal("failed to generate qr code"))?; Ok(Json(QrResponse { qr_code: format!("data:image/png;base64,{}", qr_b64), })) } #[get("/totp/status")] pub async fn get_totp_status( user: Session, svc: &State, ) -> ApiResult> { Ok(Json( svc.get_totp_status(user.uid).await? .then_some(TotpStatus::Enabled) .unwrap_or(TotpStatus::Disabled), )) } #[delete("/totp", data = "")] pub async fn disable_totp( user: Session, form: Json, svc: &State, ) -> ApiResult> { let response = svc.disable_totp(user.uid, &form.password, &form.totp_code).await?; Ok(Json(response)) } #[post("/totp/verify", data = "")] pub async fn verify_totp( claims: Claims, body: Json, svc: &State, ) -> ApiResult> { // reject if they somehow got here with a full token if claims.scope != TokenScope::TotpPending { return Err(AppError::Forbidden); } let response = svc.login_totp(claims.sub as i64, &body.code).await?; Ok(Json(response)) }