Use a single handler implementation for websocket opened by both the server or a client

[niconicoj]

I am currently running into issues where my server needs to both accept and open websocket connection to other servers. However, whether a WebSocket is created by the client or the server bears no impact on the logic that goes into handling the sockets, therefore I would very much like to only write a single implementation of that logic that would work for both.

here is an example of what I would like to be able to write :

// entry point for a connection initiated by a client
pub async fn handle_ws_connection(
    ws: WebSocketUpgrade
) -> impl IntoResponse {
    ws.on_upgrade(|socket| handle_socket(socket))
}

// entry point for a connection initiated by the server
pub async fn peer_connect(peer: String) {
    let (mut socket, _) = connect_async(peer)
        .await
        .expect(&format!("failed to connect to peer {}", &peer));
    handle_socket(socket).await;
}

// handler for sockets
async fn handle_socket(mut socket: WebSocket) {
    //...here goes some code to handle a websocket
}

But that does not work because connect_async returns a WebSocketStream from the tokio_tungestenite crate and the on_upgrade provides a WebSocket from the axum crate.

Is there currently a way to do what I am trying to do ?

[davidpdrsn]

Not sure it's possible but you can try making the function generic over a Stream + Sink.

Unfortunately we cannot expose the underlying tokio_tungstenite stream from axum since that makes tokio_tungstenite part of our public API which we don't want.

[niconicoj]

I understand not wanting to expose tokio_tungestenite through the public API. Thing is I kinda like the current WebSocket API so I don't want to throw that out just yet 😋.

What about making WebSocket generic over AsyncRead + AsyncWrite + Unpin ? Or exposing a trait for the WebSocket API that user could implements for their own types ?

[davidpdrsn]

That sounds a lot like Sink + Stream to me 😛 that is an API that users can drop down to instead.