Warp | 22 , WebSocket | Chat

Let’s implement WebSockets in Warp to build a basic chat application.

(Youtube Video )

in cargo.toml :

[dependencies]
tokio = { version = "1", features = ["full"] }
warp = { version = "0.4", features=["server","websocket"]}
futures = "0.3"

We add futures so that we can split the WebSocket into a sender and a receiver. We also enable the websocket feature in Warp.

Get Mike Code’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

in main.rs

use futures::{SinkExt, StreamExt};
use warp::{Filter, filters::ws::Message};

#[tokio::main]
async fn main() {
    let (tx, _) = tokio::sync::broadcast::channel(100);
    let app = app(tx).await;
    
    warp::serve(app)
        .run(([0,0,0,0], 3000))
        .await;
}

async fn app(tx: tokio::sync::broadcast::Sender<String>) -> impl Filter<Extract = impl warp::Reply, Error = warp::Rejection> + Clone {
    let root = warp::path!().map(|| "Home");

    let share_state = warp::any().map(move || tx.clone());

    let ws = warp::path("ws").and(warp::get()).and(warp::ws()).and(share_state).map(|ws:warp::ws::Ws, tx: tokio::sync::broadcast::Sender<String>| ws.on_upgrade(move |socket| ws_handler(socket, tx)));

    root
        .or(ws)
}

async fn ws_handler (ws: warp::ws::WebSocket, tx: tokio::sync::broadcast::Sender<String>)  {
    let (mut ws_sender, mut ws_receiver) = ws.split();


    let mut rx = tx.subscribe();

    tokio::spawn(async move{
        while let  Ok(msg) = rx.recv().await {
            ws_sender.send(Message::text(msg)).await.unwrap();
        }
    });

    while let Some(Ok(msg)) = ws_receiver.next().await {
        if msg.is_text() {
            let msg = msg.to_str().unwrap().to_string();
            tx.send(msg).unwrap();
        }
    }

}

We use Tokio to create a broadcast channel, which we use to broadcast messages to every connected user.

We define an app function to create the Warp router. In this function, we pass the broadcast channel’s sender (tx) as shared state, so it can be accessed in the WebSocket route handler.

In the WebSocket route, we use the and method to attach the shared state (tx) and warp::ws() to the route. The warp::ws() filter checks whether the incoming connection is a WebSocket connection. If it is, we use the map method and pass a closure that captures both the shared state (tx) and the WebSocket connection. Inside this closure, we call ws.on_upgrade, passing another closure that captures the socket and forwards the socket and tx to the ws_handler function.

We define a ws_handler function, where we call split on the socket to obtain a WebSocket sender and receiver. We then call subscribe on tx to create a receiver for the broadcast channel, so each connected client can receive messages broadcast on this channel.

Using Tokio, we spawn an asynchronous background task that continuously receives messages from the broadcast channel in a while loop. This ensures the task does not block the main execution flow. Whenever a message is received, we use the WebSocket sender to forward it to the connected client. In this way, every connected user receives broadcast messages.

At the same time, we run another while loop that continuously receives messages from the WebSocket receiver. Each incoming message is sent into the broadcast channel.

Overall, the flow works as follows: when a user sends a message through the WebSocket, the message is first received by the server (the user himself), then sent into the broadcast channel. The channel broadcasts the message to all subscribed receivers, and each connected client sends the message back to its WebSocket connection. As a result, all users — including the sender — receive the message in real time.