use std::time::{SystemTime, UNIX_EPOCH}; use rand::Rng; use rocket::{ Request, fs::NamedFile, http::{CookieJar, Status}, outcome::Outcome, post, request::{self, FromRequest}, serde::json::Json, }; use rocket_db_pools::{ Connection, sqlx::{self}, }; use rocket_dyn_templates::{Template, context}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use sqlx::postgres::PgQueryResult; use totp_rs::{Algorithm, Secret, TOTP}; use crate::db::DbConn; #[derive(Serialize, Deserialize)] pub struct UserCredentials { pub username: String, pub password: String, } #[get("/signup")] pub async fn signup_page() -> Template { Template::render("signup", context!()) } #[post("/signup", data = "")] pub async fn signup( cred: Json, jar: &CookieJar<'_>, mut db: Connection, ) -> Result, String> { let result = sqlx::query!( "INSERT INTO users (username, password) VALUES ($1, $2) RETURNING id", cred.username, cred.password ) .fetch_one(&mut **db) .await .map_err(|e| e.to_string())?; 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()); } jar.add_private(("session", session.token)); println!("Signup successful"); Ok(Json("Signup successful".to_string())) } #[get("/login")] pub async fn login_page() -> Template { Template::render("login", context!()) } #[post("/login", data = "")] pub async fn login( mut db: Connection, jar: &CookieJar<'_>, cred: Json, ) -> Result, String> { if let Ok(row) = sqlx::query!( "SELECT id FROM users WHERE username = $1 AND password = $2", cred.username, cred.password, ) .fetch_one(&mut **db) .await { 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()); } jar.add_private(("session", session.token)); return Ok(Json("Signup successful".to_string())); } // TODO: implement actual login logic, e.g. verify password and generate token Err("login failed".to_string()) } #[get("/totp")] pub async fn mfa_page(session: Session) -> Template { Template::render("2fa", context!()) } #[get("/api/totp.jpg")] pub async fn get_totp(s: Session) -> Option { let totp = TOTP::new( Algorithm::SHA1, 6, 1, 30, Secret::generate_secret().to_bytes().unwrap(), Some("Github".to_string()), format!("{}", s.user_id), ) .unwrap(); let qr = totp.get_qr_base64().unwrap(); Some(QrCodeImage(qr.into())) } #[derive(Debug)] pub struct Session { pub token: String, pub user_id: usize, } impl Session { pub fn new(user_id: usize) -> Self { let current_time = SystemTime::now().duration_since(UNIX_EPOCH).unwrap(); let random: u32 = rand::rng().random(); let token = format!("{}-{}", current_time.as_secs(), random); let hashed = format!("{:x}", Sha256::digest(token.as_bytes())); Self { token: hashed, user_id, } } pub async fn commit(&self, db: &mut Connection) -> Result { sqlx::query!( "INSERT INTO sessions (user_id, token) VALUES ($1, $2)", self.user_id as i32, self.token, ) .execute(&mut ***db) .await } } #[rocket::async_trait] impl<'r> FromRequest<'r> for Session { type Error = (); async fn from_request(request: &'r Request<'_>) -> request::Outcome { if let Some(c) = request.cookies().get_private("session") { let mut pool = match request.guard::>().await { Outcome::Success(pool) => pool, _ => return Outcome::Error((Status::Unauthorized, ())), }; let value = c.value(); let result = sqlx::query!( "SELECT user_id, token FROM sessions WHERE token = $1 AND expires_at > NOW()", value ) .fetch_optional(&mut **pool) .await .expect("query failed!"); if let Some(session) = result { Outcome::Success(Self { user_id: session.user_id as usize, token: session.token, }) } else { Outcome::Error((Status::Unauthorized, ())) } } else { Outcome::Error((Status::Unauthorized, ())) } } } 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() } }