Rewrite sec-websocket-protocol handling (#3620)

.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

@@ -11,7 +11,7 @@ exclude = [
 resolver = "2"
 
 [workspace.package]
-rust-version = "1.78"
+rust-version = "1.80"
 
 [workspace.lints.rust]
 unsafe_code = "forbid"

axum/CHANGELOG.md

@@ -5,6 +5,16 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+
+# 0.8.9
+
+- **added:** `WebSocketUpgrade::{requested_protocols, set_selected_protocol}` for more
+  flexible subprotocol selection ([#3597])
+- **changed:** Update minimum rust version to 1.80 ([#3620])
+
+[#3597]: https://github.com/tokio-rs/axum/pull/3597
+[#3620]: https://github.com/tokio-rs/axum/pull/3620
+
 # 0.8.8
 
 - Clarify documentation for `Router::route_layer` ([#3567])

axum/README.md

@@ -104,7 +104,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

@@ -107,8 +107,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::{
@@ -138,7 +140,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> {
@@ -242,26 +244,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
     }
@@ -277,13 +276,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.
@@ -501,7 +495,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(),