From f6bd540c99ed33d3d1b309729f7004f10a8afa4d Mon Sep 17 00:00:00 2001 From: zxq5 Date: Fri, 10 Oct 2025 15:59:38 +0100 Subject: [PATCH] . --- backend/.cargo/config.toml | 18 ++-- backend/src/auth.rs | 206 +++++++++++++++++++++++-------------- 2 files changed, 135 insertions(+), 89 deletions(-) diff --git a/backend/.cargo/config.toml b/backend/.cargo/config.toml index 97f634a..0f0ad84 100644 --- a/backend/.cargo/config.toml +++ b/backend/.cargo/config.toml @@ -1,12 +1,12 @@ -[target.x86_64-unknown-linux-gnu] -linker = "clang" -rustflags = ["-C", "link-arg=-fuse-ld=/usr/bin/mold", "-Zshare-generics=y"] +# [target.x86_64-unknown-linux-gnu] +# linker = "clang" +# rustflags = ["-C", "link-arg=-fuse-ld=/usr/bin/mold", "-Zshare-generics=y"] -[unstable] -codegen-backend = true +# [unstable] +# codegen-backend = true -[profile.dev] -codegen-backend = "cranelift" +# [profile.dev] +# codegen-backend = "cranelift" -[profile.dev.package."*"] -codegen-backend = "llvm" +# [profile.dev.package."*"] +# codegen-backend = "llvm" diff --git a/backend/src/auth.rs b/backend/src/auth.rs index 89465f5..da94ee1 100644 --- a/backend/src/auth.rs +++ b/backend/src/auth.rs @@ -7,7 +7,7 @@ use rocket::{ outcome::{Outcome, try_outcome}, post, request::{self, FromRequest}, - response::Redirect, + response::{Redirect, status}, serde::json::Json, }; use rocket_db_pools::{ @@ -93,35 +93,6 @@ pub async fn login( Err(Status::Unauthorized) } -#[get("/totp")] -pub async fn mfa_page(_session: Session) -> Template { - Template::render("2fa", context!()) -} - -#[derive(Serialize)] -pub struct QrResponse { - qr_code: String, -} - -#[get("/totp.jpg")] -pub async fn get_totp(totp: TOTPCode) -> Option> { - let totp = TOTP::new( - Algorithm::SHA1, - 6, - 1, - 30, - totp.secret.as_bytes().into(), - Some("chat.zxq5.dev".to_string()), - format!("{}", totp.user_id), - ) - .unwrap(); - - let qr = totp.get_qr_base64().unwrap(); - let data_uri = format!("data:image/png;base64,{}", qr); - - Some(Json(QrResponse { qr_code: data_uri })) -} - #[derive(Debug, Clone)] pub struct Session { pub token: String, @@ -151,56 +122,6 @@ impl Session { } } -pub struct TOTPCode { - user_id: usize, - secret: String, -} - -#[rocket::async_trait] -impl<'r> FromRequest<'r> for TOTPCode { - type Error = (); - - async fn from_request(request: &'r Request<'_>) -> request::Outcome { - let user = try_outcome!(request.guard::().await); - let mut pool = match request.guard::>().await { - Outcome::Success(pool) => pool, - _ => return Outcome::Error((Status::Unauthorized, ())), - }; - - let (enabled, mut secret) = match sqlx::query!( - "SELECT twofa_enabled, totp_secret FROM users WHERE id = $1", - user.user_id as i32, - ) - .fetch_one(&mut **pool) - .await - { - Ok(row) => (row.twofa_enabled, row.totp_secret), - Err(_) => return Outcome::Error((Status::Unauthorized, ())), - }; - - if !enabled || secret.is_none() { - secret = Some(Secret::generate_secret().to_string()); - - match sqlx::query!( - "UPDATE users SET totp_secret = $1, twofa_enabled = true WHERE id = $2", - secret.as_ref().unwrap(), - user.user_id as i32, - ) - .execute(&mut **pool) - .await - { - Ok(_) => (), - Err(_) => return Outcome::Error((Status::InternalServerError, ())), - } - } - - Outcome::Success(TOTPCode { - user_id: user.user_id, - secret: secret.unwrap(), - }) - } -} - #[rocket::async_trait] impl<'r> FromRequest<'r> for Session { type Error = (); @@ -234,3 +155,128 @@ impl<'r> FromRequest<'r> for Session { } } } + +// --------------------- TOTP 2FA Auth ----------------------- + +#[get("/totp")] +pub async fn mfa_page(_session: Session) -> Template { + Template::render("2fa", context!()) +} + +#[derive(Debug, Deserialize)] +pub struct Totp { + code: String, +} + +#[post("/totp", data = "")] +pub async fn confirm_totp( + mfa: MultiFactorEnabled, + totp: Json, + mut db: Connection, +) -> Status { + if totp.code.len() == 6 + && let Ok(code) = totp.code.parse::() + { + let secret = match sqlx::query!( + "SELECT totp_secret FROM users WHERE id = $1", + mfa.user_id as i32 + ) + .fetch_one(&mut **db) + .await + { + Err(_) => return Status::InternalServerError, + Ok(user) => user.totp_secret, + }; + } + + return Status::BadRequest; +} + +#[derive(Serialize)] +pub struct QrResponse { + qr_code: String, +} + +#[get("/totp.jpg")] +pub async fn get_totp(mfa: MultiFactorEnabled) -> Option> { + let totp = TOTP::new( + Algorithm::SHA1, + 6, + 1, + 30, + mfa.secret.as_bytes().into(), + Some("chat.zxq5.dev".to_string()), + format!("{}", mfa.user_id), + ) + .unwrap(); + + let qr = totp.get_qr_base64().unwrap(); + let data_uri = format!("data:image/png;base64,{}", qr); + + Some(Json(QrResponse { qr_code: data_uri })) +} + +pub struct MultiFactorEnabled { + user_id: usize, + secret: String, +} + +#[rocket::async_trait] +impl<'r> FromRequest<'r> for MultiFactorEnabled { + type Error = (); + + async fn from_request(request: &'r Request<'_>) -> request::Outcome { + let user = try_outcome!(request.guard::().await); + let mut pool = match request.guard::>().await { + Outcome::Success(pool) => pool, + _ => return Outcome::Error((Status::Unauthorized, ())), + }; + + let (enabled, mut secret) = match sqlx::query!( + "SELECT twofa_enabled, totp_secret FROM users WHERE id = $1", + user.user_id as i32, + ) + .fetch_one(&mut **pool) + .await + { + Ok(row) => (row.twofa_enabled, row.totp_secret), + Err(_) => return Outcome::Error((Status::Unauthorized, ())), + }; + + if !enabled || secret.is_none() { + secret = Some(Secret::generate_secret().to_string()); + + match sqlx::query!( + "UPDATE users SET totp_secret = $1 WHERE id = $2", + secret.as_ref().unwrap(), + user.user_id as i32, + ) + .execute(&mut **pool) + .await + { + Ok(_) => (), + Err(_) => return Outcome::Error((Status::InternalServerError, ())), + } + } + + Outcome::Success(MultiFactorEnabled { + user_id: user.user_id, + secret: secret.unwrap(), + }) + } +} + +impl MultiFactorEnabled { + pub async fn enable(&self, db: &mut Connection) -> Result<(), ()> { + match sqlx::query!( + "UPDATE users SET twofa_enabled = true WHERE id = $1", + self.user_id as i32, + ) + .execute(&mut ***db) + .await + { + Ok(_) => Ok(()), + Err(_) => Err(()), + } + } +}