120 lines
3.0 KiB
Rust
120 lines
3.0 KiB
Rust
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> {
|
|
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 = "<form>")]
|
|
pub async fn confirm_totp(
|
|
user: Session,
|
|
form: Json<TOTPSixDigitCode>,
|
|
svc: &State<AuthService>,
|
|
) -> ApiResult<()> {
|
|
svc.confirm_totp(user.uid, &form.code).await
|
|
}
|
|
|
|
#[post("/totp.jpg", data = "<form>")]
|
|
pub async fn get_totp(
|
|
user: Session,
|
|
form: Json<PasswordConfirmation>,
|
|
svc: &State<AuthService>,
|
|
) -> ApiResult<Json<QrResponse>> {
|
|
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<AuthService>,
|
|
) -> ApiResult<Json<TotpStatus>> {
|
|
Ok(Json(
|
|
svc.get_totp_status(user.uid).await?
|
|
.then_some(TotpStatus::Enabled)
|
|
.unwrap_or(TotpStatus::Disabled),
|
|
))
|
|
}
|
|
|
|
#[delete("/totp", data = "<form>")]
|
|
pub async fn disable_totp(
|
|
user: Session,
|
|
form: Json<PasswordAnd2fa>,
|
|
svc: &State<AuthService>,
|
|
) -> ApiResult<Json<AuthResponse>> {
|
|
let response = svc.disable_totp(user.uid, &form.password, &form.totp_code).await?;
|
|
Ok(Json(response))
|
|
}
|
|
|
|
#[post("/totp/verify", data = "<body>")]
|
|
pub async fn verify_totp(
|
|
claims: Claims,
|
|
body: Json<TotpVerifyRequest>,
|
|
svc: &State<AuthService>,
|
|
) -> ApiResult<Json<AuthResponse>> {
|
|
// 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))
|
|
} |