Custom header connect-async tokio-tungstenite

help
1

Hello

How would I add custom-headers to tokio-tungstenite connect_async?

    let url = url::Url::parse(&*format!("wss://{}/order-dashboard/notifications", *config::URL)).unwrap();

    let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel();
    tokio::spawn(read_stdin(stdin_tx));

    let (ws_stream, _) = connect_async(url).await.expect("Failed to connect");
    println!("WebSocket handshake has been successfully completed");
2

It looks like connect_async_with_config lets you do it.

3

Got an example? The only documentation I can find is this: connect_async_with_config in tokio_tungstenite - Rust

4

You can use anything implementing the IntoClientRequest trait, including a http::Request directly. You can build your own request with any additional header you want using the Request builder.

5

Doing the following:

    let request = Request::builder()
        .method("GET")
        .uri(&*format!("wss://{}/order-dashboard/notifications", *config::URL))
        .body(())
        .unwrap();

    let (ws_stream, _) = connect_async(request).await.expect("Failed to connect");

gives:

thread 'tokio-runtime-worker' panicked at 'Failed to connect: Protocol(InvalidHeader("sec-websocket-key"))', src\main.rs:100:55

Doing the following (manually upgrading HTTP to Websocket):

    let request = Request::builder()
        .method("GET")
        .uri(&*format!("https://{}/order-dashboard/notifications", *config::URL))
        .header("Upgrade", "websocket")
        .header("Connection", "upgrade")
        .header("Sec-Websocket-Key", "key123")
        .header("Sec-Websocket-Version", "13")
        .body(())
        .unwrap();

    let (ws_stream, _) = connect_async(request).await.expect("Failed to connect");

gives: thread 'tokio-runtime-worker' panicked at 'Failed to connect: Url(UnsupportedUrlScheme)', src\main.rs:100:55 note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

6

This seems to work:

    let url_str = &*format!("wss://{}/order-dashboard/notifications", *config::URL);
    let url = url::Url::parse(url_str).unwrap();
    let host = url.host_str().expect("Invalid host in WebSocket URL");
   

    let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel();
    tokio::spawn(read_stdin(stdin_tx));

    let request = Request::builder()
        .method("GET")
        .uri(url_str)
        .header("Host", host)
        .header("Upgrade", "websocket")
        .header("Connection", "upgrade")
        .header("Sec-Websocket-Key", "key123")
        .header("Sec-Websocket-Version", "13")
        .body(())
        .unwrap();

    let (ws_stream, _) = connect_async(request).await.expect("Failed to connect");
    println!("WebSocket handshake has been successfully completed");

    let (mut write, mut read) = ws_stream.split();

    let stdin_to_ws = async {
        while let Some(message) = stdin_rx.recv().await {
            write.send(message).await.expect("Failed to send message to WebSocket");
        }
    };

    let ws_to_stdout = async {
        while let Some(message) = read.next().await {
            
            let message = message.unwrap();
            if message.is_empty() {
                continue;
            }

            let data = String::from_utf8(message.into_data()).unwrap();
            let json : Value = serde_json::from_str(&data).unwrap();

            let order = api::get_order(json["id"].as_u64().unwrap()).await;

            printer::print_receipt(order);
        }
    };

    tokio::select! {
        _ = stdin_to_ws => (),
        _ = ws_to_stdout => (),
    };
1 Like
7

Oops, I didn’t expect you would end up with something like that.

First, you’re probably aware, but I’ll mention it for posterity: the value for the header Sec-Websocket-Key should not be hardcoded this way. tungstenite provides the generate_key function to generate a key properly:

use tokio_tungstenite::tungstenite::handshake::client::generate_key;

let req = Request::builder()
     /* … */
    .header("Sec-WebSocket-Key", generate_key()) // <- like this
    /* … */;

Also, maybe I can suggest you to reuse the IntoClientRequest implementation for Url instead, and to use the headers_mut method to add your own headers:

use tokio_tungstenite::tungstenite::client::IntoClientRequest;

let url = url::Url::parse(/*…*/)?;

let mut req = url.into_client_request()?;

let headers = req.headers_mut();
headers.insert(/* … */); // <- insert an additional header

let _ = connect_async(req).await?;

This way, you don’t have to insert all the WebSocket-specific headers yourself.

1 Like
8

Oh thanks mate! That is a lot more elegant then my solution!
I think you by accident used let instead of use on the first line of your second code block though.

Otherwise your code works fine!

9

Haha, indeed! I edited to fix this.

1 Like