Merge pull request 'messaging implementation using valkey pubsub and websockets' (#12) from wip/messaging-wss into main

Reviewed-on: #12

Cargo.toml

@@ -28,6 +28,8 @@ toml = "0.8"
 url = { version = "2.5", features = ["serde"] }
 uuid = { version = "1.16", features = ["serde", "v7"] }
 random-string = "1.1"
+actix-ws = "0.3.0"
+futures-util = "0.3.31"
 
 [dependencies.tokio]
 version = "1.44"

src/api/v1/servers/uuid/channels/uuid/messages.rs

@@ -11,7 +11,7 @@ struct MessageRequest {
 }
 
 #[get("{uuid}/channels/{channel_uuid}/messages")]
-pub async fn get(req: HttpRequest, path: web::Path<(Uuid, Uuid)>, message_request: web::Json<MessageRequest>, data: web::Data<Data>) -> Result<HttpResponse, Error> {
+pub async fn get(req: HttpRequest, path: web::Path<(Uuid, Uuid)>, message_request: web::Query<MessageRequest>, data: web::Data<Data>) -> Result<HttpResponse, Error> {
     let headers = req.headers();
 
     let auth_header = get_auth_header(headers);

src/api/v1/servers/uuid/channels/uuid/mod.rs

@@ -1,4 +1,5 @@
 pub mod messages;
+pub mod socket;
 
 use actix_web::{delete, get, web, Error, HttpRequest, HttpResponse};
 use crate::{api::v1::auth::check_access_token, structs::{Channel, Member}, utils::get_auth_header, Data};

src/api/v1/servers/uuid/channels/uuid/socket.rs

@@ -0,0 +1,120 @@
+use actix_web::{get, rt, web, Error, HttpRequest, HttpResponse};
+use actix_ws::AggregatedMessage;
+use futures_util::StreamExt as _;
+use uuid::Uuid;
+use log::error;
+
+use crate::{api::v1::auth::check_access_token, structs::{Channel, Member}, utils::get_auth_header, Data};
+
+#[get("{uuid}/channels/{channel_uuid}/socket")]
+pub async fn echo(req: HttpRequest, path: web::Path<(Uuid, Uuid)>, stream: web::Payload, data: web::Data<Data>) -> Result<HttpResponse, Error> {
+    // Get all headers
+    let headers = req.headers();
+
+    // Retrieve auth header
+    let auth_header = get_auth_header(headers);
+
+    if let Err(error) = auth_header {
+        return Ok(error)
+    }
+
+    // Get uuids from path
+    let (guild_uuid, channel_uuid) = path.into_inner();
+
+    // Authorize client using auth header
+    let authorized = check_access_token(auth_header.unwrap(), &data.pool).await;
+
+    if let Err(error) = authorized {
+        return Ok(error)
+    }
+
+    // Unwrap user uuid from authorization
+    let uuid = authorized.unwrap();
+
+    // Get server member from psql
+    let member = Member::fetch_one(&data.pool, uuid, guild_uuid).await;
+
+    if let Err(error) = member {
+        return Ok(error);
+    }
+
+    // Get cache for channel
+    let cache_result = data.get_cache_key(format!("{}", channel_uuid)).await;
+
+    let channel: Channel;
+
+    // Return channel cache or result from psql as `channel` variable
+    if let Ok(cache_hit) = cache_result {
+        channel = serde_json::from_str(&cache_hit).unwrap()
+    } else {
+        let channel_result = Channel::fetch_one(&data.pool, guild_uuid, channel_uuid).await;
+
+        if let Err(error) = channel_result {
+            return Ok(error)
+        }
+    
+        channel = channel_result.unwrap();
+    
+        let cache_result = data.set_cache_key(format!("{}", channel_uuid), channel.clone(), 60).await;
+    
+        if let Err(error) = cache_result {
+            error!("{}", error);
+            return Ok(HttpResponse::InternalServerError().finish());
+        }
+    }
+
+    let (res, mut session_1, stream) = actix_ws::handle(&req, stream)?;
+
+    let mut stream = stream
+        .aggregate_continuations()
+        // aggregate continuation frames up to 1MiB
+        .max_continuation_size(2_usize.pow(20));
+
+    let pubsub_result = data.cache_pool.get_async_pubsub().await;
+
+    if let Err(error) = pubsub_result {
+        error!("{}", error);
+        return Ok(HttpResponse::InternalServerError().finish())
+    }
+
+    let mut session_2 = session_1.clone();
+
+    rt::spawn(async move {
+        let mut pubsub = pubsub_result.unwrap();
+        pubsub.subscribe(channel_uuid.to_string()).await.unwrap();
+        while let Some(msg) = pubsub.on_message().next().await {
+            let payload: String = msg.get_payload().unwrap();
+            session_1.text(payload).await.unwrap();
+        }
+    });
+
+    // start task but don't wait for it
+    rt::spawn(async move {
+        let mut conn = data.cache_pool.get_multiplexed_tokio_connection().await.unwrap();
+        // receive messages from websocket
+        while let Some(msg) = stream.next().await {
+            match msg {
+                Ok(AggregatedMessage::Text(text)) => {
+                    // echo text message
+                    redis::cmd("PUBLISH").arg(&[channel_uuid.to_string(), text.to_string()]).exec_async(&mut conn).await.unwrap();
+                    channel.new_message(&data.pool, uuid, text.to_string()).await.unwrap();
+                }
+
+                Ok(AggregatedMessage::Binary(bin)) => {
+                    // echo binary message
+                    session_2.binary(bin).await.unwrap();
+                }
+
+                Ok(AggregatedMessage::Ping(msg)) => {
+                    // respond to PING frame with PONG frame
+                    session_2.pong(&msg).await.unwrap();
+                }
+
+                _ => {}
+            }
+        }
+    });
+
+    // respond immediately with response connected to WS session
+    Ok(res)
+}

src/api/v1/servers/uuid/mod.rs

@@ -17,6 +17,7 @@ pub fn web() -> Scope {
         .service(channels::uuid::get)
         .service(channels::uuid::delete)
         .service(channels::uuid::messages::get)
+        .service(channels::uuid::socket::echo)
         // Roles
         .service(roles::get)
         .service(roles::create)

src/main.rs

@@ -1,5 +1,5 @@
 use actix_cors::Cors;
-use actix_web::{App, HttpServer, web};
+use actix_web::{web, App, HttpServer};
 use argon2::Argon2;
 use clap::Parser;
 use simple_logger::SimpleLogger;
@@ -169,7 +169,6 @@ async fn main() -> Result<(), Error> {
         argon2: Argon2::default(),
         start_time: SystemTime::now(),
     };
-
     
     HttpServer::new(move || {
         // Set CORS headers
@@ -207,5 +206,6 @@ async fn main() -> Result<(), Error> {
     .bind((web.url, web.port))?
     .run()
     .await?;
+
     Ok(())
 }

src/structs.rs

@@ -171,7 +171,7 @@ impl Channel {
     }
 
     pub async fn fetch_messages(&self, pool: &Pool<Postgres>, amount: i64, offset: i64) -> Result<Vec<Message>, HttpResponse> {
-        let row = sqlx::query_as(&format!("SELECT uuid, user_uuid, message FROM channels WHERE channel_uuid = '{}' ORDER BY uuid LIMIT $1 OFFSET $2", self.uuid))
+        let row = sqlx::query_as(&format!("SELECT CAST(uuid AS VARCHAR), CAST(user_uuid AS VARCHAR), CAST(channel_uuid AS VARCHAR), message FROM messages WHERE channel_uuid = '{}' ORDER BY uuid DESC LIMIT $1 OFFSET $2", self.uuid))
             .bind(amount)
             .bind(offset)
             .fetch_all(pool)
@@ -186,6 +186,27 @@ impl Channel {
 
         Ok(message_builders.iter().map(|b| b.build()).collect())
     }
+
+    pub async fn new_message(&self, pool: &Pool<Postgres>, user_uuid: Uuid, message: String) -> Result<Message, HttpResponse> {
+        let message_uuid = Uuid::now_v7();
+
+        let row = sqlx::query(&format!("INSERT INTO messages (uuid, channel_uuid, user_uuid, message) VALUES ('{}', '{}', '{}', $1)", message_uuid, self.uuid, user_uuid))
+            .bind(&message)
+            .execute(pool)
+            .await;
+    
+        if let Err(error) = row {
+            error!("{}", error);
+            return Err(HttpResponse::InternalServerError().finish());
+        }
+
+        Ok(Message {
+            uuid: message_uuid,
+            channel_uuid: self.uuid,
+            user_uuid,
+            message,
+        })
+    }
 }
 
 #[derive(Clone, Copy)]