From 4a6c3bc49ced71ffff1fe5c4189111bb0835ce4c Mon Sep 17 00:00:00 2001 From: zxq5 Date: Fri, 10 Oct 2025 01:45:02 +0100 Subject: [PATCH] more progress on TOTP/2FA --- backend/Rocket.toml | 4 +- backend/migrations/20251009091103_auth-v2.sql | 18 +-- backend/migrations/20251010002253_auth-v3.sql | 2 + backend/src/auth.rs | 106 ++++++++++++------ backend/src/llm.rs | 4 +- backend/src/main.rs | 17 +-- backend/src/messages.rs | 52 ++++++--- backend/static/css/index.css | 8 +- backend/templates/2fa.html.tera | 68 +++++------ backend/templates/chat.html.tera | 99 ++++------------ backend/templates/login.html.tera | 4 +- backend/templates/signup.html.tera | 4 +- 12 files changed, 189 insertions(+), 197 deletions(-) create mode 100644 backend/migrations/20251010002253_auth-v3.sql diff --git a/backend/Rocket.toml b/backend/Rocket.toml index 6b21888..66a779b 100644 --- a/backend/Rocket.toml +++ b/backend/Rocket.toml @@ -1,6 +1,6 @@ [debug] secret_key = "yYhvCGnRh/TrcHtB8sZqCFifrVmJxoKFLBYw/WWBZeU=" -address = "127.0.0.1" +address = "0.0.0.0" port = 8000 [default.databases.postgres_db] @@ -8,4 +8,4 @@ url = "postgresql://chatapp:chatapp@100.118.108.58:5432/chatapp" [default] # run inside a docker container or pod address = "0.0.0.0" -port = 8082 +port = 8000 diff --git a/backend/migrations/20251009091103_auth-v2.sql b/backend/migrations/20251009091103_auth-v2.sql index 69cb4e9..5174a46 100644 --- a/backend/migrations/20251009091103_auth-v2.sql +++ b/backend/migrations/20251009091103_auth-v2.sql @@ -1,15 +1,5 @@ -- Add migration script here -DROP TABLE users; - -CREATE TABLE users ( - id SERIAL PRIMARY KEY, - username VARCHAR(50) UNIQUE NOT NULL, - display_name VARCHAR(50), - - email VARCHAR(100) NOT NULL, - password VARCHAR(100) NOT NULL, - totp_secret VARCHAR(100), - - created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP, - updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP -); +ALTER TABLE users ADD COLUMN email VARCHAR(100); +ALTER TABLE users ADD COLUMN twofa_enabled BOOLEAN DEFAULT FALSE; +ALTER TABLE users ADD COLUMN totp_secret VARCHAR(64); +ALTER TABLE users ADD COLUMN updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP; diff --git a/backend/migrations/20251010002253_auth-v3.sql b/backend/migrations/20251010002253_auth-v3.sql new file mode 100644 index 0000000..a5831ac --- /dev/null +++ b/backend/migrations/20251010002253_auth-v3.sql @@ -0,0 +1,2 @@ +-- Add migration script here +ALTER TABLE users ALTER COLUMN twofa_enabled SET NOT NULL; diff --git a/backend/src/auth.rs b/backend/src/auth.rs index 65219f0..48e9372 100644 --- a/backend/src/auth.rs +++ b/backend/src/auth.rs @@ -5,9 +5,10 @@ use rocket::{ Request, fs::NamedFile, http::{CookieJar, Status}, - outcome::Outcome, + outcome::{Outcome, try_outcome}, post, request::{self, FromRequest}, + response::Redirect, serde::json::Json, }; use rocket_db_pools::{ @@ -38,7 +39,7 @@ pub async fn signup( cred: Json, jar: &CookieJar<'_>, mut db: Connection, -) -> Result, String> { +) -> Result { let result = sqlx::query!( "INSERT INTO users (username, password) VALUES ($1, $2) RETURNING id", cred.username, @@ -46,18 +47,18 @@ pub async fn signup( ) .fetch_one(&mut **db) .await - .map_err(|e| e.to_string())?; + .map_err(|e| Status::InternalServerError)?; let session = Session::new(result.id as usize); if let Err(e) = session.commit(&mut db).await { eprintln!("Failed to create session: {}", e); - return Err(e.to_string()); + return Err(Status::InternalServerError); } jar.add_private(("session", session.token)); println!("Signup successful"); - Ok(Json("Signup successful".to_string())) + return Ok(Redirect::to("/chat")); } #[get("/login")] @@ -70,7 +71,7 @@ pub async fn login( mut db: Connection, jar: &CookieJar<'_>, cred: Json, -) -> Result, String> { +) -> Result { if let Ok(row) = sqlx::query!( "SELECT id FROM users WHERE username = $1 AND password = $2", cred.username, @@ -82,41 +83,47 @@ pub async fn login( let session = Session::new(row.id as usize); if let Err(e) = session.commit(&mut db).await { eprintln!("Failed to create session: {}", e); - return Err(e.to_string()); + return Err(Status::InternalServerError); } jar.add_private(("session", session.token)); - return Ok(Json("Signup successful".to_string())); + return Ok(Redirect::to("/chat")); } // TODO: implement actual login logic, e.g. verify password and generate token - Err("login failed".to_string()) + Err(Status::Unauthorized) } #[get("/totp")] -pub async fn mfa_page(session: Session) -> Template { +pub async fn mfa_page(_session: Session) -> Template { Template::render("2fa", context!()) } -#[get("/api/totp.jpg")] -pub async fn get_totp(s: Session) -> Option { +#[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, - Secret::generate_secret().to_bytes().unwrap(), - Some("Github".to_string()), - format!("{}", s.user_id), + 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(QrCodeImage(qr.into())) + Some(Json(QrResponse { qr_code: data_uri })) } -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct Session { pub token: String, pub user_id: usize, @@ -145,6 +152,56 @@ 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 = (); @@ -178,18 +235,3 @@ impl<'r> FromRequest<'r> for Session { } } } - -use rocket::http::ContentType; -use rocket::response::{self, Responder, Response}; -use std::io::Cursor; - -pub struct QrCodeImage(Vec); - -impl<'r> Responder<'r, 'static> for QrCodeImage { - fn respond_to(self, _: &'r rocket::Request<'_>) -> response::Result<'static> { - Response::build() - .header(ContentType::PNG) - .sized_body(self.0.len(), Cursor::new(self.0)) - .ok() - } -} diff --git a/backend/src/llm.rs b/backend/src/llm.rs index 7fda239..f1aa988 100644 --- a/backend/src/llm.rs +++ b/backend/src/llm.rs @@ -60,10 +60,10 @@ impl LlmWorker { .map_err(|_| String::from("Failed to make request to LLM server"))?; Ok(ChatMsg { - display_name: message.display_name.clone(), + display_name: Some(String::from("lmstudio")), user_id: 0, text: llm_resp.choices[0].message.content.clone(), - timestamp: chrono::Utc::now().timestamp() as usize, + timestamp: chrono::Utc::now().timestamp_millis() as usize, }) } } diff --git a/backend/src/main.rs b/backend/src/main.rs index 19bc680..2f400c8 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -34,12 +34,15 @@ async fn users(_ag: Session, mut db: Connection) -> Json> { } #[get("/users/", rank = 1)] -async fn username_for_id(id: usize, _ag: Session, mut db: Connection) -> String { - sqlx::query!("SELECT username FROM users WHERE id = $1", id as i32) - .fetch_one(&mut **db) - .await - .map(|row| row.username) - .unwrap_or_else(|_| "User not found".to_string()) +async fn display_name(id: usize, _ag: Session, mut db: Connection) -> String { + sqlx::query!( + "SELECT display_name, username FROM users WHERE id = $1", + id as i32 + ) + .fetch_one(&mut **db) + .await + .map(|row| row.display_name.unwrap_or(row.username)) + .unwrap_or_else(|_| "User not found".to_string()) } #[launch] @@ -80,7 +83,7 @@ fn rocket() -> Rocket { messages::post_message, messages::event_stream, users, - username_for_id, + display_name, auth::signup, auth::login, auth::get_totp, diff --git a/backend/src/messages.rs b/backend/src/messages.rs index 8062a68..4be21d1 100644 --- a/backend/src/messages.rs +++ b/backend/src/messages.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use rocket::{ + Shutdown, response::stream::{Event, EventStream}, serde::json::Json, time::OffsetDateTime, @@ -9,9 +10,9 @@ use rocket_db_pools::Connection; use rocket_dyn_templates::{Template, context}; use serde::{Deserialize, Serialize}; use sqlx::prelude::FromRow; -use tokio::sync::broadcast; +use tokio::{select, sync::broadcast}; -use crate::{auth::Session, db::DbConn, llm::LlmWorker}; +use crate::{auth::Session, db::DbConn, display_name, llm::LlmWorker}; /// ---------- shared broadcaster ---------- pub struct ChatBroadcaster { @@ -54,14 +55,25 @@ pub async fn post_message( let chat = chat.inner().clone(); + let display_name = sqlx::query!( + "SELECT display_name, username FROM users WHERE id = $1", + session.user_id as i32 + ) + .fetch_one(&mut **db) + .await + .map(|row| row.display_name.unwrap_or(row.username)) + .unwrap_or_else(|_| "Unknown".to_string()); + msg.user_id = session.user_id; + msg.display_name = Some(display_name); chat.publish(msg.clone().into_inner()).await; sqlx::query!( - "INSERT INTO messages (channel_id, user_id, content) VALUES ($1, $2, $3)", + "INSERT INTO messages (channel_id, user_id, content, created_at) VALUES ($1, $2, $3, $4)", CHANNEL_ID, msg.user_id as i32, - msg.text + msg.text, + OffsetDateTime::from_unix_timestamp_nanos(msg.timestamp as i128 * 1_000_000).unwrap() ) .execute(&mut **db) .await @@ -71,15 +83,15 @@ pub async fn post_message( tokio::spawn(async move { let response = LlmWorker::new(LMSTUDIO_URI.to_string()).query(&msg).await; - if let Ok(message) = response { - chat.publish(message.clone()).await; + if let Ok(reply) = response { + chat.publish(reply.clone()).await; sqlx::query!( "INSERT INTO messages (channel_id, user_id, content, created_at) VALUES ($1, $2, $3, $4)", CHANNEL_ID, - message.user_id as i32, - message.text, - OffsetDateTime::from_unix_timestamp(message.timestamp as i64).unwrap() + reply.user_id as i32, + reply.text, + OffsetDateTime::from_unix_timestamp_nanos(reply.timestamp as i128 * 1_000_000).unwrap() ) .execute(&mut **db) .await @@ -106,7 +118,7 @@ pub async fn get_messages(mut db: Connection, _session: Session) -> Json display_name: Some(msg.display_name.unwrap_or(msg.username)), user_id: msg.id as usize, text: msg.content, - timestamp: msg.created_at.unwrap().unix_timestamp() as usize, + timestamp: (msg.created_at.unwrap().unix_timestamp_nanos() / 1_000_000) as usize, }) .collect(), ) @@ -118,6 +130,7 @@ pub async fn event_stream( chat: &rocket::State>, db: Connection, ag: Session, + mut shutdown: Shutdown, ) -> EventStream![] { let mut rx = chat.subscribe(); @@ -128,18 +141,23 @@ pub async fn event_stream( } loop { - match rx.recv().await { - Ok(msg) => yield Event::json(&msg), - Err(broadcast::error::RecvError::Lagged(_)) => { - yield Event::comment("RecvError::Lagged"); - } - Err(broadcast::error::RecvError::Closed) => break, + select!{ + // exit early on shutdown + _ = &mut shutdown => break, + + msg = rx.recv() => match msg { + Ok(msg) => yield Event::json(&msg), + Err(broadcast::error::RecvError::Lagged(_)) => { + yield Event::comment("RecvError::Lagged"); + } + Err(broadcast::error::RecvError::Closed) => break, + }, } } } } -#[get("/")] +#[get("/chat")] pub async fn chat_page(session: Session) -> Template { Template::render("chat", context!(user_id: session.user_id)) } diff --git a/backend/static/css/index.css b/backend/static/css/index.css index 8560511..6aee4c6 100644 --- a/backend/static/css/index.css +++ b/backend/static/css/index.css @@ -8,9 +8,9 @@ body { font-family: "Segoe UI", Tahoma, Geneva, Verdana, sans-serif; background: #0a0a0a; color: #e0e0e0; - min-height: 100vh; - min-width: 100vw; - height: 100vh; + max-height: 100dvh; + min-width: 100dvw; + height: 100dvh; overflow: hidden; } @@ -33,7 +33,7 @@ body { .chat-container { display: flex; flex-direction: column; - height: 100vh; + height: 100dvh; min-width: 100vw; margin: 0 0; background: #121212; diff --git a/backend/templates/2fa.html.tera b/backend/templates/2fa.html.tera index a75d859..5a5cf30 100644 --- a/backend/templates/2fa.html.tera +++ b/backend/templates/2fa.html.tera @@ -3,54 +3,46 @@ - Discord Clone - Group Chat + Discord Clone - Sign Up -
- - -
-
- -

Chat title

-
+