Rewrite sec-websocket-protocol handling (#3620)

Author: jplatte

.github/workflows/CI.yml

@@ -2,7 +2,7 @@ name: CI
 
 env:
   CARGO_TERM_COLOR: always
-  MSRV: '1.78'
+  MSRV: '1.80'
 
 on:
   push:

Cargo.toml

@@ -3,7 +3,7 @@ members = ["axum", "axum-*"]
 resolver = "2"
 
 [workspace.package]
-rust-version = "1.78"
+rust-version = "1.80"
 
 [workspace.lints.rust]
 unsafe_code = "forbid"

axum/CHANGELOG.md

@@ -16,8 +16,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   (because it was already never terminating if that method wasn't used) ([#3601])
 - **added:** New `ListenerExt::limit_connections` allows limiting concurrent `axum::serve` connections ([#3489])
 - **added:** `MethodRouter::method_filter` ([#3586])
+- **added:** `WebSocketUpgrade::{requested_protocols, set_selected_protocol}` for more
+  flexible subprotocol selection ([#3597])
 - **changed:** `serve` has an additional generic argument and can now work with any response body
   type, not just `axum::body::Body` ([#3205])
+- **changed:** Update minimum rust version to 1.80 ([#3620])
 
 [#3158]: https://github.com/tokio-rs/axum/pull/3158
 [#3261]: https://github.com/tokio-rs/axum/pull/3261
@@ -26,6 +29,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [#3601]: https://github.com/tokio-rs/axum/pull/3601
 [#3489]: https://github.com/tokio-rs/axum/pull/3489
 [#3586]: https://github.com/tokio-rs/axum/pull/3586
+[#3597]: https://github.com/tokio-rs/axum/pull/3597
+[#3620]: https://github.com/tokio-rs/axum/pull/3620
 
 # 0.8.8
 

axum/README.md

@@ -111,7 +111,7 @@ This crate uses `#![forbid(unsafe_code)]` to ensure everything is implemented in
 
 ## Minimum supported Rust version
 
-axum's MSRV is 1.78.
+axum's MSRV is 1.80.
 
 ## Examples
 

axum/src/extract/ws.rs

@@ -106,8 +106,10 @@ use hyper_util::rt::TokioIo;
 use sha1::{Digest, Sha1};
 use std::{
     borrow::Cow,
+    collections::BTreeSet,
     future::Future,
     pin::Pin,
+    str,
     task::{ready, Context, Poll},
 };
 use tokio_tungstenite::{
@@ -137,7 +139,7 @@ pub struct WebSocketUpgrade<F = DefaultOnFailedUpgrade> {
     sec_websocket_key: Option<HeaderValue>,
     on_upgrade: hyper::upgrade::OnUpgrade,
     on_failed_upgrade: F,
-    sec_websocket_protocol: Option<HeaderValue>,
+    sec_websocket_protocol: BTreeSet<HeaderValue>,
 }
 
 impl<F> std::fmt::Debug for WebSocketUpgrade<F> {
@@ -241,26 +243,23 @@ impl<F> WebSocketUpgrade<F> {
         I: IntoIterator,
         I::Item: Into<Cow<'static, str>>,
     {
-        if let Some(req_protocols) = self
-            .sec_websocket_protocol
-            .as_ref()
-            .and_then(|p| p.to_str().ok())
-        {
-            self.protocol = protocols
-                .into_iter()
-                // FIXME: This will often allocate a new `String` and so is less efficient than it
-                // could be. But that can't be fixed without breaking changes to the public API.
-                .map(Into::into)
-                .find(|protocol| {
-                    req_protocols
-                        .split(',')
-                        .any(|req_protocol| req_protocol.trim() == protocol)
-                })
-                .map(|protocol| match protocol {
-                    Cow::Owned(s) => HeaderValue::from_str(&s).unwrap(),
-                    Cow::Borrowed(s) => HeaderValue::from_static(s),
-                });
-        }
+        self.protocol = protocols
+            .into_iter()
+            .map(Into::into)
+            .find(|proto| {
+                // FIXME: When https://github.com/hyperium/http/pull/814
+                //        is merged + released, we can look use
+                //        `contains(proto.as_bytes())` without converting
+                //        to `HeaderValue` first.
+                let Ok(proto) = HeaderValue::from_str(proto) else {
+                    return false;
+                };
+                self.sec_websocket_protocol.contains(&proto)
+            })
+            .map(|protocol| match protocol {
+                Cow::Owned(s) => HeaderValue::from_str(&s).unwrap(),
+                Cow::Borrowed(s) => HeaderValue::from_static(s),
+            });
 
         self
     }
@@ -276,13 +275,8 @@ impl<F> WebSocketUpgrade<F> {
     /// ```
     ///
     /// this method returns an iterator yielding `"soap"` and `"wamp"`.
-    pub fn requested_protocols(&self) -> impl Iterator<Item = &str> {
-        self.sec_websocket_protocol
-            .as_ref()
-            .and_then(|p| p.to_str().ok())
-            .into_iter()
-            .flat_map(|s| s.split(','))
-            .map(|s| s.trim())
+    pub fn requested_protocols(&self) -> impl Iterator<Item = &HeaderValue> {
+        self.sec_websocket_protocol.iter()
     }
 
     /// Set the chosen WebSocket subprotocol.
@@ -500,7 +494,16 @@ where
             .remove::<hyper::upgrade::OnUpgrade>()
             .ok_or(ConnectionNotUpgradable)?;
 
-        let sec_websocket_protocol = parts.headers.get(header::SEC_WEBSOCKET_PROTOCOL).cloned();
+        let sec_websocket_protocol = parts
+            .headers
+            .get_all(header::SEC_WEBSOCKET_PROTOCOL)
+            .iter()
+            .flat_map(|val| val.as_bytes().split(|&b| b == b','))
+            .map(|proto| {
+                HeaderValue::from_bytes(proto.trim_ascii())
+                    .expect("substring of HeaderValue is valid HeaderValue")
+            })
+            .collect();
 
         Ok(Self {
             config: Default::default(),