WebSocket Gateways

WebSocket connections in Toni are handled through gateways — structs that implement the GatewayTrait and handle connection lifecycle events and messages. Gateways are registered in the DI container like any other provider, so they have full access to injected services.

Enabling WebSocket support

Use the WebSocket adapter alongside your HTTP adapter:

use toni::ToniFactory;
use toni_axum::{AxumAdapter, AxumWebSocketAdapter};

#[tokio::main]
async fn main() {
    let mut app = ToniFactory::create(AppModule, AxumAdapter::new()).await;
    app.use_websocket_adapter(AxumWebSocketAdapter::new()).unwrap();
    app.listen(3000, "127.0.0.1").await;
}

The GatewayTrait

use toni::async_trait;
use toni::websocket::{GatewayTrait, WsClient, WsMessage};

#[async_trait]
pub trait GatewayTrait: Send + Sync {
    fn path(&self) -> &str;
    async fn on_connect(&self, client: &WsClient);
    async fn on_disconnect(&self, client: &WsClient, reason: DisconnectReason);
    async fn on_message(&self, client: &WsClient, message: WsMessage);
}

Basic chat gateway

use toni::{injectable, async_trait};
use toni::websocket::{GatewayTrait, WsClient, WsMessage, DisconnectReason};

#[injectable(pub struct ChatGateway {})]
impl ChatGateway {
    pub fn new() -> Self { Self {} }
}

#[async_trait]
impl GatewayTrait for ChatGateway {
    fn path(&self) -> &str {
        "/chat"
    }

    async fn on_connect(&self, client: &WsClient) {
        println!("Client {} connected", client.id());
        client.send(WsMessage::text("Welcome to the chat!")).await;
    }

    async fn on_disconnect(&self, client: &WsClient, reason: DisconnectReason) {
        println!("Client {} disconnected: {:?}", client.id(), reason);
    }

    async fn on_message(&self, client: &WsClient, message: WsMessage) {
        if let Some(text) = message.as_text() {
            println!("Message from {}: {}", client.id(), text);
            // Echo back
            client.send(WsMessage::text(format!("You said: {text}"))).await;
        }
    }
}

WsClient API

WsClient represents an individual connected client:

client.id()                          // unique client identifier
client.send(WsMessage).await         // send a message to this client
client.close().await                 // close the connection
client.get_metadata::<T>()           // retrieve typed metadata (set during handshake)

WsMessage types

WsMessage::text("hello")             // text frame
WsMessage::binary(vec![0u8, 1, 2])  // binary frame
WsMessage::ping()                    // ping control frame
WsMessage::pong()                    // pong control frame

Registering a gateway

Gateways are regular providers. List them in a module:

#[module(
    providers: [ChatGateway],
)]
pub struct ChatModule;

Toni automatically detects that ChatGateway implements GatewayTrait and registers it with the WebSocket adapter.

Authentication via handshake

Access the HTTP handshake request to authenticate connections before they're established:

use toni::websocket::{GatewayTrait, WsHandshake};

#[async_trait]
impl GatewayTrait for SecureGateway {
    fn path(&self) -> &str { "/secure" }

    async fn on_handshake(&self, handshake: &WsHandshake) -> bool {
        // Return false to reject the connection
        handshake.headers
            .get("authorization")
            .map(|h| self.auth_service.verify(h))
            .unwrap_or(false)
    }

    async fn on_connect(&self, client: &WsClient) { /* ... */ }
    async fn on_disconnect(&self, client: &WsClient, reason: DisconnectReason) { /* ... */ }
    async fn on_message(&self, client: &WsClient, message: WsMessage) { /* ... */ }
}