.

backend/.cargo/config.toml

@@ -1,12 +1,12 @@
-[target.x86_64-unknown-linux-gnu]
+# [target.x86_64-unknown-linux-gnu]
-linker = "clang"
+# linker = "clang"
-rustflags = ["-C", "link-arg=-fuse-ld=/usr/bin/mold", "-Zshare-generics=y"]
+# rustflags = ["-C", "link-arg=-fuse-ld=/usr/bin/mold", "-Zshare-generics=y"]
 
-[unstable]
+# [unstable]
-codegen-backend = true
+# codegen-backend = true
 
-[profile.dev]
+# [profile.dev]
-codegen-backend = "cranelift"
+# codegen-backend = "cranelift"
 
-[profile.dev.package."*"]
+# [profile.dev.package."*"]
-codegen-backend = "llvm"
+# codegen-backend = "llvm"

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<Json<QrResponse>> {
-    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<Self, Self::Error> {
-        let user = try_outcome!(request.guard::<Session>().await);
-        let mut pool = match request.guard::<Connection<DbConn>>().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 = "<totp>")]
+pub async fn confirm_totp(
+    mfa: MultiFactorEnabled,
+    totp: Json<Totp>,
+    mut db: Connection<DbConn>,
+) -> Status {
+    if totp.code.len() == 6
+        && let Ok(code) = totp.code.parse::<usize>()
+    {
+        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<Json<QrResponse>> {
+    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<Self, Self::Error> {
+        let user = try_outcome!(request.guard::<Session>().await);
+        let mut pool = match request.guard::<Connection<DbConn>>().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<DbConn>) -> 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(()),
+        }
+    }
+}