WebSockets with axum::extract::ws

1. Welcome and Icebreaker:

Have you ever used a real-time chat or live update feature on a website? What was it?

2. Outline the Session:

Agenda:

1. Concept Explanation: What are WebSockets and how does Axum support them?
2. Live Coding: Implement a WebSocket endpoint in Axum.
3. Tasks: Practice by building your own WebSocket handler.

Learning Outcomes:

• Understand WebSockets and their use cases.
• Implement a WebSocket endpoint using Axum.

Expectations: Try sending and receiving messages with your WebSocket server!

3. Concept Explanation:

WebSockets provide full-duplex communication between client and server. Axum supports WebSockets via the axum::extract::ws module, allowing you to build real-time features like chat or notifications.

4. Live Coding Demonstration:

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

async fn ws_handler(ws: WebSocketUpgrade) -> impl axum::response::IntoResponse {
    ws.on_upgrade(handle_socket)
}

async fn handle_socket(mut socket: WebSocket) {
    while let Some(Ok(msg)) = socket.recv().await {
        if let Message::Text(text) = msg {
            if socket.send(Message::Text(format!("Echo: {}", text).into())).await.is_err() {
                break; // Client disconnected.
            }
        }
    }
}

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

Enable the opt-in WebSocket feature: axum = { version = "0.8", features = ["ws"] }.

Best Practices:

• Handle errors and closed connections gracefully.
• Limit message size to prevent abuse.

Common Pitfalls:

• Not handling all message types (e.g., binary, close).
• Blocking the async runtime in the handler.

5. Tasks

• Build a simple chat server that broadcasts messages to all clients.
• Add authentication to your WebSocket endpoint.
• Try sending binary messages and handle them appropriately.