fastwebsockets

To use fastwebsockets with the Axum web framework, enable the upgrade and with_axum features in your Cargo.toml:

fastwebsockets = { version = "0.9.0", features = ["upgrade", "with_axum"] }

In your Axum handler, use upgrade::IncomingUpgrade to receive the upgrade request. Call .upgrade() on the incoming upgrade to get a response and an UpgradeFut. You can then wrap the resulting stream in a FragmentCollector to handle full messages easily.

use axum::{response::IntoResponse, routing::get, Router};
use fastwebsockets::upgrade;
use fastwebsockets::OpCode;
use fastwebsockets::WebSocketError;

#[tokio::main]
async fn main() {
  let app = Router::new().route("/", get(ws_handler));

  let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
  axum::serve(listener, app).await.unwrap();
}

async fn handle_client(fut: upgrade::UpgradeFut) -> Result<(), WebSocketError> {
  let mut ws = fastwebsockets::FragmentCollector::new(fut.await?);

  loop {
    let frame = ws.read_frame().await?;
    match frame.opcode {
      OpCode::Close => break,
      OpCode::Text | OpCode::Binary => {
        ws.write_frame(frame).await?;
      }
      _ => {}
    }
  }

  Ok
}

async fn ws_handler(ws: upgrade::IncomingUpgrade) -> impl IntoResponse {
  let (response, fut) = ws.upgrade().unwrap();

  tokio::task::spawn(async move {
    if let Err(e) = handle_client(fut).await {
      eprintln!("Error in websocket connection: {}", e);
    }
  });

  response
}

To connect to a WebSocket server as a client, use the fastwebsockets::handshake module.

1. Construct a hyper::Request with the necessary WebSocket headers:
   • UPGRADE: "websocket"
   • CONNECTION: "upgrade"
   • Sec-WebSocket-Key: Generated via handshake::generate_key()
   • Sec-WebSocket-Version: "13"
2. Use handshake::client to perform the handshake. This requires an executor (e.g., a SpawnExecutor that wraps tokio::task::spawn) to manage the asynchronous tasks.

use fastwebsockets::handshake;
use fastwebsockets::WebSocket;
use hyper::{Request, body::Bytes, upgrade::Upgraded, header::{UPGRADE, CONNECTION}};
use http_body_util::Empty;
use tokio::net::TcpStream;
use std::future::Future;
use anyhow::Result;

async fn connect() -> Result<WebSocket<Upgraded>> {
  let stream = TcpStream::connect("localhost:9001").await?;

  let req = Request::builder()
    .method("GET")
    .uri("http://localhost:9001/")
    .header("Host", "localhost:9001")
    .header(UPGRADE, "websocket")
    .header(CONNECTION, "upgrade")
    .header(
      "Sec-WebSocket-Key",
      handshake::generate_key(),
    )
    .header("Sec-WebSocket-Version", "13")
    .body(Empty::<Bytes>::new())?;

  let (ws, _) = handshake::client(&SpawnExecutor, req, stream).await?;
  Ok(ws)
}

// Tie hyper's executor to tokio runtime
struct SpawnExecutor;

impl<Fut> hyper::rt::Executor<Fut> for SpawnExecutor
where
  Fut: Future + Send + 'static,
  Fut::Output: Send + 'static,
{
  fn execute(&self, fut: Fut) {
    tokio::task::spawn(fut);
  }
}

The FragmentCollector is a wrapper for a WebSocket that automatically handles fragmented messages. Instead of receiving multiple individual frames for a single logical message, FragmentCollector buffers the fragments in memory and returns a single, complete Frame once the final fragment (with the fin bit set) is received.

Key behaviors:

• Memory Usage: The entire message payload is buffered in memory until completion. Use this only when streaming is not an option or messages are reasonably sized.
• UTF-8 Validation: For OpCode::Text messages, the collector ensures the resulting payload is valid UTF-8.
• Automatic Management: It handles control frames (like Pongs) and ensures the connection state remains consistent during collection.

To use it, wrap your WebSocket instance with FragmentCollector::new(ws) and call read_frame() instead of reading directly from the raw WebSocket.

use tokio::net::TcpStream;
use fastwebsockets::{WebSocket, FragmentCollector, OpCode, Role};
use anyhow::Result;

async fn handle_client(
  socket: TcpStream,
) -> Result<()> {
  let ws = WebSocket::after_handshake(socket, Role::Server);
  let mut ws = FragmentCollector::new(ws);

  loop {
    let frame = ws.read_frame().await?;
    match frame.opcode {
      OpCode::Close => break,
      OpCode::Text | OpCode::Binary => {
        ws.write_frame(frame).await?;
      }
      _ => {}
    }
  }
  Ok(())
}

You can use fastwebsockets as a full-fledged WebSocket server. After performing a handshake on a TcpStream, you can wrap the socket in a WebSocket instance.

Key configuration methods for the WebSocket instance include:

• set_writev(bool): Enables vectored writes.
• set_auto_close(bool): Automatically handles closing.
• set_auto_pong(bool): Automatically responds to PING frames with PONG.

To handle messages, loop over ws.read_frame().await? and match on the OpCode.

use fastwebsockets::{Frame, OpCode, WebSocket};
use tokio::net::TcpStream;

async fn handle_client(
  mut socket: TcpStream,
) -> Result<(), WebSocketError> {
  handshake(&mut socket).await?;

  let mut ws = WebSocket::after_handshake(socket);
  ws.set_writev(true);
  ws.set_auto_close(true);
  ws.set_auto_pong(true);

  loop {
    let frame = ws.read_frame().await?;

    match frame {
      OpCode::Close => break,
      OpCode::Text | OpCode::Binary => {
        let frame = Frame::new(true, frame.opcode, None, frame.payload);
        ws.write_frame(frame).await?;
      }
    }
  }

  Ok(())
}

You can convert between raw u16 status codes and the CloseCode enum using the From trait.

u16 to CloseCode

Converting a u16 to a CloseCode maps standard WebSocket status codes to their respective enum variants:

• 1000 -> Normal
• 1001 -> Away
• 1002 -> Protocol
• 1003 -> Unsupported
• 1005 -> Status
• 1006 -> Abnormal
• 1007 -> Invalid
• 1008 -> Policy
• 1009 -> Size
• 1010 -> Extension
• 1011 -> Error
• 1012 -> Restart
• 1013 -> Again
• 1015 -> Tls
• 1..=999 -> Bad(code)
• 1016..=2999 -> Reserved(code)
• 3000..=3999 -> Iana(code)
• 4000..=4999 -> Library(code)
• All other values -> Bad(code)

CloseCode to u16

Converting a CloseCode back to a u16 returns the corresponding numeric status code used in the WebSocket protocol.

// Convert u16 to CloseCode
let code = CloseCode::from(1000); // Returns CloseCode::Normal

// Convert CloseCode to u16
let numeric_code: u16 = CloseCode::Normal.into(); // Returns 1000

If the with_axum feature is enabled, you can use IncomingUpgrade as an Axum extractor (FromRequestParts). This simplifies the handshake process within an Axum handler.

IncomingUpgrade automatically:

• Extracts the Sec-WebSocket-Key.
• Validates the Sec-WebSocket-Version is 13.
• Retrieves the hyper::upgrade::OnUpgrade extension.
• Generates the correct Sec-WebSocket-Accept header for the response.

// Example Axum handler using IncomingUpgrade
async fn ws_handler(
    upgrade: IncomingUpgrade,
) -> impl axum::response::IntoResponse {
    // 1. Generate the response using the helper
    let (response, upgrade_fut) = upgrade.upgrade().unwrap();
    
    // 2. Spawn a task to handle the actual websocket connection
    tokio::spawn(async move {
        let ws = upgrade_fut.await.unwrap();
        // use ws...
    });

    // 3. Return the response to the client
    response
}