Get tokio_tungstenite::WebSocketStream from axum::WebSocket

[mdeloof]

Hi

I was wondering if there is a way to get a tokio_tungstenite::WebSocketStream from a axum::WebSocket? If I understand correctly from looking at the source code, the latter is a wrapper around the former? But I couldn't find a way to access the inner WebSocketStream. Or would I be better off writing a custom extractor in this case?

Thanks! :)

[davidpdrsn]

That's intentionally not possible as that would make tungstenite part of axum's public api, which we don't want.

I'd recommend abstracting over the Stream and Sink trait instead if possible.

[mdeloof]

Ah okay!

I am trying to make it generic over Sink and Stream, but I'm not sure what type Item should be set to as Message is not the same for axum and tungstenite?

[davidpdrsn]

You should be able to copy this function into your code to map from axum's message type to tungstenite's.

[rlidwka]

You should be able to copy this function into your code to map from axum's message type to tungstenite's.

your link is unfortunately no longer valid because it wasn't pinned to a specific commit

anyways, here's what I ended up with while facing similar issue:

use futures::{SinkExt, StreamExt};

fn axum_to_tungstenite(message: Result<axum::extract::ws::Message, axum::Error>) -> Result<tungstenite::Message, tungstenite::Error> {
    match message {
        Ok(msg) => Ok(match msg {
            axum::extract::ws::Message::Text(text) => tungstenite::Message::Text(text),
            axum::extract::ws::Message::Binary(bin) => tungstenite::Message::Binary(bin),
            axum::extract::ws::Message::Ping(ping) => tungstenite::Message::Ping(ping),
            axum::extract::ws::Message::Pong(pong) => tungstenite::Message::Pong(pong),
            axum::extract::ws::Message::Close(close) => tungstenite::Message::Close(close.map(|c| {
                tungstenite::protocol::CloseFrame {
                    code: c.code.into(),
                    reason: c.reason,
                }
            })),
        }),
        Err(_err) => Err(tungstenite::Error::ConnectionClosed),
    }
}

fn tungstenite_to_axum(message: tungstenite::Message) -> Pin<Box<dyn Future<Output = Result<axum::extract::ws::Message, axum::Error>> + Send + Sync>> {
    Box::pin(async {
        Ok(match message {
            tungstenite::Message::Text(text) => axum::extract::ws::Message::Text(text),
            tungstenite::Message::Binary(bin) => axum::extract::ws::Message::Binary(bin),
            tungstenite::Message::Ping(ping) => axum::extract::ws::Message::Ping(ping),
            tungstenite::Message::Pong(pong) => axum::extract::ws::Message::Pong(pong),
            tungstenite::Message::Frame(_frame) => unreachable!(),
            tungstenite::Message::Close(close) => axum::extract::ws::Message::Close(close.map(|c| {
                axum::extract::ws::CloseFrame {
                    code: c.code.into(),
                    reason: c.reason,
                }
            })),
        })
    })
}

async fn ws_handler(ws: WebSocketUpgrade) -> impl IntoResponse {
    ws.on_upgrade(|socket| async {
        serve_tungstenite_stream(
            socket.map(axum_to_tungstenite).with(tungstenite_to_axum),
        ).await.unwrap();
    })
}