diff --git a/axum/src/extract/ws.rs b/axum/src/extract/ws.rs index 2f367e43..11ed5740 100644 --- a/axum/src/extract/ws.rs +++ b/axum/src/extract/ws.rs @@ -265,6 +265,15 @@ impl WebSocketUpgrade { self } + /// Return the selected WebSocket subprotocol, if one has been chosen. + /// + /// If [`protocols()`][Self::protocols] has been called and a matching + /// protocol has been selected, the return value will be `Some` containing + /// said protocol. Otherwise, it will be `None`. + pub fn selected_protocol(&self) -> Option<&HeaderValue> { + self.protocol.as_ref() + } + /// Provide a callback to call if upgrading the connection fails. /// /// The connection upgrade is performed in a background task. If that fails this callback @@ -1154,6 +1163,7 @@ mod tests { fn echo_app() -> Router { async fn handle_socket(mut socket: WebSocket) { + assert_eq!(socket.protocol().unwrap(), "echo"); while let Some(Ok(msg)) = socket.recv().await { match msg { Message::Text(_) | Message::Binary(_) | Message::Close(_) => { @@ -1171,7 +1181,9 @@ mod tests { Router::new().route( "/echo", any(|ws: WebSocketUpgrade| { - ready(ws.protocols(["echo2", "echo"]).on_upgrade(handle_socket)) + let ws = ws.protocols(["echo2", "echo"]); + assert_eq!(ws.selected_protocol().unwrap(), "echo"); + ready(ws.on_upgrade(handle_socket)) }), ) }