megacommit

backend/src/auth.rs

@@ -1,6 +1,21 @@
-use rocket::{post, serde::json::Json};
-use rocket_db_pools::{Connection, Database, sqlx};
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use rand::Rng;
+use rocket::{
+    Request,
+    http::{CookieJar, Status},
+    outcome::Outcome,
+    post,
+    request::{self, FromRequest},
+    serde::json::Json,
+};
+use rocket_db_pools::{
+    Connection, Database,
+    sqlx::{self},
+};
+use rocket_dyn_templates::{Template, context};
 use serde::{Deserialize, Serialize};
+use sha2::{Digest, Sha256};
 
 #[derive(Database)]
 #[database("postgres_db")]
@@ -14,12 +29,44 @@ pub struct UserCredentials {
 
 #[post("/signup", data = "<cred>")]
 pub async fn signup(
-    conn: Connection<DbConn>,
     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 = SessionToken::new(result.id as usize);
+    let result = sqlx::query!(
+        "INSERT INTO sessions (user_id, token) VALUES ($1, $2)",
+        result.id,
+        session.token,
+    )
+    .execute(&mut **db)
+    .await;
+
+    if let Err(e) = result {
+        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("/signup")]
+pub async fn signup_page() -> Template {
+    Template::render("signup", context!())
+}
+
 #[post("/login", data = "<cred>")]
 pub async fn login(
     conn: Connection<DbConn>,
@@ -28,3 +75,58 @@ pub async fn login(
     // TODO: implement actual login logic, e.g. verify password and generate token
     Ok(Json("Login successful".to_string()))
 }
+
+pub struct SessionToken {
+    pub token: String,
+    pub user_id: usize,
+}
+
+impl SessionToken {
+    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()));
+        SessionToken {
+            token: hashed,
+            user_id,
+        }
+    }
+}
+
+type UserID = usize;
+
+#[derive(Debug)]
+pub struct AuthGuard(pub UserID);
+
+#[rocket::async_trait]
+impl<'r> FromRequest<'r> for AuthGuard {
+    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 FROM sessions WHERE token = $1 AND expires_at > NOW()",
+                value
+            )
+            .fetch_optional(&mut **pool)
+            .await
+            .expect("query failed!");
+
+            if let Some(token) = result {
+                let user_id = token.user_id;
+                Outcome::Success(AuthGuard(user_id as usize))
+            } else {
+                Outcome::Error((Status::Unauthorized, ()))
+            }
+        } else {
+            Outcome::Error((Status::Unauthorized, ()))
+        }
+    }
+}

backend/src/cdn.rs

@@ -0,0 +1,25 @@
+use std::path::{Path, PathBuf};
+
+use rocket::fs::NamedFile;
+
+pub fn routes() -> Vec<rocket::Route> {
+    routes![profile_pic, general]
+}
+
+#[get("/profile/<user_id>")]
+pub async fn profile_pic(user_id: usize) -> Option<NamedFile> {
+    if let Ok(image) =
+        NamedFile::open(Path::new("./cdn/profiles/").join(format!("{}.jpg", user_id))).await
+    {
+        Some(image)
+    } else {
+        Some(NamedFile::open("./cdn/profiles/default.jpg").await.ok()?)
+    }
+}
+
+#[get("/<file_name..>")]
+pub async fn general(file_name: PathBuf) -> Option<NamedFile> {
+    NamedFile::open(Path::new("./cdn/").join(file_name))
+        .await
+        .ok()
+}

backend/src/llm.rs

@@ -1,6 +1,8 @@
 // src/llm.rs
 use serde::{Deserialize, Serialize};
 
+use crate::messages::ChatMsg;
+
 #[derive(Serialize)]
 struct LlmRequest {
     model: String,
@@ -13,36 +15,51 @@ struct Message {
     content: String,
 }
 
-pub async fn query_llm(text: &str) -> Result<String, String> {
-    let client = reqwest::Client::new();
-
-    // Build the request body
-    let payload = LlmRequest {
-        model: "gemma2-9b-it".into(), // whatever model you run locally
-        messages: vec![Message {
-            role: "user".into(),
-            content: text.into(),
-        }],
-    };
-
-    // POST to lm‑studio (default 127.0.0.1:1234)
-    let resp = client
-        .post("http://127.0.0.1:1234/v1/chat/completions")
-        .json(&payload)
-        .send()
-        .await
-        .unwrap();
-
-    // The API returns a JSON with `choices[].message.content`
-    #[derive(Deserialize)]
-    struct LlmResponse {
-        choices: Vec<Choice>,
-    }
-    #[derive(Deserialize)]
-    struct Choice {
-        message: Message,
-    }
-
-    let llm_resp: LlmResponse = resp.json().await.unwrap();
-    Ok(llm_resp.choices[0].message.content.clone())
+pub struct LlmWorker {
+    uri: String,
+}
+
+impl LlmWorker {
+    pub fn new(uri: String) -> Self {
+        Self { uri }
+    }
+
+    pub async fn query(&self, message: &ChatMsg) -> Result<ChatMsg, String> {
+        let client = reqwest::Client::new();
+
+        // Build the request body
+        let payload = LlmRequest {
+            model: "gemma2-9b-it".into(), // whatever model you run locally
+            messages: vec![Message {
+                role: "user".into(),
+                content: message.text.clone().into(),
+            }],
+        };
+
+        // POST to lm‑studio (default 127.0.0.1:1234)
+        let resp = client
+            .post(self.uri.clone())
+            .json(&payload)
+            .send()
+            .await
+            .unwrap();
+
+        // The API returns a JSON with `choices[].message.content`
+        #[derive(Deserialize)]
+        struct LlmResponse {
+            choices: Vec<Choice>,
+        }
+        #[derive(Deserialize)]
+        struct Choice {
+            message: Message,
+        }
+
+        let llm_resp: LlmResponse = resp.json().await.unwrap();
+
+        Ok(ChatMsg {
+            user_id: 0,
+            text: llm_resp.choices[0].message.content.clone(),
+            timestamp: chrono::Local::now().timestamp() as usize,
+        })
+    }
 }

backend/src/main.rs

@@ -3,134 +3,44 @@
 extern crate rocket;
 
 use rocket::fairing::Fairing;
+use rocket::fs::FileServer;
 use rocket::http::Method;
 use rocket::response::stream::{Event, EventStream};
 use rocket::serde::json::Json;
 use rocket::{Build, Rocket};
 use rocket_cors::{AllowedOrigins, CorsOptions};
+use rocket_db_pools::{Connection, Database};
+use rocket_dyn_templates::{Template, context};
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
 use tokio::sync::broadcast;
 
-use crate::llm::query_llm;
+use crate::auth::{AuthGuard, DbConn};
+use crate::llm::LlmWorker;
+use crate::messages::ChatBroadcaster;
 
 pub mod auth;
+pub mod cdn;
 pub mod llm;
+pub mod messages;
 
-/// ---------- shared broadcaster ----------
-struct ChatBroadcaster {
-    sender: broadcast::Sender<ChatMsg>,
+#[get("/users", rank = 2)]
+async fn users(_ag: AuthGuard, mut db: Connection<DbConn>) -> Json<Vec<i32>> {
+    sqlx::query!("SELECT id FROM users")
+        .fetch_all(&mut **db)
+        .await
+        .map(|rows| rows.into_iter().map(|row| row.id).collect())
+        .unwrap_or_else(|_| Vec::new())
+        .into()
 }
 
-impl ChatBroadcaster {
-    fn new(buffer_size: usize) -> Self {
-        let (sender, _rx) = broadcast::channel::<ChatMsg>(buffer_size);
-        Self { sender }
-    }
-
-    async fn publish(&self, msg: ChatMsg) {
-        let _ = self.sender.send(msg);
-    }
-
-    fn subscribe(&self) -> broadcast::Receiver<ChatMsg> {
-        self.sender.subscribe()
-    }
-}
-
-/// ---------- Rocket routes ----------
-#[derive(Debug, Serialize, Deserialize, Clone)]
-struct ChatMsg {
-    userid: usize,
-    text: String,
-    timestamp: usize,
-}
-
-#[post("/chat", format = "json", data = "<msg>")]
-async fn post_message(
-    msg: Json<ChatMsg>,
-    chat: &rocket::State<Arc<ChatBroadcaster>>,
-) -> &'static str {
-    chat.publish(msg.into_inner()).await;
-    "Message sent"
-}
-
-#[get("/events")]
-async fn event_stream(chat: &rocket::State<Arc<ChatBroadcaster>>) -> EventStream![] {
-    let mut rx = chat.subscribe();
-
-    EventStream! {
-        loop {
-            match rx.recv().await {
-                Ok(msg) => yield Event::json(&msg),
-                Err(broadcast::error::RecvError::Lagged(_)) => {
-                    yield Event::comment("lagged");
-                }
-                Err(broadcast::error::RecvError::Closed) => break,
-            }
-        }
-    }
-}
-
-// ---------- LLM worker ----------
-async fn start_llm_worker(chat: Arc<ChatBroadcaster>) {
-    let mut rx = chat.subscribe();
-
-    loop {
-        match rx.recv().await {
-            Ok(msg) => {
-                if msg.userid == 0 {
-                    // ignore bot messages
-                    continue;
-                }
-
-                let user_text = msg.text.clone();
-                let chat_clone = chat.clone();
-
-                rocket::tokio::spawn(async move {
-                    match query_llm(&user_text).await {
-                        Ok(reply) => {
-                            let bot_msg = ChatMsg {
-                                userid: 0,
-                                text: reply,
-                                timestamp: chrono::Local::now().timestamp() as usize,
-                            };
-                            chat_clone.publish(bot_msg).await;
-                        }
-                        Err(e) => eprintln!("LLM error: {}", e),
-                    }
-                });
-            }
-            Err(_) => break, // channel closed
-        }
-    }
-}
-
-pub struct LlmWorkerFairing;
-
-#[rocket::async_trait]
-impl Fairing for LlmWorkerFairing {
-    fn info(&self) -> rocket::fairing::Info {
-        rocket::fairing::Info {
-            name: "LLM background worker",
-            kind: rocket::fairing::Kind::Ignite,
-        }
-    }
-
-    async fn on_ignite(&self, rocket: Rocket<Build>) -> rocket::fairing::Result {
-        // Grab the shared broadcaster from state
-        let chat = rocket
-            .state::<Arc<ChatBroadcaster>>()
-            .expect("ChatBroadcaster not managed");
-        // Clone it so we can move into async block
-        let chat_clone = Arc::clone(chat);
-
-        // Spawn the background worker **inside** on_ignite
-        tokio::spawn(async move {
-            start_llm_worker(chat_clone).await;
-        });
-
-        Ok(rocket)
-    }
+#[get("/users/<id>", rank = 1)]
+async fn username_for_id(id: usize, _ag: AuthGuard, mut db: Connection<DbConn>) -> 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())
 }
 
 /// ---------- launch ----------
@@ -151,9 +61,22 @@ fn rocket() -> Rocket<Build> {
     rocket::build()
         .manage(chat)
         .attach(cors.to_cors().unwrap())
-        .attach(LlmWorkerFairing {})
+        .attach(DbConn::init())
+        .attach(Template::fairing())
+        .mount("/static", FileServer::from("static"))
+        .mount("/cdn", cdn::routes())
         .mount(
             "/",
-            routes![post_message, event_stream, auth::signup, auth::login],
+            routes![
+                users,
+                username_for_id,
+                messages::chat_page,
+                messages::get_messages,
+                messages::post_message,
+                messages::event_stream,
+                auth::signup,
+                auth::signup_page,
+                auth::login
+            ],
         )
 }

backend/src/messages.rs

@@ -0,0 +1,136 @@
+use std::sync::Arc;
+
+use rocket::{
+    response::stream::{Event, EventStream},
+    serde::json::Json,
+};
+use rocket_db_pools::Connection;
+use rocket_dyn_templates::{Template, context};
+use serde::{Deserialize, Serialize};
+use tokio::sync::broadcast;
+
+use crate::{
+    auth::{AuthGuard, DbConn},
+    llm::LlmWorker,
+};
+
+/// ---------- shared broadcaster ----------
+pub struct ChatBroadcaster {
+    sender: broadcast::Sender<ChatMsg>,
+}
+
+impl ChatBroadcaster {
+    pub fn new(buffer_size: usize) -> Self {
+        let (sender, _rx) = broadcast::channel::<ChatMsg>(buffer_size);
+        Self { sender }
+    }
+
+    pub async fn publish(&self, msg: ChatMsg) {
+        let _ = self.sender.send(msg);
+    }
+
+    pub fn subscribe(&self) -> broadcast::Receiver<ChatMsg> {
+        self.sender.subscribe()
+    }
+}
+
+/// ---------- Rocket routes ----------
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct ChatMsg {
+    pub user_id: usize,
+    pub text: String,
+    pub timestamp: usize,
+}
+
+#[post("/chat", format = "json", data = "<msg>")]
+pub async fn post_message(
+    mut msg: Json<ChatMsg>,
+    chat: &rocket::State<Arc<ChatBroadcaster>>,
+    mut db: Connection<DbConn>,
+    ag: AuthGuard,
+) -> Result<(), String> {
+    const CHANNEL_ID: i32 = 1;
+    const LMSTUDIO_URI: &'static str = "http://127.0.0.1:1234/v1/chat/completions";
+
+    let chat = chat.inner().clone();
+
+    msg.user_id = ag.0;
+    chat.publish(msg.clone().into_inner()).await;
+
+    sqlx::query!(
+        "INSERT INTO messages (channel_id, user_id, content) VALUES ($1, $2, $3)",
+        CHANNEL_ID,
+        msg.user_id as i32,
+        msg.text
+    )
+    .execute(&mut **db)
+    .await
+    .map_err(|_| "Failed".to_string())?;
+
+    // get response
+    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;
+
+            sqlx::query!(
+                "INSERT INTO messages (channel_id, user_id, content) VALUES ($1, $2, $3)",
+                CHANNEL_ID,
+                message.user_id as i32,
+                message.text
+            )
+            .execute(&mut **db)
+            .await
+            .map_err(|_| "Failed".to_string())
+            .unwrap();
+        }
+    });
+
+    Ok(())
+}
+
+#[get("/messages")]
+pub async fn get_messages(mut db: Connection<DbConn>, _ag: AuthGuard) -> Json<Vec<ChatMsg>> {
+    Json(
+        sqlx::query!(
+            "SELECT user_id, content, created_at FROM messages ORDER BY created_at DESC LIMIT 100"
+        )
+        .fetch_all(&mut **db)
+        .await
+        .unwrap_or_else(|_| Vec::new())
+        .into_iter()
+        .rev()
+        .map(|row| ChatMsg {
+            user_id: row.user_id as usize,
+            text: row.content,
+            timestamp: row.created_at.unwrap().unix_timestamp() as usize,
+        })
+        .collect(),
+    )
+}
+
+#[get("/events")]
+pub async fn event_stream(
+    chat: &rocket::State<Arc<ChatBroadcaster>>,
+    _ag: AuthGuard,
+) -> EventStream![] {
+    let mut rx = chat.subscribe();
+
+    EventStream! {
+        loop {
+            match rx.recv().await {
+                Ok(msg) => yield Event::json(&msg),
+                Err(broadcast::error::RecvError::Lagged(_)) => {
+                    yield Event::comment("lagged");
+                }
+                Err(broadcast::error::RecvError::Closed) => break,
+            }
+        }
+    }
+}
+
+#[get("/")]
+pub async fn chat_page(ag: AuthGuard) -> Template {
+    Template::render("chat", context!(user_id: ag.0))
+}