Module axum::extract::ws

Available on crate feature ws only.

Expand description

Handle WebSocket connections.

Example

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

let app = Router::new().route("/ws", get(handler));

async fn handler(ws: WebSocketUpgrade) -> Response {
    ws.on_upgrade(handle_socket)
}

async fn handle_socket(mut socket: WebSocket) {
    while let Some(msg) = socket.recv().await {
        let msg = if let Ok(msg) = msg {
            msg
        } else {
            // client disconnected
            return;
        };

        if socket.send(msg).await.is_err() {
            // client disconnected
            return;
        }
    }
}

Passing data and/or state to an `on_upgrade` callback

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

#[derive(Clone)]
struct State {
    // ...
}

async fn handler(ws: WebSocketUpgrade, Extension(state): Extension<State>) -> Response {
    ws.on_upgrade(|socket| handle_socket(socket, state))
}

async fn handle_socket(socket: WebSocket, state: State) {
    // ...
}

let app = Router::new()
    .route("/ws", get(handler))
    .layer(Extension(State { /* ... */ }));

Read and write concurrently

If you need to read and write concurrently from a WebSocket you can use StreamExt::split:

use axum::{Error, extract::ws::{WebSocket, Message}};
use futures::{sink::SinkExt, stream::{StreamExt, SplitSink, SplitStream}};

async fn handle_socket(mut socket: WebSocket) {
    let (mut sender, mut receiver) = socket.split();

    tokio::spawn(write(sender));
    tokio::spawn(read(receiver));
}

async fn read(receiver: SplitStream<WebSocket>) {
    // ...
}

async fn write(sender: SplitSink<WebSocket, Message>) {
    // ...
}