# How Buzz Uses Axum for Its WebSocket Server: A Deep Dive into the Relay Implementation

> Explore how Buzz implements its Nostr relay server using Axum. Discover its use of WebSocketUpgrade, Router API, and middleware for CORS and authentication in the buzz-relay crate.

- Repository: [Block Open Source/buzz](https://github.com/block/buzz)
- Tags: deep-dive
- Published: 2026-08-28

---

**Buzz leverages Axum's `WebSocketUpgrade` extractor and composable `Router` API to expose Nostr relay endpoints, handling real-time message streaming through the `buzz-relay` crate while applying unified middleware for CORS and authentication across all routes.**

The [block/buzz](https://github.com/block/buzz) repository implements a high-performance Nostr relay in Rust, utilizing the Axum web framework to power its WebSocket server and HTTP API. By integrating Axum's native WebSocket support, Buzz creates a cohesive routing system that manages persistent connections for event propagation alongside standard REST endpoints.

## Constructing the WebSocket Router in [`router.rs`](https://github.com/block/buzz/blob/main/router.rs)

The foundation of Buzz's WebSocket server resides in [`crates/buzz-relay/src/router.rs`](https://github.com/block/buzz/blob/main/crates/buzz-relay/src/router.rs), where the `build_router` function assembles the Axum application.

This module defines the route hierarchy that mounts the WebSocket endpoint together with REST API routes, health checks, and Prometheus metrics. Structural comments at the top of the file (lines 1-6) outline the middleware stack and the order in which routes are composed.

```rust
// crates/buzz-relay/src/router.rs
let app = Router::new()
    .route("/ws", get(ws_handler))          // WebSocket upgrade endpoint
    .route("/events", post(event_handler))  // NIP-29 POST endpoint
    .layer(cors_layer)                      // CORS middleware applied uniformly
    .with_state(app_state);

```

The router uses Axum's method-specific routing to distinguish between HTTP `GET` requests—which trigger WebSocket upgrades—and other verbs. By mounting the `/ws` path with `get(ws_handler)`, Buzz ensures that only GET requests initiate the WebSocket handshake while maintaining the ability to reject improper methods with appropriate status codes.

## Handling WebSocket Upgrades in [`connection.rs`](https://github.com/block/buzz/blob/main/connection.rs)

When a client connects to `/ws`, Axum dispatches the request to the handler defined in [`crates/buzz-relay/src/connection.rs`](https://github.com/block/buzz/blob/main/crates/buzz-relay/src/connection.rs). This module contains the core logic for managing WebSocket lifecycles and converting raw frames into application events.

### The Upgrade Pattern

The handler uses `axum::extract::ws::WebSocketUpgrade` to perform the protocol switch. According to the imports at lines 9-12, the module brings in `axum::extract::ws::Message` to handle incoming and outgoing frames:

```rust
// crates/buzz-relay/src/connection.rs
async fn ws_handler(
    ws: axum::extract::ws::WebSocketUpgrade,
    State(state): State<Arc<AppState>>,
) -> impl IntoResponse {
    ws.on_upgrade(|socket| handle_socket(socket, state))
}

```

The `on_upgrade` callback receives the `WebSocket` struct, which implements `Stream` and `Sink` for asynchronous message processing. Inside `handle_socket`, Buzz converts these `axum::extract::ws::Message` variants—Text, Binary, Close, Ping, and Pong—into internal relay events before dispatching them to the subscription management system.

## WebSocket State Management and Protocol Compliance

Robust error handling requires explicit close codes and message type definitions, which Buzz centralizes in [`crates/buzz-relay/src/state.rs`](https://github.com/block/buzz/blob/main/crates/buzz-relay/src/state.rs).

### Close Codes and Shutdown Signals

Lines 52-53 and 511-513 of [`state.rs`](https://github.com/block/buzz/blob/main/state.rs) define specific WebSocket close codes including `POLICY` and `RESTART`. These constants allow the relay to communicate shutdown reasons to clients according to the WebSocket specification:

- **POLICY**: Indicates the connection violated relay policy (rate limiting, banned content)
- **RESTART**: Signals an administrative restart or graceful shutdown

By standardizing these codes in the state module, Buzz ensures consistent behavior across the connection lifecycle, whether closing due to policy enforcement or server maintenance.

## Launching the Server with `axum::serve`

The final assembly occurs in [`crates/buzz-relay/src/main.rs`](https://github.com/block/buzz/blob/main/crates/buzz-relay/src/main.rs), where the constructed router binds to TCP listeners and enters the async runtime.

At lines 1273, 1369, 1378, and 1403, the codebase calls `axum::serve` to spawn the server. This function drives both the primary listener (serving the WebSocket and HTTP endpoints) and separate health-probe listeners for container orchestration:

```rust
// crates/buzz-relay/src/main.rs
let router = build_router(state.clone());
axum::serve(listener, router).await?;

```

Using `axum::serve` unifies the connection handling for both WebSocket upgrades and standard HTTP requests under a single hyper-powered server. This approach allows Buzz to share middleware—such as authentication layers and request tracing—across all endpoints without duplicating logic.

## Summary

- **Buzz uses `axum::extract::ws::WebSocketUpgrade`** in [`crates/buzz-relay/src/connection.rs`](https://github.com/block/buzz/blob/main/crates/buzz-relay/src/connection.rs) to handle protocol upgrades from HTTP to WebSocket.
- **The router in [`crates/buzz-relay/src/router.rs`](https://github.com/block/buzz/blob/main/crates/buzz-relay/src/router.rs)** composes WebSocket routes with REST endpoints using `Router::new().route("/ws", get(...))`, enabling unified middleware application.
- **Close codes like POLICY and RESTART** defined in [`crates/buzz-relay/src/state.rs`](https://github.com/block/buzz/blob/main/crates/buzz-relay/src/state.rs) (lines 52-53 and 511-513) provide standardized connection termination semantics for Nostr clients.
- **`axum::serve` in [`crates/buzz-relay/src/main.rs`](https://github.com/block/buzz/blob/main/crates/buzz-relay/src/main.rs)** (lines 1273, 1369, 1378, 1403) launches the server, driving both WebSocket streams and health check endpoints through a single async runtime.

## Frequently Asked Questions

### What Axum extractors does Buzz use for WebSocket handling?

Buzz relies on `axum::extract::ws::WebSocketUpgrade` to initiate the protocol handshake and `axum::extract::ws::WebSocket` (obtained via `on_upgrade`) for the actual message stream. The `State` extractor injects the shared application state into handlers, allowing access to relay configuration and subscription managers.

### How does Buzz handle WebSocket message types?

Incoming frames are received as `axum::extract::ws::Message` enums in [`crates/buzz-relay/src/connection.rs`](https://github.com/block/buzz/blob/main/crates/buzz-relay/src/connection.rs). The handler pattern-matches on these variants—distinguishing between Text, Binary, Ping, Pong, and Close frames—before converting them to internal relay events for processing by the Nostr event pipeline.

### Where is the WebSocket route defined in the Buzz codebase?

The route definition resides in [`crates/buzz-relay/src/router.rs`](https://github.com/block/buzz/blob/main/crates/buzz-relay/src/router.rs) within the `build_router` function. The WebSocket endpoint mounts at the `/ws` path using Axum's `get()` handler, which differentiates upgrade requests from standard HTTP methods on the same route.

### Can Buzz apply middleware to WebSocket connections?

Yes. Because Axum applies middleware at the router level, Buzz can layer authentication, CORS, and tracing middleware in [`router.rs`](https://github.com/block/buzz/blob/main/router.rs) using `.layer()`. These middleware components execute for all routes—including `/ws`—ensuring consistent request validation before the WebSocket upgrade occurs.