diff --git a/CHANGELOG.md b/CHANGELOG.md index d03069b8..c9b0efd7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Breaking - **`ConnectOptions` has a new public field, `websocket_config`.** Code that builds `ConnectOptions` with a struct literal must add it. The field exists only with the `transport-websocket` feature, which is on by default. +- **`WebSocketConfig` has a new public field, `request_modifier`.** Code that builds `WebSocketConfig` with a struct literal must add it. ### Fixed @@ -34,12 +35,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - a subprotocol is not a valid HTTP token, or is listed more than once (RFC 6455 ยง4.1). - Custom headers are sent, and validated, in name order, so the request and the error a bad configuration reports no longer depend on `HashMap` iteration order. - `WebSocketConfig`'s `Debug` output lists custom header names but not their values, which often carry credentials. +- **`MqttClient` returns WebSocket transport errors unchanged.** It wrapped every one in `MqttError::ConnectionError("WebSocket connect failed: ...")`, so a handshake that timed out was not `MqttError::Timeout` and a configuration the transport rejected was not `MqttError::Configuration`. The transport's own connection errors still start with "WebSocket connect failed". +- **`MqttClient` no longer resolves the host of a ws:// or wss:// address itself.** The handshake dials the URL's host, so the separate lookup only added a second resolution. A host that does not resolve now fails with the transport's "WebSocket connect failed" error instead of "Failed to resolve address". ### Added - **`ConnectOptions::with_websocket_config`, so `MqttClient` can send custom WebSocket handshake headers** (#165). `MqttClient` built each connection's `WebSocketConfig` from the URL alone, so its ws:// and wss:// connections could not carry custom headers, subprotocols or a user agent. The configuration set here is applied on every connection attempt, including automatic reconnects, with the URL taken from the address passed to `connect`; for wss://, its TLS configuration, when it has one, is used in place of the one stored with `set_tls_config` or `connect_with_tls`. A connection that follows a server redirect (`UseAnotherServer` or `ServerMoved`) does not use it, so its headers are never sent to a host the broker names. - `WebSocketConfig` implements `Clone`. +- **`WebSocketConfig::with_request_modifier`, an async callback that adjusts the WebSocket upgrade request before each connection attempt** (#165). It receives the validated request and returns the one to send, so it can add values computed at connect time, such as a short-lived signed token, or replace the URI with a presigned one. Set on the `WebSocketConfig` given to `ConnectOptions::with_websocket_config`, it runs before every `MqttClient` connection attempt, including automatic reconnects. It runs within the connect timeout, and an error fails the attempt with `MqttError::ConnectionError`, logged at error level. The returned request must keep the URI's scheme, the handshake headers the transport sets, and a `Sec-WebSocket-Protocol` header that offers `mqtt`, and must not add a header `with_header` reserves; otherwise the attempt fails with `MqttError::Configuration` before dialing. Because `MqttClient` does not resolve a WebSocket address, the address passed to `connect` can name a placeholder host that the callback replaces. - `TlsConfig::client_config`, which returns the rustls client configuration a `TlsConfig` describes. +- `RequestModifier` and `RequestModifierError` are exported from `mqtt5::transport`. - `WebSocketConfig::build_handshake_request`, which returns the upgrade request the configuration produces. ## [mqtt5 0.47.3] - 2026-10-09 diff --git a/crates/mqtt5/src/client/inner.rs b/crates/mqtt5/src/client/inner.rs index 2818d32a..ba225eec 100644 --- a/crates/mqtt5/src/client/inner.rs +++ b/crates/mqtt5/src/client/inner.rs @@ -195,12 +195,8 @@ impl MqttClient { pub(crate) fn select_addresses_for_connection<'a>( addrs: &'a [std::net::SocketAddr], host: &str, - transport_type: &ClientTransportType, ) -> &'a [std::net::SocketAddr] { - if transport_type.dials_by_url() { - tracing::debug!("WebSocket transport dials the URL host, using one attempt"); - &addrs[0..1] - } else if Self::is_aws_iot_endpoint(host) { + if Self::is_aws_iot_endpoint(host) { tracing::debug!("AWS IoT endpoint detected, limiting to first resolved address"); &addrs[0..1] } else { @@ -351,10 +347,7 @@ impl MqttClient { #[cfg(feature = "transport-websocket")] async fn open_websocket(config: WebSocketConfig) -> Result { let mut ws_transport = WebSocketTransport::new(config); - ws_transport - .connect() - .await - .map_err(|e| MqttError::ConnectionError(format!("WebSocket connect failed: {e}")))?; + ws_transport.connect().await?; Ok(TransportType::WebSocket(Box::new(ws_transport))) } @@ -545,9 +538,15 @@ impl MqttClient { } else { requested_transport_type.redirected() }; - let addrs = Self::resolve_addresses(host, port).await?; - let addresses_to_try = - Self::select_addresses_for_connection(&addrs, host, &client_transport_type); + let addrs = if client_transport_type.dials_by_url() { + vec![std::net::SocketAddr::new( + std::net::Ipv4Addr::UNSPECIFIED.into(), + port, + )] + } else { + Self::resolve_addresses(host, port).await? + }; + let addresses_to_try = Self::select_addresses_for_connection(&addrs, host); match self .try_connect_to_addresses(addresses_to_try, client_transport_type, host) diff --git a/crates/mqtt5/src/transport.rs b/crates/mqtt5/src/transport.rs index ef461198..04444f1c 100644 --- a/crates/mqtt5/src/transport.rs +++ b/crates/mqtt5/src/transport.rs @@ -43,7 +43,7 @@ pub use tcp::{TcpConfig, TcpTransport}; #[cfg(not(target_arch = "wasm32"))] pub use tls::{TlsConfig, TlsTransport}; #[cfg(all(not(target_arch = "wasm32"), feature = "transport-websocket"))] -pub use websocket::{WebSocketConfig, WebSocketTransport}; +pub use websocket::{RequestModifier, RequestModifierError, WebSocketConfig, WebSocketTransport}; #[cfg(not(target_arch = "wasm32"))] pub enum TransportType { diff --git a/crates/mqtt5/src/transport/websocket.rs b/crates/mqtt5/src/transport/websocket.rs index a35b51bf..0daf8131 100644 --- a/crates/mqtt5/src/transport/websocket.rs +++ b/crates/mqtt5/src/transport/websocket.rs @@ -48,7 +48,9 @@ use crate::Transport; use bytes::{Buf, Bytes, BytesMut}; use futures_util::{stream::SplitSink, stream::SplitStream, StreamExt}; use std::collections::HashMap; +use std::future::Future; use std::net::SocketAddr; +use std::pin::Pin; use std::sync::Arc; use tokio::net::TcpStream; use tokio_tungstenite::{ @@ -96,6 +98,14 @@ const RESERVED_HEADERS: [&str; 14] = [ "expect", ]; +const TRANSPORT_HEADERS: [&str; 5] = [ + "host", + "connection", + "upgrade", + "sec-websocket-version", + "sec-websocket-key", +]; + /// WebSocket transport configuration #[derive(Clone)] pub struct WebSocketConfig { @@ -112,6 +122,37 @@ pub struct WebSocketConfig { pub user_agent: Option, /// TLS configuration for secure WebSocket connections (wss://) pub tls_config: Option, + /// Adjusts the upgrade request before each connection attempt; see + /// [`WebSocketConfig::with_request_modifier`] + pub request_modifier: Option, +} + +/// Error a [`RequestModifier`] fails with +pub type RequestModifierError = Box; + +type ModifiedRequest = + Pin, RequestModifierError>> + Send>>; + +/// A callback that adjusts the WebSocket upgrade request before each +/// connection attempt, set with [`WebSocketConfig::with_request_modifier`] +#[derive(Clone)] +pub struct RequestModifier(Arc) -> ModifiedRequest + Send + Sync>); + +impl RequestModifier { + async fn apply(&self, request: Request<()>, secure: bool) -> Result> { + let modified = (self.0)(request).await.map_err(|e| { + error!(error = %e, "WebSocket request modifier failed"); + MqttError::ConnectionError(format!("WebSocket request modifier failed: {e}")) + })?; + validate_modified_request(&modified, secure)?; + Ok(modified) + } +} + +impl std::fmt::Debug for RequestModifier { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("RequestModifier(..)") + } } impl std::fmt::Debug for WebSocketConfig { @@ -125,6 +166,7 @@ impl std::fmt::Debug for WebSocketConfig { .field("headers", &header_names) .field("user_agent", &self.user_agent) .field("tls_config", &self.tls_config) + .field("request_modifier", &self.request_modifier) .finish() } } @@ -143,6 +185,7 @@ impl WebSocketConfig { headers: HashMap::new(), user_agent: Some(concat!("mqtt5/", env!("CARGO_PKG_VERSION")).to_string()), tls_config: None, + request_modifier: None, }) } @@ -217,6 +260,61 @@ impl WebSocketConfig { self } + /// Sets a callback that adjusts the upgrade request before each connection + /// attempt + /// + /// The callback receives the validated request that + /// [`build_handshake_request`](Self::build_handshake_request) produces and + /// returns the request to send. It can add values computed at connect + /// time, such as a short-lived signed token, or replace the URI, such as + /// with a presigned URL. Through + /// [`ConnectOptions::with_websocket_config`](crate::ConnectOptions::with_websocket_config) + /// it runs before every connection attempt an `MqttClient` makes, + /// including automatic reconnects, so those values never go stale. + /// + /// The connection is made to the returned request's URI, so a callback + /// that points it at another host must update `Host` too; `MqttClient` + /// does not resolve the host it was given for a WebSocket connection, so + /// that host may be a placeholder. The returned request must keep the + /// URI's scheme, exactly one of each header the transport sets (`Host`, + /// `Connection`, `Upgrade`, `Sec-WebSocket-Version`, `Sec-WebSocket-Key`), + /// none of the other headers [`with_header`](Self::with_header) reserves, + /// and one `Sec-WebSocket-Protocol` header that still offers `mqtt`; + /// otherwise the attempt fails with `MqttError::Configuration` before + /// dialing. Other headers the callback adds are sent unchecked. + /// The callback runs within the [`timeout`](Self::with_timeout), and an + /// error fails the attempt with `MqttError::ConnectionError`. + /// + /// ```rust,no_run + /// # use mqtt5::transport::websocket::WebSocketConfig; + /// # async fn sign() -> Result { Ok(String::new()) } + /// # fn example() -> mqtt5::Result<()> { + /// let config = WebSocketConfig::new("wss://broker.example.com/mqtt")? + /// .with_header("x-amz-customauthorizer-name", "my-authorizer") + /// .with_request_modifier(|mut request| async move { + /// let signature = sign().await?; + /// request + /// .headers_mut() + /// .insert("x-amz-customauthorizer-signature", signature.parse()?); + /// Ok::<_, Box>(request) + /// }); + /// # Ok(()) + /// # } + /// ``` + #[must_use] + pub fn with_request_modifier(mut self, modifier: F) -> Self + where + F: Fn(Request<()>) -> Fut + Send + Sync + 'static, + Fut: Future, E>> + Send + 'static, + E: Into, + { + self.request_modifier = Some(RequestModifier(Arc::new(move |request| { + let pending = modifier(request); + Box::pin(async move { pending.await.map_err(Into::into) }) + }))); + self + } + /// Sets a custom TLS configuration for wss:// connections /// /// Its root certificates, system-roots setting, server-certificate @@ -502,6 +600,45 @@ impl WebSocketConfig { } } +fn validate_modified_request(request: &Request<()>, secure: bool) -> Result<()> { + let scheme = if secure { "wss" } else { "ws" }; + if request.uri().scheme_str() != Some(scheme) { + return Err(MqttError::Configuration(format!( + "WebSocket request modifier changed the URI scheme; it must stay {scheme}" + ))); + } + + let headers = request.headers(); + for name in RESERVED_HEADERS { + let expected = usize::from(TRANSPORT_HEADERS.contains(&name)); + let found = headers.get_all(name).iter().count(); + if found != expected { + return Err(MqttError::Configuration(format!( + "WebSocket request modifier left {found} {name:?} headers; the handshake needs {expected}" + ))); + } + } + + let mut protocols = headers.get_all(SEC_WEBSOCKET_PROTOCOL).iter(); + let (Some(offered), None) = (protocols.next(), protocols.next()) else { + return Err(MqttError::Configuration( + "WebSocket request modifier must leave exactly one Sec-WebSocket-Protocol header" + .to_string(), + )); + }; + let offers_mqtt = offered.to_str().is_ok_and(|offered| { + offered + .split(',') + .any(|subprotocol| subprotocol.trim() == DEFAULT_SUBPROTOCOL) + }); + if !offers_mqtt { + return Err(MqttError::Configuration(format!( + "WebSocket request modifier must keep offering the {DEFAULT_SUBPROTOCOL:?} subprotocol" + ))); + } + Ok(()) +} + /// Parses a handshake header value, naming the header in the error. fn header_value(name: &str, value: &str) -> Result { HeaderValue::from_str(value).map_err(|e| { @@ -724,11 +861,20 @@ impl Transport for WebSocketTransport { _ => None, }; - let ws_result = tokio::time::timeout( - self.config.timeout, - tokio_tungstenite::connect_async_tls_with_config(request, None, false, connector), - ) - .await; + let secure = self.config.is_secure(); + let handshake = async { + let request = match &self.config.request_modifier { + Some(modifier) => modifier.apply(request, secure).await?, + None => request, + }; + tokio_tungstenite::connect_async_tls_with_config(request, None, false, connector) + .await + .map_err(|e| { + error!(error = %e, "WebSocket connection failed"); + MqttError::ConnectionError(format!("WebSocket connect failed: {e}")) + }) + }; + let ws_result = tokio::time::timeout(self.config.timeout, Box::pin(handshake)).await; match ws_result { Ok(Ok((ws_stream, response))) => { @@ -744,10 +890,7 @@ impl Transport for WebSocketTransport { debug!("WebSocket connection established"); Ok(()) } - Ok(Err(e)) => { - error!(error = %e, "WebSocket connection failed"); - Err(MqttError::ConnectionError(e.to_string())) - } + Ok(Err(e)) => Err(e), Err(_) => { error!("WebSocket connection timed out"); Err(MqttError::Timeout) diff --git a/crates/mqtt5/tests/client_websocket_config.rs b/crates/mqtt5/tests/client_websocket_config.rs index 22f5b1b5..1b65e751 100644 --- a/crates/mqtt5/tests/client_websocket_config.rs +++ b/crates/mqtt5/tests/client_websocket_config.rs @@ -10,6 +10,7 @@ mod wss_test_server; use futures::{SinkExt, StreamExt}; +use mqtt5::error::MqttError; use mqtt5::time::Duration; use mqtt5::transport::websocket::WebSocketConfig; use mqtt5::{ConnectOptions, MqttClient}; @@ -169,6 +170,102 @@ async fn websocket_config_is_not_sent_to_a_redirect_target() { let _ = client.disconnect().await; } +#[tokio::test] +async fn request_modifier_runs_again_on_reconnect() { + let (url, mut requests) = closing_mqtt_server(CONNACK.to_vec()).await; + let calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let counter = std::sync::Arc::clone(&calls); + let websocket = WebSocketConfig::new("ws://placeholder.invalid/mqtt") + .expect("config") + .with_request_modifier(move |mut request| { + let call = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1; + async move { + request + .headers_mut() + .insert("x-token", format!("token-{call}").parse()?); + Ok::<_, Box>(request) + } + }); + let options = ConnectOptions::new("ws-modifier-client") + .with_websocket_config(websocket) + .with_reconnect_delay(Duration::from_millis(100), Duration::from_millis(500)); + let client = MqttClient::with_options(options); + + timeout(WAIT, client.connect(&url)) + .await + .expect("connect timed out") + .expect("connect"); + + for (attempt, token) in [("connect", "token-1"), ("reconnect", "token-2")] { + let headers = timeout(WAIT, requests.recv()) + .await + .unwrap_or_else(|_| panic!("no upgrade request for the {attempt}")) + .expect("server stopped"); + assert_eq!(header(&headers, "x-token"), Some(token), "{attempt}"); + } + + let _ = client.disconnect().await; +} + +#[tokio::test] +async fn request_modifier_can_point_a_placeholder_host_at_the_broker() { + let (url, mut requests) = closing_mqtt_server(CONNACK.to_vec()).await; + let broker: tokio_tungstenite::tungstenite::http::Uri = url.parse().expect("broker uri"); + let websocket = WebSocketConfig::new("ws://placeholder.invalid/mqtt") + .expect("config") + .with_request_modifier(move |mut request| { + let broker = broker.clone(); + async move { + let host = broker.authority().expect("authority").as_str().parse()?; + request.headers_mut().insert("host", host); + *request.uri_mut() = broker; + Ok::<_, Box>(request) + } + }); + let client = MqttClient::with_options( + ConnectOptions::new("ws-placeholder-client") + .with_websocket_config(websocket) + .with_automatic_reconnect(false), + ); + + timeout(WAIT, client.connect("ws://placeholder.invalid/mqtt")) + .await + .expect("connect timed out") + .expect("connect"); + let headers = timeout(WAIT, requests.recv()) + .await + .expect("no upgrade request") + .expect("server stopped"); + assert_eq!( + header(&headers, "host"), + url.strip_prefix("ws://") + .and_then(|rest| rest.strip_suffix("/mqtt")) + ); + + let _ = client.disconnect().await; +} + +#[tokio::test] +async fn client_reports_a_rejected_websocket_config_as_a_configuration_error() { + let websocket = WebSocketConfig::new("ws://127.0.0.1:9/mqtt") + .expect("config") + .with_header("Content-Length", "5"); + let client = MqttClient::with_options( + ConnectOptions::new("ws-rejected-config-client") + .with_websocket_config(websocket) + .with_automatic_reconnect(false), + ); + + let err = timeout(WAIT, client.connect("ws://127.0.0.1:9/mqtt")) + .await + .expect("connect timed out") + .expect_err("connect accepted a reserved header"); + assert!( + matches!(err, MqttError::Configuration(_)), + "unexpected error: {err:?}" + ); +} + #[tokio::test] async fn websocket_config_tls_is_used_for_wss() { let (addr, server) = wss_server(server_config( diff --git a/crates/mqtt5/tests/websocket_request_modifier.rs b/crates/mqtt5/tests/websocket_request_modifier.rs new file mode 100644 index 00000000..85e6b319 --- /dev/null +++ b/crates/mqtt5/tests/websocket_request_modifier.rs @@ -0,0 +1,287 @@ +#![cfg(feature = "transport-websocket")] + +//! Checks that `WebSocketConfig::with_request_modifier` adjusts the upgrade +//! request the transport sends, as seen by a server. + +use mqtt5::error::MqttError; +use mqtt5::transport::websocket::{WebSocketConfig, WebSocketTransport}; +use mqtt5::{ConnectOptions, MqttClient, Transport}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::net::TcpListener; +use tokio::time::timeout; +use tokio_tungstenite::tungstenite::handshake::server::{ + Callback, ErrorResponse, Request, Response, +}; +use tokio_tungstenite::tungstenite::http::{HeaderMap, HeaderValue, Uri}; + +const WAIT: Duration = Duration::from_secs(5); + +/// What the server saw of the upgrade request. +type Seen = Arc>>; + +/// Records the upgrade request and selects "mqtt", as a broker would. +struct Capture(Seen); + +impl Callback for Capture { + fn on_request(self, req: &Request, mut resp: Response) -> Result { + *self.0.lock().expect("lock") = Some((req.uri().clone(), req.headers().clone())); + resp.headers_mut() + .insert("Sec-WebSocket-Protocol", HeaderValue::from_static("mqtt")); + Ok(resp) + } +} + +/// Connects a transport with `config` to a one-connection loopback server +/// and returns the request the server received. +async fn received(config: impl FnOnce(String) -> WebSocketConfig) -> (Uri, HeaderMap) { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("addr"); + let seen = Seen::default(); + let capture = Capture(Arc::clone(&seen)); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept"); + tokio_tungstenite::accept_hdr_async(stream, capture) + .await + .expect("server handshake") + }); + + let mut transport = WebSocketTransport::new(config(format!("ws://{addr}/mqtt"))); + timeout(WAIT, transport.connect()) + .await + .expect("connect timed out") + .expect("client handshake"); + let _server_ws = timeout(WAIT, server) + .await + .expect("server timed out") + .expect("server task"); + transport.close().await.expect("close"); + + let request = seen.lock().expect("lock").take(); + request.expect("server saw no upgrade request") +} + +#[tokio::test] +async fn modifier_adds_headers_computed_at_connect_time() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&calls); + let (_, headers) = received(|url| { + WebSocketConfig::new(&url) + .expect("config") + .with_header("X-Static", "configured") + .with_request_modifier(move |mut request| { + let call = counter.fetch_add(1, Ordering::SeqCst) + 1; + async move { + request + .headers_mut() + .insert("x-token", HeaderValue::from(call)); + Ok::<_, std::convert::Infallible>(request) + } + }) + }) + .await; + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!( + headers.get("x-token").and_then(|v| v.to_str().ok()), + Some("1") + ); + assert_eq!( + headers.get("x-static").and_then(|v| v.to_str().ok()), + Some("configured") + ); +} + +#[tokio::test] +async fn modifier_can_replace_the_uri() { + let (uri, _) = received(|url| { + WebSocketConfig::new(&url) + .expect("config") + .with_request_modifier(|mut request| async move { + let presigned = format!("{}?X-Amz-Signature=abc123", request.uri()); + *request.uri_mut() = presigned.parse()?; + Ok::<_, Box>(request) + }) + }) + .await; + + assert_eq!(uri.path(), "/mqtt"); + assert_eq!(uri.query(), Some("X-Amz-Signature=abc123")); +} + +#[tokio::test] +async fn modifier_error_fails_connect_before_dialing() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("addr"); + let config = WebSocketConfig::new(&format!("ws://{addr}/mqtt")) + .expect("config") + .with_request_modifier(|_request| async { + Err::(std::io::Error::other("token service unavailable")) + }); + let mut transport = WebSocketTransport::new(config); + + let err = timeout(WAIT, transport.connect()) + .await + .expect("connect timed out") + .expect_err("connect succeeded despite the modifier failing"); + assert!( + matches!(&err, MqttError::ConnectionError(msg) + if msg.contains("request modifier failed") && msg.contains("token service unavailable")), + "unexpected error: {err}" + ); + assert!(!transport.is_connected()); + assert!( + timeout(Duration::from_millis(200), listener.accept()) + .await + .is_err(), + "the transport dialed although the modifier failed" + ); +} + +#[tokio::test] +async fn modifier_runs_within_the_connect_timeout() { + let config = WebSocketConfig::new("ws://127.0.0.1:9/mqtt") + .expect("config") + .with_timeout(mqtt5::time::Duration::from_millis(200)) + .with_request_modifier(|_request| { + std::future::pending::>() + }); + let mut transport = WebSocketTransport::new(config); + + let err = timeout(WAIT, transport.connect()) + .await + .expect("the modifier was not bounded by the connect timeout") + .expect_err("connect succeeded with a modifier that never finishes"); + assert!(matches!(err, MqttError::Timeout), "unexpected error: {err}"); +} + +async fn rejected_before_dialing(modify: fn(&mut Request)) -> MqttError { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("addr"); + let config = WebSocketConfig::new(&format!("ws://{addr}/mqtt")) + .expect("config") + .with_request_modifier(move |mut request| { + modify(&mut request); + async move { Ok::<_, std::convert::Infallible>(request) } + }); + let mut transport = WebSocketTransport::new(config); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept"); + tokio_tungstenite::accept_hdr_async(stream, Capture(Seen::default())).await + }); + + let err = timeout(WAIT, transport.connect()) + .await + .expect("connect timed out") + .expect_err("connect sent a request the modifier made invalid"); + assert!( + timeout(Duration::from_millis(200), server).await.is_err(), + "the transport dialed although the modified request is invalid" + ); + err +} + +fn assert_configuration_error(err: &MqttError, expected: &str) { + assert!( + matches!(err, MqttError::Configuration(msg) if msg.contains(expected)), + "unexpected error: {err}" + ); +} + +#[tokio::test] +async fn modifier_cannot_add_a_reserved_header() { + let err = rejected_before_dialing(|request| { + request + .headers_mut() + .insert("content-length", HeaderValue::from_static("5")); + }) + .await; + assert_configuration_error(&err, "content-length"); +} + +#[tokio::test] +async fn modifier_cannot_repeat_a_handshake_header() { + let err = rejected_before_dialing(|request| { + request + .headers_mut() + .append("host", HeaderValue::from_static("elsewhere.example")); + }) + .await; + assert_configuration_error(&err, "host"); +} + +#[tokio::test] +async fn modifier_cannot_stop_offering_mqtt() { + let err = rejected_before_dialing(|request| { + request + .headers_mut() + .insert("sec-websocket-protocol", HeaderValue::from_static("mqttv5")); + }) + .await; + assert_configuration_error(&err, "mqtt"); +} + +#[tokio::test] +async fn modifier_cannot_remove_the_subprotocol_header() { + let err = rejected_before_dialing(|request| { + request.headers_mut().remove("sec-websocket-protocol"); + }) + .await; + assert_configuration_error(&err, "Sec-WebSocket-Protocol"); +} + +#[tokio::test] +async fn modifier_cannot_change_the_scheme() { + let err = rejected_before_dialing(|request| { + let secure = request.uri().to_string().replacen("ws://", "wss://", 1); + *request.uri_mut() = secure.parse().expect("uri"); + }) + .await; + assert_configuration_error(&err, "scheme"); +} + +fn client_with(config: WebSocketConfig) -> MqttClient { + MqttClient::with_options( + ConnectOptions::new("ws-modifier-errors") + .with_websocket_config(config) + .with_automatic_reconnect(false), + ) +} + +#[tokio::test] +async fn client_reports_a_hung_modifier_as_a_timeout() { + let config = WebSocketConfig::new("ws://127.0.0.1:9/mqtt") + .expect("config") + .with_timeout(mqtt5::time::Duration::from_millis(200)) + .with_request_modifier(|_request| { + std::future::pending::>() + }); + let client = client_with(config); + + let err = timeout(WAIT, client.connect("ws://127.0.0.1:9/mqtt")) + .await + .expect("connect finishes") + .expect_err("connect succeeded with a modifier that never finishes"); + assert!(matches!(err, MqttError::Timeout), "unexpected error: {err}"); +} + +#[tokio::test] +async fn client_reports_a_modifier_error_once() { + let config = WebSocketConfig::new("ws://127.0.0.1:9/mqtt") + .expect("config") + .with_request_modifier(|_request| async { + Err::(std::io::Error::other("token service unavailable")) + }); + let client = client_with(config); + + let err = timeout(WAIT, client.connect("ws://127.0.0.1:9/mqtt")) + .await + .expect("connect finishes") + .expect_err("connect succeeded despite the modifier failing"); + assert!( + matches!(&err, MqttError::ConnectionError(msg) + if msg == "WebSocket request modifier failed: token service unavailable"), + "unexpected error: {err:?}" + ); +}