chat

src/main.rs

@@ -0,0 +1,82 @@
+use axum::{
+    extract::{
+        ws::{Message, WebSocket, WebSocketUpgrade},
+        State,
+    },
+    response::{Html, IntoResponse},
+    routing::get,
+    Router,
+};
+use futures::{sink::SinkExt, stream::StreamExt};
+use std::sync::Arc;
+use tokio::sync::broadcast;
+
+// Define application state to hold the broadcast channel
+struct AppState {
+    tx: broadcast::Sender<String>,
+}
+
+#[tokio::main]
+async fn main() {
+    // Create a broadcast channel with a capacity of 100 messages
+    let (tx, _rx) = broadcast::channel(100);
+    let app_state = Arc::new(AppState { tx });
+
+    // Build application with routes
+    let app = Router::new()
+        .route("/", get(index))
+        .route("/ws", get(websocket_handler))
+        .with_state(app_state);
+
+    // Run the app using hyper, listening globally on port 3000
+    let listener = tokio::net::TcpListener::bind("0.0.0.0:3001").await.unwrap();
+    println!("Chat server listening on http://0.0.0.0:3001");
+    axum::serve(listener, app).await.unwrap();
+}
+
+// Handler to serve the HTML file
+async fn index() -> Html<&'static str> {
+    Html(include_str!("../index.html"))
+}
+
+// Handler to upgrade HTTP connection to WebSocket
+async fn websocket_handler(
+    ws: WebSocketUpgrade,
+    State(state): State<Arc<AppState>>,
+) -> impl IntoResponse {
+    ws.on_upgrade(|socket| websocket(socket, state))
+}
+
+// The actual WebSocket connection handler
+async fn websocket(stream: WebSocket, state: Arc<AppState>) {
+    // Split the stream into sender and receiver
+    let (mut sender, mut receiver) = stream.split();
+
+    // Subscribe to the broadcast channel
+    let mut rx = state.tx.subscribe();
+
+    // Spawn a task to send messages from the broadcast channel to this client
+    let mut send_task = tokio::spawn(async move {
+        while let Ok(msg) = rx.recv().await {
+            // In a real app, you might want to handle errors better
+            if sender.send(Message::Text(msg.into())).await.is_err() {
+                break;
+            }
+        }
+    });
+
+    // Spawn a task to receive messages from this client and broadcast them
+    let tx = state.tx.clone();
+    let mut recv_task = tokio::spawn(async move {
+        while let Some(Ok(Message::Text(text))) = receiver.next().await {
+            // Broadcast the message to all subscribers
+            let _ = tx.send(text.to_string());
+        }
+    });
+
+    // If either task finishes (connection closed or error), abort the other
+    tokio::select! {
+        _ = (&mut send_task) => recv_task.abort(),
+        _ = (&mut recv_task) => send_task.abort(),
+    };
+}