WebSockets

Echo server and a broadcast chat room, using Tokio's broadcast channel.

Call req.upgrade_websocket in a handler and return the result. Your closure receives a socket you can recv() from and send() to. For a chat room, share a tokio::sync::broadcast channel through state and select! between the socket and the channel.

src/main.rs

use axeon::ws::Message;
use axeon::Server;
use tokio::sync::broadcast;

fn main() {
    let (tx, _) = broadcast::channel::<String>(64);

    let mut app = Server::new();
    app.state(tx);

    // Echo every message back.
    app.get("/echo", |mut req| async move {
        req.upgrade_websocket(|mut socket| async move {
            while let Some(Ok(message)) = socket.recv().await {
                if (message.is_text() || message.is_binary()) && socket.send(message).await.is_err() {
                    break;
                }
            }
        })
    });

    // Everything one client says goes to every client.
    app.get("/chat", |mut req| async move {
        let tx = req.state::<broadcast::Sender<String>>().unwrap().clone();
        req.upgrade_websocket(move |mut socket| async move {
            let mut rx = tx.subscribe();
            loop {
                tokio::select! {
                    incoming = socket.recv() => match incoming {
                        Some(Ok(message)) if message.is_text() => {
                            let _ = tx.send(message.as_text().unwrap_or_default().to_string());
                        }
                        Some(Ok(message)) if message.is_close() => break,
                        Some(Ok(_)) => {}
                        _ => break,
                    },
                    outgoing = rx.recv() => match outgoing {
                        Ok(text) => {
                            if socket.send(Message::text(text)).await.is_err() {
                                break;
                            }
                        }
                        Err(broadcast::error::RecvError::Lagged(_)) => {}
                        Err(_) => break,
                    },
                }
            }
        })
    });

    app.listen("127.0.0.1:3000").expect("server failed to start");
}

Terminal

# in two terminals
websocat ws://127.0.0.1:3000/chat

Notes#

Authenticate before upgrading: check a cookie, session or token in the handler and return an error instead of calling upgrade_websocket.
A bounded broadcast channel drops the oldest messages for slow readers (Lagged) instead of growing without limit.
On graceful shutdown Axeon closes open WebSockets for you.