Rewrite sec-websocket-protocol handling (#3620)

This commit is contained in:
Jonas Platte
2026-04-03 08:49:48 +02:00
committed by Alice Ryhl
parent 8019ae0786
commit 39b9f727d7
5 changed files with 45 additions and 32 deletions
+1 -1
View File
@@ -2,7 +2,7 @@ name: CI
env: env:
CARGO_TERM_COLOR: always CARGO_TERM_COLOR: always
MSRV: '1.78' MSRV: '1.80'
on: on:
push: push:
+1 -1
View File
@@ -11,7 +11,7 @@ exclude = [
resolver = "2" resolver = "2"
[workspace.package] [workspace.package]
rust-version = "1.78" rust-version = "1.80"
[workspace.lints.rust] [workspace.lints.rust]
unsafe_code = "forbid" unsafe_code = "forbid"
+10
View File
@@ -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/), 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). 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 # 0.8.8
- Clarify documentation for `Router::route_layer` ([#3567]) - Clarify documentation for `Router::route_layer` ([#3567])
+1 -1
View File
@@ -104,7 +104,7 @@ This crate uses `#![forbid(unsafe_code)]` to ensure everything is implemented in
## Minimum supported Rust version ## Minimum supported Rust version
axum's MSRV is 1.78. axum's MSRV is 1.80.
## Examples ## Examples
+32 -29
View File
@@ -107,8 +107,10 @@ use hyper_util::rt::TokioIo;
use sha1::{Digest, Sha1}; use sha1::{Digest, Sha1};
use std::{ use std::{
borrow::Cow, borrow::Cow,
collections::BTreeSet,
future::Future, future::Future,
pin::Pin, pin::Pin,
str,
task::{ready, Context, Poll}, task::{ready, Context, Poll},
}; };
use tokio_tungstenite::{ use tokio_tungstenite::{
@@ -138,7 +140,7 @@ pub struct WebSocketUpgrade<F = DefaultOnFailedUpgrade> {
sec_websocket_key: Option<HeaderValue>, sec_websocket_key: Option<HeaderValue>,
on_upgrade: hyper::upgrade::OnUpgrade, on_upgrade: hyper::upgrade::OnUpgrade,
on_failed_upgrade: F, on_failed_upgrade: F,
sec_websocket_protocol: Option<HeaderValue>, sec_websocket_protocol: BTreeSet<HeaderValue>,
} }
impl<F> std::fmt::Debug for WebSocketUpgrade<F> { impl<F> std::fmt::Debug for WebSocketUpgrade<F> {
@@ -242,26 +244,23 @@ impl<F> WebSocketUpgrade<F> {
I: IntoIterator, I: IntoIterator,
I::Item: Into<Cow<'static, str>>, I::Item: Into<Cow<'static, str>>,
{ {
if let Some(req_protocols) = self self.protocol = protocols
.sec_websocket_protocol .into_iter()
.as_ref() .map(Into::into)
.and_then(|p| p.to_str().ok()) .find(|proto| {
{ // FIXME: When https://github.com/hyperium/http/pull/814
self.protocol = protocols // is merged + released, we can look use
.into_iter() // `contains(proto.as_bytes())` without converting
// FIXME: This will often allocate a new `String` and so is less efficient than it // to `HeaderValue` first.
// could be. But that can't be fixed without breaking changes to the public API. let Ok(proto) = HeaderValue::from_str(proto) else {
.map(Into::into) return false;
.find(|protocol| { };
req_protocols self.sec_websocket_protocol.contains(&proto)
.split(',') })
.any(|req_protocol| req_protocol.trim() == protocol) .map(|protocol| match protocol {
}) Cow::Owned(s) => HeaderValue::from_str(&s).unwrap(),
.map(|protocol| match protocol { Cow::Borrowed(s) => HeaderValue::from_static(s),
Cow::Owned(s) => HeaderValue::from_str(&s).unwrap(), });
Cow::Borrowed(s) => HeaderValue::from_static(s),
});
}
self self
} }
@@ -277,13 +276,8 @@ impl<F> WebSocketUpgrade<F> {
/// ``` /// ```
/// ///
/// this method returns an iterator yielding `"soap"` and `"wamp"`. /// this method returns an iterator yielding `"soap"` and `"wamp"`.
pub fn requested_protocols(&self) -> impl Iterator<Item = &str> { pub fn requested_protocols(&self) -> impl Iterator<Item = &HeaderValue> {
self.sec_websocket_protocol self.sec_websocket_protocol.iter()
.as_ref()
.and_then(|p| p.to_str().ok())
.into_iter()
.flat_map(|s| s.split(','))
.map(|s| s.trim())
} }
/// Set the chosen WebSocket subprotocol. /// Set the chosen WebSocket subprotocol.
@@ -501,7 +495,16 @@ where
.remove::<hyper::upgrade::OnUpgrade>() .remove::<hyper::upgrade::OnUpgrade>()
.ok_or(ConnectionNotUpgradable)?; .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 { Ok(Self {
config: Default::default(), config: Default::default(),