add WebSocket support for server events, including /api/wss route and connection handling logic

Author: JrTimha

Cargo.lock

@@ -5,7 +5,7 @@ edition = "2024"
 
 [dependencies]
 log = "0.4.29"
-axum = { version = "0.8.8", features = ["multipart"] }
+axum = { version = "0.8.8", features = ["multipart", "ws"] }
 tokio = {version = "1.49.0", features = ["full"]}
 tower = "0.5.3"
 config = "0.15.18"

Cargo.toml

@@ -1,6 +1,6 @@
 ism_url = "localhost"
 ism_port= 5403
-log_level = "info"
+log_level = "debug"
 cors_origin = "http://localhost:4200"
 use_kafka = true
 push_notification_access_token="oiqhfriuhf"

default.config.toml

@@ -2,19 +2,25 @@ use std::sync::Arc;
 use std::time::Duration;
 use axum::{Extension, Json};
 use axum::extract::{Query, State};
-use axum::response::{Sse};
+use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
+use axum::response::{IntoResponse, Sse};
 use axum::response::sse::Event;
+use bytes::Bytes;
 use chrono::{DateTime, Utc};
 use futures::Stream;
-use log::error;
+use tokio::time;
+use log::{debug, error};
 use serde::Deserialize;
+use tokio::sync::broadcast::error::RecvError;
 use tokio_stream::wrappers::BroadcastStream;
 use tokio_stream::wrappers::errors::BroadcastStreamRecvError;
+use tracing::warn;
 use uuid::Uuid;
 use crate::broadcast::{BroadcastChannel, Notification};
 use crate::core::AppState;
 use crate::errors::{AppError, AppResponse};
 use crate::keycloak::decode::KeycloakToken;
+use crate::keycloak::layer::KeycloakAuthLayer;
 
 struct ConnectionGuard {
     user_id: Uuid,
@@ -64,6 +70,69 @@ pub async fn stream_server_events(
     )
 }
 
+
+pub async fn websocket_server_events(
+    websocket: WebSocketUpgrade,
+    Extension(token): Extension<KeycloakToken<String>>
+) -> impl IntoResponse {
+
+    websocket
+        .on_failed_upgrade(|error| warn!("Error upgrading websocket: {}", error))
+        .on_upgrade(move |socket| handle_socket(socket, token.subject.clone()))
+}
+
+async fn handle_socket(mut socket: WebSocket, user_id: Uuid) {
+
+    let mut broadcast_events = BroadcastChannel::get().subscribe_to_user_events(user_id.clone()).await;
+    let _guard = ConnectionGuard { user_id };
+    let mut ping_interval = time::interval(Duration::from_secs(30));
+
+    loop {
+        tokio::select! {
+            // 1. Handle new broadcasting event:
+            notification_result = broadcast_events.recv() => {
+                match notification_result {
+                    Ok(event) => {
+                        let json_msg = serde_json::to_string(&event).unwrap();
+                        let ws_message = Message::text(json_msg);
+
+                        if socket.send(ws_message).await.is_err() {
+                            error!("Failed to send message to client");
+                        }
+                    }
+                    Err(RecvError::Closed) => {
+                        debug!("Client disconnected or channel closed");
+                        break;
+                    }
+                    Err(RecvError::Lagged(_)) => {
+                        debug!("Client is too slow!")
+                    }
+                }
+            }
+
+            // 2. Regular ping from ism:
+            _ = ping_interval.tick() => {
+                if socket.send(Message::Ping(Bytes::new())).await.is_err() { // connection is dead when we can't send ping
+                    break;
+                }
+            }
+
+            // 3. Receive messages from the client:
+            client_msg = socket.recv() => {
+                match client_msg {
+                    Some(Ok(Message::Close(_))) | None => break, //client is closing connection
+                    Some(Err(_)) => break, //client error
+                    Some(Ok(Message::Pong(_))) => {
+                        debug!("Client has sent Pong");
+                    }
+                    _ => {} //for the future
+                }
+            }
+        }
+    }
+}
+
+
 #[derive(Deserialize)]
 pub struct NotificationQueryParam {
     timestamp: DateTime<Utc>

src/messaging/notifications.rs

@@ -1,13 +1,14 @@
 use std::sync::Arc;
 use axum::Router;
-use axum::routing::{get, post};
+use axum::routing::{any, get, post};
 use crate::core::AppState;
 use crate::messaging::handler::handle_send_message;
-use crate::messaging::notifications::{get_latest_notification_events, stream_server_events};
+use crate::messaging::notifications::{get_latest_notification_events, stream_server_events, websocket_server_events};
 
 pub fn create_messaging_routes() -> Router<Arc<AppState>> {
     Router::new() //add new routes here
         .route("/api/notifications", get(get_latest_notification_events))
         .route("/api/sse", get(stream_server_events))
+        .route("/api/wss", any(websocket_server_events))
         .route("/api/send-msg", post(handle_send_message))
 }