Repository: tokio-rs/axum
File: axum/src/extract/ws.rs (Rust)

Code before:
    ///
    /// This API is more flexible than [`protocols`](Self::protocols) and allows
    /// custom selection logic, such as Token-encoded subprotocols .
    ///
    /// # Examples
    ///
    /// ## Accept a fixed protocol with no allocation
    ///
    /// ```rust
    /// use http::HeaderValue;
    /// use axum::{
    ///     extract::ws::{WebSocketUpgrade, WebSocket},
    ///     routing::any,
    ///     response::{IntoResponse, Response},
    ///     Router,
    /// };
    ///
    /// let app = Router::new().route("/ws", any(handler));
    ///
    /// async fn handler(ws: WebSocketUpgrade) -> Response {
    ///     ws.select_protocol(|proto| {
    ///             (proto == "my-protocol")
    ///                 .then(|| HeaderValue::from_static("my-protocol"))
    ///         }).on_upgrade(|socket| async {
    ///             // ...
    ///         })
    /// }
    /// # let _: Router = app;
    /// ```
    ///
    /// ## Accept token based auth
    /// Suppose that the HTTP header is set as:
    /// ```txt
    /// Sec-WebSocket-Protocol: token.subprotocol.bearer.bXl0b2tlbgac7G, token.subprotocol.auth
    /// ```
    ///
    /// where the `bXl0b2tlbgac7G` is the credential token
    ///
    /// ```rust
    /// use http::HeaderValue;
    /// use axum::{
    ///     extract::ws::{WebSocketUpgrade, WebSocket},
    ///     routing::any,
    ///     response::{IntoResponse, Response},
    ///     Router,
    /// };
    ///
    /// let app = Router::new().route("/ws", any(handler));
    /// async fn handler(ws: WebSocketUpgrade) -> impl IntoResponse {
    ///     let ws = ws.select_protocol(|proto| {
    ///         proto

Review comment (suggestion):
```suggestion
    ///     ws.select_protocol(|proto| {
    ///         (proto == "my-protocol")
    ///             .then(|| HeaderValue::from_static("my-protocol"))
    ///     }).on_upgrade(|socket| async {
    ///         // ...
    ///     })
```

Code after:
            .and_then(|p| p.to_str().ok())
            .into_iter()
            .flat_map(|s| s.split(','))
            .map(|s| s.trim())
    }

    /// Set the chosen WebSocket subprotocol.
    ///
    /// Another method, [`protocols()`][Self::protocols], also sets the chosen WebSocket
    /// subprotocol. If both methods are called, only the latter call takes effect.
    ///
    /// # Notes
    ///
    /// - The chosen protocol is echoed back in the WebSocket upgrade
    ///   response as required by RFC 6455. Some browsers may reject a
    ///   value that was not present in the client's request.
    pub fn set_selected_protocol(&mut self, protocol: HeaderValue) {
        self.protocol = Some(protocol);
    }

    /// Return the selected WebSocket subprotocol, if one has been chosen.
    ///
    /// If [`protocols()`][Self::protocols] selects a matching protocol, or
    /// [`set_selected_protocol()`][Self::set_selected_protocol] has been called, the return
    /// value will be `Some` containing the selected 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
    /// will be called.
    ///
    /// By default any errors will be silently ignored.
    ///
    /// # Example
    ///
    /// ```
    /// use axum::{
    ///     extract::{WebSocketUpgrade},
    ///     response::Response,
    /// };
    ///
    /// async fn handler(ws: WebSocketUpgrade) -> Response {
    ///     ws.on_failed_upgrade(|error| {
    ///         report_error(error);
    ///     })
    ///     .on_upgrade(|socket| async { /* ... */ })
    /// }
    /// #