Pipe bytes between the WebSocket and an in-memory `MultiplexService` stream.
(
ws: axum::extract::ws::WebSocket,
state: Arc<ServerState>,
)
| 52 | |
| 53 | /// Pipe bytes between the WebSocket and an in-memory `MultiplexService` stream. |
| 54 | async fn handle_ws_tunnel( |
| 55 | ws: axum::extract::ws::WebSocket, |
| 56 | state: Arc<ServerState>, |
| 57 | ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> { |
| 58 | let service = crate::MultiplexService::new(state); |
| 59 | let (tunnel_stream, service_stream) = tokio::io::duplex(64 * 1024); |
| 60 | debug!("WS tunnel: spawned in-memory multiplex connection"); |
| 61 | |
| 62 | let (ws_sink, ws_source) = ws.split(); |
| 63 | let (tunnel_read, tunnel_write) = tokio::io::split(tunnel_stream); |
| 64 | |
| 65 | let service_task = tokio::spawn(async move { |
| 66 | if let Err(e) = service.serve(service_stream).await { |
| 67 | debug!(error = %e, "WS tunnel: multiplex service error"); |
| 68 | } |
| 69 | }); |
| 70 | let mut tunnel_to_ws = tokio::spawn(copy_reader_to_ws(tunnel_read, ws_sink)); |
| 71 | let mut ws_to_tunnel = tokio::spawn(copy_ws_to_writer(ws_source, tunnel_write)); |
| 72 | |
| 73 | tokio::select! { |
| 74 | res = &mut tunnel_to_ws => { |
| 75 | if let Ok(Err(e)) = res { |
| 76 | debug!(error = %e, "WS tunnel: tunnel->ws error"); |
| 77 | } |
| 78 | ws_to_tunnel.abort(); |
| 79 | } |
| 80 | res = &mut ws_to_tunnel => { |
| 81 | if let Ok(Err(e)) = res { |
| 82 | debug!(error = %e, "WS tunnel: ws->tunnel error"); |
| 83 | } |
| 84 | tunnel_to_ws.abort(); |
| 85 | } |
| 86 | } |
| 87 | service_task.abort(); |
| 88 | |
| 89 | Ok(()) |
| 90 | } |
| 91 | |
| 92 | /// Copy bytes from an async reader into WebSocket binary frames. |
| 93 | async fn copy_reader_to_ws<R>( |
no test coverage detected