add WebSocketConfig::with_request_modifier (#165)

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

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<TransportType> {
         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)

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 {

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<String>,
     /// TLS configuration for secure WebSocket connections (wss://)
     pub tls_config: Option<TlsConfig>,
+    /// Adjusts the upgrade request before each connection attempt; see
+    /// [`WebSocketConfig::with_request_modifier`]
+    pub request_modifier: Option<RequestModifier>,
+}
+
+/// Error a [`RequestModifier`] fails with
+pub type RequestModifierError = Box<dyn std::error::Error + Send + Sync>;
+
+type ModifiedRequest =
+    Pin<Box<dyn Future<Output = std::result::Result<Request<()>, 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<dyn Fn(Request<()>) -> ModifiedRequest + Send + Sync>);
+
+impl RequestModifier {
+    async fn apply(&self, request: Request<()>, secure: bool) -> Result<Request<()>> {
+        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<String, std::io::Error> { 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<dyn std::error::Error + Send + Sync>>(request)
+    ///     });
+    /// # Ok(())
+    /// # }
+    /// ```
+    #[must_use]
+    pub fn with_request_modifier<F, Fut, E>(mut self, modifier: F) -> Self
+    where
+        F: Fn(Request<()>) -> Fut + Send + Sync + 'static,
+        Fut: Future<Output = std::result::Result<Request<()>, E>> + Send + 'static,
+        E: Into<RequestModifierError>,
+    {
+        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> {
     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)