WebSockets and Server-Sent Events in Axum

WebSockets and Server-Sent Events (SSE) keep a connection open so the server can push updates. This lesson builds a live order-status feed for the Storefront with both.

WebSocket: two-way messages over an upgraded HTTP connection. Reconnect is written by the application.
SSE: one-way events over one long HTTP response. The browser EventSource reconnects on its own.

Project setup

Cargo.toml

TOML

[package]
name = "storefront-live"
version = "0.1.0"
edition = "2021"

[dependencies]
axum = { version = "0.8", features = ["ws"] }
tokio = { version = "1", features = ["full"] }
tokio-stream = { version = "0.1", features = ["sync"] }
futures-util = "0.3"
serde = { version = "1", features = ["derive"] }
serde_json = "1"

No database or broker is needed. Events live in memory. For a WebSocket client from the terminal, install websocat.

Example 1: The smallest WebSocket handler

A WebSocket starts as an HTTP GET with Upgrade: websocket. The server answers 101 Switching Protocols, and the same TCP connection then carries WebSocket frames (text, binary, ping, pong, close).

src/main.rs

Rust

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

async fn ws_handler(upgrade: WebSocketUpgrade) -> Response {
    upgrade.on_upgrade(handle_socket)
}

async fn handle_socket(mut socket: WebSocket) {
    while let Some(Ok(message)) = socket.recv().await {
        match message {
            Message::Text(text) => {
                let reply = format!("echo: {text}");
                if socket.send(Message::Text(reply.into())).await.is_err() {
                    break;
                }
            }
            Message::Close(_) => break,
            _ => {}
        }
    }
}

#[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();
}

WebSocketUpgrade is an extractor. on_upgrade returns the 101 response at once and runs handle_socket as a separate task afterwards.
recv() returns Option<Result<Message, Error>>. None means the connection closed.
Message::Text holds Utf8Bytes in Axum 0.8. .into() converts from String.
A failed send means the client is gone, so the loop ends.
Axum answers ping frames automatically. A peer that vanishes without a close frame is detected only by a failed read or write, so servers send their own pings on a timer.

Example 2: Broadcast to many WebSocket clients

A tokio::sync::broadcast channel has one sender and one receiver per connection. Each receiver sees every message sent after it subscribed. WebSocket implements Stream and Sink, so split() gives a sender half and a receiver half that run in separate tasks.

This example replaces src/main.rs. POST /publish sends the request body to every connected client.

src/main.rs

Rust

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        State,
    },
    response::Response,
    routing::{get, post},
    Router,
};
use futures_util::{SinkExt, StreamExt};
use tokio::sync::broadcast::{self, error::RecvError};

async fn ws_handler(
    upgrade: WebSocketUpgrade,
    State(tx): State<broadcast::Sender<String>>,
) -> Response {
    upgrade.on_upgrade(move |socket| handle_socket(socket, tx))
}

async fn handle_socket(socket: WebSocket, tx: broadcast::Sender<String>) {
    let (mut sender, mut receiver) = socket.split();
    let mut events = tx.subscribe();

    let mut send_task = tokio::spawn(async move {
        loop {
            match events.recv().await {
                Ok(text) => {
                    if sender.send(Message::Text(text.into())).await.is_err() {
                        return;
                    }
                }
                Err(RecvError::Lagged(_)) => {
                    let _ = sender.send(Message::Text("resync".into())).await;
                    return;
                }
                Err(RecvError::Closed) => return,
            }
        }
    });

    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(message)) = receiver.next().await {
            if matches!(message, Message::Close(_)) {
                return;
            }
        }
    });

    tokio::select! {
        _ = &mut send_task => recv_task.abort(),
        _ = &mut recv_task => send_task.abort(),
    }
}

async fn publish(State(tx): State<broadcast::Sender<String>>, body: String) {
    let _ = tx.send(body);
}

#[tokio::main]
async fn main() {
    let (tx, _) = broadcast::channel::<String>(64);
    let app = Router::new()
        .route("/ws", get(ws_handler))
        .route("/publish", post(publish))
        .with_state(tx);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

The channel keeps the last 64 messages. A receiver that falls further behind gets RecvError::Lagged(n). This handler tells the client to resync and closes the connection.
select! waits for the first task to finish and aborts the other. Without it, a closed client leaves the send task running until its next failed send.
tx.send returns an error when no receiver exists. let _ = ignores it.
This design is per process. With several instances, each instance subscribes to a shared source (Redis pub/sub or a topic from messaging and Kafka) and forwards into its local channel.

Example 3: Server-Sent Events

SSE is a response with Content-Type: text/event-stream that never ends. Each event is text lines ended by a blank line.

Text

id: 42
event: order_status
data: {"order_id":7,"status":"paid"}

data: the payload. event: a name that EventSource dispatches to a listener. id: the id the browser remembers. retry: reconnect delay in milliseconds. A line that starts with : is a comment used as keep-alive.
EventSource sends only GET and cannot set headers, so it authenticates with cookies. Clients that need an Authorization header read the stream with fetch.

This example replaces src/main.rs with the same broadcast channel served as SSE.

src/main.rs

Rust

use axum::{
    extract::State,
    response::sse::{Event, KeepAlive, Sse},
    routing::{get, post},
    Router,
};
use std::convert::Infallible;
use tokio::sync::broadcast;
use tokio_stream::{wrappers::BroadcastStream, Stream, StreamExt};

async fn events(
    State(tx): State<broadcast::Sender<String>>,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    let stream = BroadcastStream::new(tx.subscribe())
        .map_while(|item| item.ok())
        .map(|text| Ok(Event::default().event("order_status").data(text)));
    Sse::new(stream).keep_alive(KeepAlive::default())
}

async fn publish(State(tx): State<broadcast::Sender<String>>, body: String) {
    let _ = tx.send(body);
}

#[tokio::main]
async fn main() {
    let (tx, _) = broadcast::channel::<String>(64);
    let app = Router::new()
        .route("/events", get(events))
        .route("/publish", post(publish))
        .with_state(tx);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

map_while ends the stream on a lag error. The browser then reconnects, which Example 5 turns into a resume.
KeepAlive writes comment lines at intervals so proxies do not close an idle response. The proxy idle timeout must exceed that interval.
Choose SSE when the client only listens and HTTP-only proxies are in the path. Choose WebSocket for frequent client messages or binary payloads.

Example 4: Origin check and ticket authentication

Checks run during the handshake, before on_upgrade, because a WebSocket outlives the request.

Browsers attach Origin to WebSocket handshakes and apply no same-origin policy. Without an allow list, a page on another site can open a socket with the user's cookies (cross-site WebSocket hijacking).
Browser JavaScript cannot set Authorization on a WebSocket. A short-lived, single-use ticket in the query string solves this: an authenticated REST call issues it, and the handshake redeems it. Token checks follow authentication and authorization and OAuth, JWT and mTLS.

src/auth.rs

Rust

use axum::http::{header, HeaderMap, StatusCode};
use std::{collections::HashMap, sync::Mutex};

const ALLOWED_ORIGINS: [&str; 1] = ["http://localhost:3000"];

pub fn check_origin(headers: &HeaderMap) -> Result<(), StatusCode> {
    let origin = headers
        .get(header::ORIGIN)
        .and_then(|value| value.to_str().ok())
        .ok_or(StatusCode::FORBIDDEN)?;
    if ALLOWED_ORIGINS.contains(&origin) {
        Ok(())
    } else {
        Err(StatusCode::FORBIDDEN)
    }
}

#[derive(Default)]
pub struct Tickets {
    inner: Mutex<(u64, HashMap<String, u64>)>,
}

impl Tickets {
    pub fn issue(&self, user_id: u64) -> String {
        let mut guard = self.inner.lock().unwrap();
        guard.0 += 1;
        let ticket = format!("t{}", guard.0);
        guard.1.insert(ticket.clone(), user_id);
        ticket
    }

    pub fn redeem(&self, ticket: &str) -> Option<u64> {
        self.inner.lock().unwrap().1.remove(ticket)
    }
}

redeem removes the ticket, so it works once. Production tickets are random, expire in seconds, and live in Redis (Redis caching). The counter here keeps the demo output fixed.
The user id returned by redeem is also the identity for per-message authorization: a subscribe request for an order must check that the user owns it, as a REST handler does.

Example 5: Resume with ids and a backlog

After a drop, EventSource reconnects with the header Last-Event-ID. A WebSocket client sends its last id itself, here as ?after=N. The server replays stored events after that id, then continues live.

Two orderings prevent gaps and duplicates. Subscribing and copying the backlog happen under one lock, and publishing sends on the channel under the same lock.

src/state.rs

Rust

use serde::Serialize;
use std::sync::{Arc, Mutex};
use tokio::sync::broadcast;

#[derive(Clone, Serialize)]
pub struct OrderEvent {
    pub id: u64,
    pub order_id: u64,
    pub status: String,
}

#[derive(Clone)]
pub struct Feed {
    log: Arc<Mutex<Vec<OrderEvent>>>,
    tx: broadcast::Sender<OrderEvent>,
}

impl Feed {
    pub fn new() -> Self {
        let (tx, _) = broadcast::channel(64);
        Feed { log: Arc::new(Mutex::new(Vec::new())), tx }
    }

    pub fn publish(&self, order_id: u64, status: String) -> OrderEvent {
        let mut log = self.log.lock().unwrap();
        let event = OrderEvent { id: log.len() as u64 + 1, order_id, status };
        log.push(event.clone());
        let _ = self.tx.send(event.clone());
        event
    }

    pub fn subscribe_from(&self, after: u64) -> (Vec<OrderEvent>, broadcast::Receiver<OrderEvent>) {
        let log = self.log.lock().unwrap();
        let receiver = self.tx.subscribe();
        let backlog = log.iter().filter(|event| event.id > after).cloned().collect();
        (backlog, receiver)
    }
}

Ids come from a counter here. In the Storefront they come from the outbox sequence in SQL transactions.

Backpressure and slow clients

Backpressure means a full buffer slows or stops the producer instead of growing without limit. Options for a client that cannot keep up:

Drop old messages and keep the latest. Fits state updates.
Send a resync message so the client refetches state. Fits after Lagged.
Close the connection and let the client reconnect. Fits stalled clients.
Wrap every send in a timeout. One client must never block a loop that serves others.
Cap message size (max_message_size), connections per user, and incoming message rate (resilience patterns).

Long-lived connections never end on their own. Graceful shutdown (health, shutdown and configuration) must signal connection tasks to close, and a rolling deploy (Docker, Kubernetes and CI) drops every connection, so clients need reconnect logic.

Complete example

Text

storefront-live/
  Cargo.toml
  src/
    main.rs
    auth.rs
    state.rs

Cargo.toml

TOML

[package]
name = "storefront-live"
version = "0.1.0"
edition = "2021"

[dependencies]
axum = { version = "0.8", features = ["ws"] }
tokio = { version = "1", features = ["full"] }
tokio-stream = { version = "0.1", features = ["sync"] }
futures-util = "0.3"
serde = { version = "1", features = ["derive"] }
serde_json = "1"

src/state.rs

Rust

use serde::Serialize;
use std::sync::{Arc, Mutex};
use tokio::sync::broadcast;

#[derive(Clone, Serialize)]
pub struct OrderEvent {
    pub id: u64,
    pub order_id: u64,
    pub status: String,
}

#[derive(Clone)]
pub struct Feed {
    log: Arc<Mutex<Vec<OrderEvent>>>,
    tx: broadcast::Sender<OrderEvent>,
}

impl Feed {
    pub fn new() -> Self {
        let (tx, _) = broadcast::channel(64);
        Feed { log: Arc::new(Mutex::new(Vec::new())), tx }
    }

    pub fn publish(&self, order_id: u64, status: String) -> OrderEvent {
        let mut log = self.log.lock().unwrap();
        let event = OrderEvent { id: log.len() as u64 + 1, order_id, status };
        log.push(event.clone());
        let _ = self.tx.send(event.clone());
        event
    }

    pub fn subscribe_from(&self, after: u64) -> (Vec<OrderEvent>, broadcast::Receiver<OrderEvent>) {
        let log = self.log.lock().unwrap();
        let receiver = self.tx.subscribe();
        let backlog = log.iter().filter(|event| event.id > after).cloned().collect();
        (backlog, receiver)
    }
}

src/auth.rs

Rust

use axum::http::{header, HeaderMap, StatusCode};
use std::{collections::HashMap, sync::Mutex};

const ALLOWED_ORIGINS: [&str; 1] = ["http://localhost:3000"];

pub fn check_origin(headers: &HeaderMap) -> Result<(), StatusCode> {
    let origin = headers
        .get(header::ORIGIN)
        .and_then(|value| value.to_str().ok())
        .ok_or(StatusCode::FORBIDDEN)?;
    if ALLOWED_ORIGINS.contains(&origin) {
        Ok(())
    } else {
        Err(StatusCode::FORBIDDEN)
    }
}

#[derive(Default)]
pub struct Tickets {
    inner: Mutex<(u64, HashMap<String, u64>)>,
}

impl Tickets {
    pub fn issue(&self, user_id: u64) -> String {
        let mut guard = self.inner.lock().unwrap();
        guard.0 += 1;
        let ticket = format!("t{}", guard.0);
        guard.1.insert(ticket.clone(), user_id);
        ticket
    }

    pub fn redeem(&self, ticket: &str) -> Option<u64> {
        self.inner.lock().unwrap().1.remove(ticket)
    }
}

src/main.rs

Rust

mod auth;
mod state;

use auth::{check_origin, Tickets};
use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Path, Query, State,
    },
    http::{HeaderMap, StatusCode},
    response::{
        sse::{Event, KeepAlive, Sse},
        Response,
    },
    routing::{get, post},
    Json, Router,
};
use futures_util::{stream::SplitSink, SinkExt, StreamExt};
use serde::Deserialize;
use state::{Feed, OrderEvent};
use std::{convert::Infallible, future::ready, sync::Arc, time::Duration};
use tokio::sync::broadcast::error::RecvError;
use tokio_stream::{wrappers::BroadcastStream, Stream};

#[derive(Clone)]
struct AppState {
    feed: Feed,
    tickets: Arc<Tickets>,
}

#[derive(Deserialize)]
struct StatusBody {
    status: String,
}

#[derive(Deserialize)]
struct TicketParams {
    user: u64,
}

#[derive(Deserialize)]
struct WsParams {
    ticket: String,
    after: Option<u64>,
}

async fn set_status(
    State(state): State<AppState>,
    Path(order_id): Path<u64>,
    Json(body): Json<StatusBody>,
) -> Json<OrderEvent> {
    Json(state.feed.publish(order_id, body.status))
}

async fn issue_ticket(
    State(state): State<AppState>,
    Query(params): Query<TicketParams>,
) -> Json<serde_json::Value> {
    Json(serde_json::json!({ "ticket": state.tickets.issue(params.user) }))
}

fn sse_event(event: &OrderEvent) -> Event {
    Event::default()
        .id(event.id.to_string())
        .event("order_status")
        .data(serde_json::to_string(event).unwrap())
}

async fn events(
    State(state): State<AppState>,
    headers: HeaderMap,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    let last_id: u64 = headers
        .get("last-event-id")
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.parse().ok())
        .unwrap_or(0);
    let (backlog, receiver) = state.feed.subscribe_from(last_id);
    let backlog = tokio_stream::iter(backlog.into_iter().map(|e| Ok(sse_event(&e))));
    let live = BroadcastStream::new(receiver)
        .take_while(|item| ready(item.is_ok()))
        .filter_map(|item| ready(item.ok()))
        .map(|e| Ok(sse_event(&e)));
    Sse::new(backlog.chain(live)).keep_alive(KeepAlive::default())
}

async fn ws_handler(
    upgrade: WebSocketUpgrade,
    headers: HeaderMap,
    Query(params): Query<WsParams>,
    State(state): State<AppState>,
) -> Result<Response, StatusCode> {
    check_origin(&headers)?;
    let _user_id = state
        .tickets
        .redeem(&params.ticket)
        .ok_or(StatusCode::UNAUTHORIZED)?;
    let after = params.after.unwrap_or(0);
    Ok(upgrade
        .max_message_size(64 * 1024)
        .on_upgrade(move |socket| handle_socket(socket, state.feed, after)))
}

async fn send_event(sender: &mut SplitSink<WebSocket, Message>, event: &OrderEvent) -> bool {
    let text = serde_json::to_string(event).unwrap();
    let sent = tokio::time::timeout(
        Duration::from_secs(5),
        sender.send(Message::Text(text.into())),
    )
    .await;
    matches!(sent, Ok(Ok(())))
}

async fn handle_socket(socket: WebSocket, feed: Feed, after: u64) {
    let (mut sender, mut receiver) = socket.split();
    let (backlog, mut live) = feed.subscribe_from(after);

    let mut send_task = tokio::spawn(async move {
        for event in backlog {
            if !send_event(&mut sender, &event).await {
                return;
            }
        }
        loop {
            match live.recv().await {
                Ok(event) => {
                    if !send_event(&mut sender, &event).await {
                        return;
                    }
                }
                Err(RecvError::Lagged(_)) => {
                    let _ = sender.send(Message::Text(r#"{"type":"resync"}"#.into())).await;
                    return;
                }
                Err(RecvError::Closed) => return,
            }
        }
    });

    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(message)) = receiver.next().await {
            if matches!(message, Message::Close(_)) {
                return;
            }
        }
    });

    tokio::select! {
        _ = &mut send_task => recv_task.abort(),
        _ = &mut recv_task => send_task.abort(),
    }
}

#[tokio::main]
async fn main() {
    let state = AppState {
        feed: Feed::new(),
        tickets: Arc::new(Tickets::default()),
    };
    let app = Router::new()
        .route("/orders/{id}/status", post(set_status))
        .route("/tickets", post(issue_ticket))
        .route("/events", get(events))
        .route("/ws", get(ws_handler))
        .with_state(state);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

Notes on the code:

Only futures_util::StreamExt is imported. Importing tokio_stream::StreamExt as well would make next, map, and chain ambiguous.
take_while ends the SSE stream on a lag error, and filter_map drops the wrapped Ok. The browser reconnects with Last-Event-ID, which replays the missed events from the log.

Run:

Terminal

cargo run

Requests (each in its own terminal, in this order):

Terminal

# terminal 1: SSE stream, stays open
curl -N http://127.0.0.1:3000/events

# terminal 2: publish two events
curl -s -X POST http://127.0.0.1:3000/orders/7/status \
  -H 'content-type: application/json' -d '{"status":"paid"}'
curl -s -X POST http://127.0.0.1:3000/orders/7/status \
  -H 'content-type: application/json' -d '{"status":"shipped"}'

Responses of the two POST calls:

Output

{"id":1,"order_id":7,"status":"paid"}
{"id":2,"order_id":7,"status":"shipped"}

Terminal 1 prints each event as it arrives (keep-alive comments appear every 15 seconds):

Output

id: 1
event: order_status
data: {"id":1,"order_id":7,"status":"paid"}

id: 2
event: order_status
data: {"id":2,"order_id":7,"status":"shipped"}

Resume after a drop: curl -N -H 'Last-Event-ID: 1' http://127.0.0.1:3000/events prints only event 2, then stays open.

WebSocket with a ticket and origin check:

Terminal

curl -s -X POST 'http://127.0.0.1:3000/tickets?user=5'
websocat -H 'Origin: http://localhost:3000' 'ws://127.0.0.1:3000/ws?ticket=t1&after=0'

Output

{"ticket":"t1"}
{"id":1,"order_id":7,"status":"paid"}
{"id":2,"order_id":7,"status":"shipped"}

Rejected handshakes return plain HTTP errors and never upgrade:

Terminal

# wrong origin: 403 Forbidden
websocat -H 'Origin: http://evil.example' 'ws://127.0.0.1:3000/ws?ticket=t1'
# reused or unknown ticket with a valid origin: 401 Unauthorized
websocat -H 'Origin: http://localhost:3000' 'ws://127.0.0.1:3000/ws?ticket=t1'

websocat reports the failed handshake with the status code. Ticket t1 is spent after the first successful connection, so the second command fails with 401. Request a new ticket for each connection.

Summary

A WebSocket starts as an HTTP request with Upgrade: websocket and the server answers 101 Switching Protocols.
WebSocketUpgrade with on_upgrade runs the message loop after the handshake. split() gives independent send and receive halves.
Browsers apply no same-origin policy to WebSockets, so the server compares Origin with an allow list.
Browser JavaScript cannot set headers on a WebSocket. Cookies with an origin check, or a single-use ticket, carry identity.
tokio::sync::broadcast fans events out within one process. Several instances need a shared pub/sub source.
SSE id fields and the Last-Event-ID header resume a stream. Subscribing and reading the backlog under one lock prevents gaps.
Bounded channels, send timeouts, message size limits, and connection limits handle slow clients.
Shutdown and deploys close long-lived connections, so clients must reconnect and resume.