chatapp/backend/src/auth.rs

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 = "<cred>")]
pub async fn signup(
    cred: Json<UserCredentials>,
    jar: &CookieJar<'_>,
    mut db: Connection<DbConn>,
) -> Result<Json<String>, 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 = "<cred>")]
pub async fn login(
    mut db: Connection<DbConn>,
    jar: &CookieJar<'_>,
    cred: Json<UserCredentials>,
) -> Result<Json<String>, 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<QrCodeImage> {
    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<DbConn>) -> Result<PgQueryResult, sqlx::Error> {
        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<Self, Self::Error> {
        if let Some(c) = request.cookies().get_private("session") {
            let mut pool = match request.guard::<Connection<DbConn>>().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<u8>);

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()
    }
}