Creating WebSocket Servers in Axum: A Hands-On Rust Guide

Amza

Amza

— 4 min read

If you're building real-time applications in Rust, you might have come across Axum, the powerful web framework built on top of hyper and tokio. While Axum is known for its simplicity and performance, one thing that might seem tricky is integrating WebSockets, a key technology for creating real-time, bidirectional communication between clients and servers.

In this blog, we will explore how to implement WebSocket support in Axum from the ground up. By the end of this article, you'll be able to handle WebSocket connections in your Axum-based Rust application and understand how the underlying WebSocket protocol works.

Let’s dive in!

Step 1: Setting Up Your Rust Project

First, let's start by setting up a new Rust project.

cargo new axum_websocket_example --bin
cd axum_websocket_example

Now, we need to add dependencies to your Cargo.toml file for Axum:

[dependencies]
axum = { version = "0.6", features = ["ws"] }
tokio = { version = "1", features = ["full"] }

The axum crate now includes WebSocket support natively, so we no longer need an external WebSocket library like tokio-tungstenite. The WebSocket functionality is built into Axum's core, making it easier to handle WebSocket connections directly.

Step 2: Writing the Basic Axum Server

Let’s get a basic server running using Axum. Start by creating a simple route in main.rs that will serve our WebSocket endpoint.

use axum::{routing::get, Router};
use std::net::SocketAddr;

#[tokio::main]
async fn main() {
    // Initialize the Axum router
    let app = Router::new().route("/", get(root_handler));

    // Specify the address to bind to
    let addr = SocketAddr::from(([127, 0, 0, 1], 8080));

    // Start the Axum server
    println!("Server running at {}", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn root_handler() -> &'static str {
    "Hello, world!"
}

This simple program creates an HTTP server using Axum that listens on localhost:8080 and responds with "Hello, world!" when accessed. To check if everything works so far, run:

cargo run

Now, navigate to http://localhost:8080 in your browser, and you should see the “Hello, world!” message.

Step 3: Introducing WebSocket Support

Now that we have our basic Axum app running, it’s time to introduce WebSocket support. First, we need to create a route for WebSocket connections and upgrade the HTTP connection to a WebSocket connection.

In Axum, you can upgrade an HTTP connection using the ws method from the axum::extrcat::ws. Let's define this route in our main.rs.

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

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

async fn handle_socket(mut socket: WebSocket) {
    // Send a greeting message to the client
    if let Err(e) = socket
        .send(Message::Text("Hello from the server!".to_string()))
        .await
    {
        eprintln!("Error sending message: {}", e);
        return;
    }

    // Loop to keep the connection alive
    while let Some(Ok(msg)) = socket.recv().await {
        match msg {
            Message::Text(msg) => {
                println!("Received message: {}", msg);
                if let Err(e) = socket.send(Message::Text(format!("Echo: {}", msg))).await {
                    eprintln!("Error sending message: {}", e);
                }
            }
            Message::Close(_) => {
                println!("Closing WebSocket connection.");
                break;
            }
            _ => {}
        }
    }
}

Step 4: Updating the Axum Router

Now we need to update our router to include this new WebSocket handler. Modify the main function as follows:

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/", get(root_handler))
        .route("/ws", get(websocket_handler));  // Add WebSocket route

    let addr = SocketAddr::from(([127, 0, 0, 1], 8080));
    println!("Server running at {}", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

Now your Axum server can handle WebSocket connections at /ws.

Step 5: Testing the WebSocket

To test your WebSocket, you can use a WebSocket client, such as the browser’s developer tools console or an external WebSocket testing tool.

For the browser:

  1. Open the developer tools (F12).
  2. Go to the Console tab.
  3. Paste this JavaScript code into the console:
const socket = new WebSocket('ws://localhost:8080/ws');

socket.onopen = () => {
    console.log('Connected to server!');
    socket.send('Hello, server!');
};

socket.onmessage = (event) => {
    console.log('Received from server:', event.data);
};

You should see the connection message and an echoed message back from the server.

Concepts and Explanations

  • WebSocket Protocol: A WebSocket connection allows for bidirectional communication between the server and client over a single, long-lived connection. This is especially useful for applications like real-time chats, live feeds, or notifications.

  • WebSocketUpgrade: This is part of the axum::extract::ws module and is used to upgrade an HTTP connection to a WebSocket connection. This is essential for handling WebSocket requests in Axum.

  • Message Types: We used Message::Text to send and receive textual messages. WebSocket messages can also be binary (Message::Binary), or you can handle closing the connection (Message::Close).

  • Tokio and Asynchronous Programming: The server runs asynchronously using the tokio runtime. WebSocket communication requires asynchronous code because sending and receiving data over WebSockets is inherently non-blocking.

Challenges or Questions

  1. Modify the WebSocket handler to send a "Goodbye" message when the connection is closed. Try to handle any errors gracefully.
  2. How would you implement a WebSocket client in a Rust application? Could you create a simple client that connects to your server and sends messages?

Recap and Conclusion

Congratulations! You've just built a real-time WebSocket server in Rust using Axum. We started with a basic HTTP server and gradually added WebSocket support using axum::extract::ws. Here's a quick recap:

  • We set up a basic Axum server.
  • We added a WebSocket handler using axum::extract::ws.
  • We upgraded HTTP connections to WebSockets and sent/received messages.

If you're interested in diving deeper into Axum and WebSockets, consider exploring more advanced features like broadcasting messages to multiple clients or managing WebSocket connections in a more structured way.

Happy coding!

Read more