WebSockets

use std::time::Duration;

use axum::{
    Router,
    extract::{
        State,
        ws::{Message, Utf8Bytes, WebSocket, WebSocketUpgrade},
    },
    http::{HeaderMap, StatusCode, header},
    response::{IntoResponse, Response},
    routing::{any, post},
};
use futures::{SinkExt, StreamExt};
use tokio::sync::broadcast;

const PING_EVERY: Duration = Duration::from_secs(30);
const IDLE_TIMEOUT: Duration = Duration::from_secs(60);
const ALLOWED_ORIGINS: [&str; 2] = ["https://bookshelf.example", "http://localhost:3000"];

#[derive(Clone)]
struct AppState {
    // one Sender, cloned into state: broadcast::Sender is already cheap to clone and thread-safe
    reviews: broadcast::Sender<String>,
}

async fn ws_handler(State(state): State<AppState>, headers: HeaderMap, ws: WebSocketUpgrade) -> Response {
    // Browsers send Origin on the upgrade; reject cross-site pages (CORS does not apply to WebSockets).
    let origin = headers.get(header::ORIGIN).and_then(|o| o.to_str().ok());
    if origin.is_some_and(|o| !ALLOWED_ORIGINS.contains(&o)) {
        return StatusCode::FORBIDDEN.into_response();
    }
    // (authenticate here with the session cookie or a one-time ticket; see below)
    ws.protocols(["bookshelf.reviews.v1"]) // subprotocol negotiation
        .max_message_size(64 * 1024)
        .on_upgrade(move |socket| handle_socket(socket, state))
}

async fn handle_socket(socket: WebSocket, state: AppState) {
    let (mut sink, mut stream) = socket.split();
    let mut feed = state.reviews.subscribe();

    // writer: broadcast messages and periodic pings -> client
    let mut writer = tokio::spawn(async move {
        let mut ping = tokio::time::interval(PING_EVERY);
        loop {
            tokio::select! {
                msg = feed.recv() => match msg {
                    Ok(review) => {
                        if sink.send(Message::Text(Utf8Bytes::from(review))).await.is_err() {
                            break; // client gone
                        }
                    }
                    // a slow client missed messages: tell it, keep going (or close it)
                    Err(broadcast::error::RecvError::Lagged(n)) => {
                        let notice = format!("{{\"type\":\"lagged\",\"missed\":{n}}}");
                        if sink.send(Message::Text(notice.into())).await.is_err() {
                            break;
                        }
                    }
                    Err(broadcast::error::RecvError::Closed) => break,
                },
                _ = ping.tick() => {
                    if sink.send(Message::Ping(axum::body::Bytes::from_static(b"keepalive"))).await.is_err() {
                        break;
                    }
                }
            }
        }
    });

    // reader: client -> server, with an idle timeout (pongs count as activity)
    let mut reader = tokio::spawn(async move {
        loop {
            match tokio::time::timeout(IDLE_TIMEOUT, stream.next()).await {
                Ok(Some(Ok(Message::Text(text)))) => tracing::debug!(%text, "client message"),
                Ok(Some(Ok(Message::Close(_)))) | Ok(None) | Ok(Some(Err(_))) => break,
                Ok(Some(Ok(_))) => {} // Ping (answered automatically), Pong, Binary
                Err(_) => break,      // idle for too long
            }
        }
    });

    // when either side finishes, stop the other
    tokio::select! {
        _ = &mut writer => reader.abort(),
        _ = &mut reader => writer.abort(),
    }
}

async fn post_review(State(state): State<AppState>, body: String) -> StatusCode {
    let _ = state.reviews.send(body); // Err only if nobody is subscribed
    StatusCode::ACCEPTED
}

#[tokio::main]
async fn main() {
    let (reviews, _) = broadcast::channel(256);
    let app = Router::new()
        .route("/ws", any(ws_handler)) // any(): HTTP/1.1 GET upgrades and HTTP/2 CONNECT (RFC 8441)
        .route("/api/v1/reviews", post(post_review))
        .with_state(AppState { reviews });
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

$ python3 -c '
import asyncio, websockets
async def main():
    async with websockets.connect("ws://localhost:3000/ws", subprotocols=["bookshelf.reviews.v1"]) as ws:
        print("subprotocol:", ws.subprotocol)
        print("received:", await ws.recv())
asyncio.run(main())' &
$ curl -s -d 'Dune: 5 stars' localhost:3000/api/v1/reviews
subprotocol: bookshelf.reviews.v1
received: Dune: 5 stars

use std::time::Duration;

use axum::{
    Router,
    extract::ws::{Message, WebSocket, WebSocketUpgrade},
    response::Response,
    routing::get,
};

const IDLE: Duration = Duration::from_secs(60);

async fn ws_handler(ws: WebSocketUpgrade) -> Response {
    ws.max_message_size(64 * 1024) // reviews are small: refuse anything larger
        .max_frame_size(64 * 1024)
        .on_upgrade(handle_socket)
}

async fn handle_socket(mut socket: WebSocket) {
    loop {
        match tokio::time::timeout(IDLE, socket.recv()).await {
            Err(_elapsed) => {
                // nothing received for IDLE (not even a pong): close and free the task
                let _ = socket.send(Message::Close(None)).await;
                break;
            }
            Ok(Some(Ok(Message::Text(text)))) => {
                if socket.send(Message::Text(text)).await.is_err() {
                    break;
                }
            }
            Ok(Some(Ok(_))) => {}                 // ping, pong, binary
            Ok(Some(Err(_)) | None) => break,     // protocol error, oversized message or closed
        }
    }
}

#[tokio::main]
async fn main() {
    let app = Router::new().route("/ws", get(ws_handler));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}