add ConnectOptions::with_websocket_config for MqttClient (#165)

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### 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.
+
 ### Fixed
 
 - **`wss://` connections now use the TLS configuration on `WebSocketConfig`** (#190). With server-certificate verification on, `WebSocketTransport::connect` used tokio-tungstenite's own client configuration built from the platform's root certificates, so a CA set with `with_ca_cert_from_file` or `with_ca_cert_from_bytes` was never trusted (a server it issued failed with `UnknownIssuer`), a client certificate set with `with_client_auth_from_files` or `with_client_auth_from_bytes` was never presented, and `use_system_roots` and ALPN protocols had no effect. With verification off, the client certificate and ALPN protocols were dropped. The handshake now uses the rustls configuration `TlsTransport` builds from a `TlsConfig`, in both cases, except that `use_system_roots` adds the platform's root certificates, as a `wss://` connection without a TLS configuration already did, rather than the bundled `webpki-roots`.
@@ -32,6 +36,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### 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`.
 - `TlsConfig::client_config`, which returns the rustls client configuration a `TlsConfig` describes.
 - `WebSocketConfig::build_handshake_request`, which returns the upgrade request the configuration produces.
 

crates/mqtt5/src/client/inner.rs

@@ -20,6 +20,8 @@ use crate::transport::{TcpTransport, TlsTransport, TransportType};
 use super::connection::ConnectionEvent;
 use super::direct::{AutomaticReconnectLifecycle, StoredSubscription};
 use super::state::ClientTransportType;
+#[cfg(feature = "transport-websocket")]
+use super::state::WebSocketSettings;
 use super::MqttClient;
 
 impl MqttClient {
@@ -71,7 +73,10 @@ impl MqttClient {
             {
                 let (host, port) = Self::split_host_port(rest, 80)?;
                 Ok((
-                    ClientTransportType::WebSocket(address.to_string()),
+                    ClientTransportType::WebSocket(
+                        address.to_string(),
+                        WebSocketSettings::Configured,
+                    ),
                     host,
                     port,
                 ))
@@ -88,7 +93,10 @@ impl MqttClient {
             {
                 let (host, port) = Self::split_host_port(rest, 443)?;
                 Ok((
-                    ClientTransportType::WebSocketSecure(address.to_string()),
+                    ClientTransportType::WebSocketSecure(
+                        address.to_string(),
+                        WebSocketSettings::Configured,
+                    ),
                     host,
                     port,
                 ))
@@ -208,10 +216,13 @@ impl MqttClient {
             ClientTransportType::Tcp => Self::connect_tcp(addr).await,
             ClientTransportType::Tls => self.connect_tls(addr, host).await,
             #[cfg(feature = "transport-websocket")]
-            ClientTransportType::WebSocket(url) => Self::connect_websocket(&url).await,
+            ClientTransportType::WebSocket(url, settings) => {
+                self.connect_websocket(&url, settings).await
+            }
             #[cfg(feature = "transport-websocket")]
-            ClientTransportType::WebSocketSecure(url) => {
-                self.connect_websocket_secure(addr, host, &url).await
+            ClientTransportType::WebSocketSecure(url, settings) => {
+                self.connect_websocket_secure(addr, host, &url, settings)
+                    .await
             }
             #[cfg(feature = "transport-quic")]
             ClientTransportType::Quic => self.connect_quic(addr, host).await,
@@ -259,16 +270,37 @@ impl MqttClient {
         Ok(TransportType::Tls(Box::new(tls_transport)))
     }
 
+    /// The WebSocket configuration for a connection to `url`: the one set with
+    /// `ConnectOptions::with_websocket_config`, read on every attempt so
+    /// reconnects use it too, or a default one, which a server redirect
+    /// always gets.
     #[cfg(feature = "transport-websocket")]
-    async fn connect_websocket(url: &str) -> Result<TransportType> {
-        let config = WebSocketConfig::new(url)
-            .map_err(|e| MqttError::ConnectionError(format!("Invalid WebSocket URL: {e}")))?;
-        let mut ws_transport = WebSocketTransport::new(config);
-        ws_transport
-            .connect()
-            .await
-            .map_err(|e| MqttError::ConnectionError(format!("WebSocket connect failed: {e}")))?;
-        Ok(TransportType::WebSocket(Box::new(ws_transport)))
+    async fn websocket_config_for(
+        &self,
+        url: &str,
+        settings: WebSocketSettings,
+    ) -> Result<WebSocketConfig> {
+        let configured = match settings {
+            WebSocketSettings::Configured => {
+                self.inner.read().await.options.websocket_config.clone()
+            }
+            WebSocketSettings::Default => None,
+        };
+        match configured {
+            Some(config) => WebSocketConfig::clone(&config).with_url(url),
+            None => WebSocketConfig::new(url),
+        }
+        .map_err(|e| MqttError::ConnectionError(format!("Invalid WebSocket URL: {e}")))
+    }
+
+    #[cfg(feature = "transport-websocket")]
+    async fn connect_websocket(
+        &self,
+        url: &str,
+        settings: WebSocketSettings,
+    ) -> Result<TransportType> {
+        let config = self.websocket_config_for(url, settings).await?;
+        Self::open_websocket(config).await
     }
 
     #[cfg(feature = "transport-websocket")]
@@ -277,10 +309,17 @@ impl MqttClient {
         addr: std::net::SocketAddr,
         host: &str,
         url: &str,
+        settings: WebSocketSettings,
     ) -> Result<TransportType> {
         let insecure = self.transport_config.read().await.insecure_tls;
-        let mut config = WebSocketConfig::new(url)
-            .map_err(|e| MqttError::ConnectionError(format!("Invalid WebSocket URL: {e}")))?;
+        let mut config = self.websocket_config_for(url, settings).await?;
+
+        if let Some(tls_config) = config.tls_config.as_mut() {
+            if insecure {
+                tls_config.verify_server_cert = false;
+            }
+            return Self::open_websocket(config).await;
+        }
 
         // As in connect_tls: a stored TLS config applies to wss:// too. With
         // none, an insecure connection still needs one to disable
@@ -304,6 +343,11 @@ impl MqttClient {
             config = config.with_tls_config(tls_config);
         }
 
+        Self::open_websocket(config).await
+    }
+
+    #[cfg(feature = "transport-websocket")]
+    async fn open_websocket(config: WebSocketConfig) -> Result<TransportType> {
         let mut ws_transport = WebSocketTransport::new(config);
         ws_transport
             .connect()
@@ -471,10 +515,18 @@ impl MqttClient {
     }
 
     pub(crate) async fn connect_internal(&self, address: &str) -> Result<ConnectResult> {
+        Box::pin(self.connect_following_redirects(address, 0)).await
+    }
+
+    async fn connect_following_redirects(
+        &self,
+        address: &str,
+        redirects_followed: u8,
+    ) -> Result<ConnectResult> {
         const MAX_REDIRECTS: u8 = 3;
         self.inner.read().await.options.validate_deferred_ack()?;
         let mut current_address = address.to_string();
-        let mut redirect_count: u8 = 0;
+        let mut redirect_count = redirects_followed;
 
         loop {
             let client_id = self.inner.read().await.options.client_id.clone();
@@ -485,7 +537,12 @@ impl MqttClient {
                 "Connection attempt"
             );
 
-            let (client_transport_type, host, port) = Self::parse_address(&current_address)?;
+            let (requested_transport_type, host, port) = Self::parse_address(&current_address)?;
+            let client_transport_type = if redirect_count == 0 {
+                requested_transport_type
+            } else {
+                requested_transport_type.redirected()
+            };
             let addrs = Self::resolve_addresses(host, port).await?;
             let addresses_to_try = Self::select_addresses_for_connection(&addrs, host);
 
@@ -569,7 +626,7 @@ impl MqttClient {
                         reason = ?reason,
                         "Following server redirect from TLS connection"
                     );
-                    return Box::pin(self.connect_internal(&url)).await;
+                    return Box::pin(self.connect_following_redirects(&url, 1)).await;
                 }
                 Err(MqttError::ConnectionError(
                     "Server redirect missing server reference".to_string(),

crates/mqtt5/src/client/mod.rs

@@ -1344,23 +1344,23 @@ mod tests {
             let (transport, host, port) = MqttClient::parse_address("ws://localhost").unwrap();
             assert!(matches!(
                 transport,
-                state::ClientTransportType::WebSocket(_)
+                state::ClientTransportType::WebSocket(..)
             ));
             assert_eq!(host, "localhost");
             assert_eq!(port, 80);
 
             let (transport, host, port) = MqttClient::parse_address("ws://localhost:8080").unwrap();
             assert!(matches!(
                 transport,
-                state::ClientTransportType::WebSocket(_)
+                state::ClientTransportType::WebSocket(..)
             ));
             assert_eq!(host, "localhost");
             assert_eq!(port, 8080);
 
             let (transport, host, port) = MqttClient::parse_address("wss://secure.broker").unwrap();
             assert!(matches!(
                 transport,
-                state::ClientTransportType::WebSocketSecure(_)
+                state::ClientTransportType::WebSocketSecure(..)
             ));
             assert_eq!(host, "secure.broker");
             assert_eq!(port, 443);
@@ -1369,14 +1369,14 @@ mod tests {
                 MqttClient::parse_address("wss://secure.broker:8443").unwrap();
             assert!(matches!(
                 transport,
-                state::ClientTransportType::WebSocketSecure(_)
+                state::ClientTransportType::WebSocketSecure(..)
             ));
             assert_eq!(host, "secure.broker");
             assert_eq!(port, 8443);
 
             let (transport, host, port) =
                 MqttClient::parse_address("ws://broker.emqx.io:8083/mqtt").unwrap();
-            if let state::ClientTransportType::WebSocket(url) = transport {
+            if let state::ClientTransportType::WebSocket(url, _) = transport {
                 assert_eq!(url, "ws://broker.emqx.io:8083/mqtt");
             } else {
                 panic!("Expected WebSocket transport type");
@@ -1386,7 +1386,7 @@ mod tests {
 
             let (transport, host, port) =
                 MqttClient::parse_address("wss://broker.hivemq.com:8884/mqtt").unwrap();
-            if let state::ClientTransportType::WebSocketSecure(url) = transport {
+            if let state::ClientTransportType::WebSocketSecure(url, _) = transport {
                 assert_eq!(url, "wss://broker.hivemq.com:8884/mqtt");
             } else {
                 panic!("Expected WebSocketSecure transport type");

crates/mqtt5/src/client/state.rs

@@ -12,20 +12,39 @@ use super::direct::AutomaticReconnectLifecycle;
 use super::direct::{StoredSubscription, SubscriptionPersistence};
 use super::MqttClient;
 
+#[cfg(feature = "transport-websocket")]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum WebSocketSettings {
+    Configured,
+    Default,
+}
+
 #[derive(Debug, Clone)]
 pub(crate) enum ClientTransportType {
     Tcp,
     Tls,
     #[cfg(feature = "transport-websocket")]
-    WebSocket(String),
+    WebSocket(String, WebSocketSettings),
     #[cfg(feature = "transport-websocket")]
-    WebSocketSecure(String),
+    WebSocketSecure(String, WebSocketSettings),
     #[cfg(feature = "transport-quic")]
     Quic,
     #[cfg(feature = "transport-quic")]
     QuicSecure,
 }
 
+impl ClientTransportType {
+    pub(crate) fn redirected(self) -> Self {
+        match self {
+            #[cfg(feature = "transport-websocket")]
+            Self::WebSocket(url, _) => Self::WebSocket(url, WebSocketSettings::Default),
+            #[cfg(feature = "transport-websocket")]
+            Self::WebSocketSecure(url, _) => Self::WebSocketSecure(url, WebSocketSettings::Default),
+            other => other,
+        }
+    }
+}
+
 impl MqttClient {
     pub(crate) async fn trigger_connection_event(&self, event: ConnectionEvent) {
         super::fire_connection_event(&self.connection_event_callbacks, event).await;

crates/mqtt5/src/transport/websocket.rs

@@ -97,6 +97,7 @@ const RESERVED_HEADERS: [&str; 14] = [
 ];
 
 /// WebSocket transport configuration
+#[derive(Clone)]
 pub struct WebSocketConfig {
     /// WebSocket URL (ws:// or wss://)
     pub url: Url,
@@ -135,20 +136,8 @@ impl WebSocketConfig {
     ///
     /// Returns an error if the URL is invalid or uses an unsupported scheme
     pub fn new(url: &str) -> Result<Self> {
-        let parsed_url = Url::parse(url)
-            .map_err(|e| MqttError::ProtocolError(format!("Invalid WebSocket URL: {e}")))?;
-
-        match parsed_url.scheme() {
-            "ws" | "wss" => {}
-            scheme => {
-                return Err(MqttError::ProtocolError(format!(
-                    "Unsupported WebSocket scheme: {scheme}. Use 'ws' or 'wss'"
-                )));
-            }
-        }
-
         Ok(Self {
-            url: parsed_url,
+            url: parse_websocket_url(url)?,
             timeout: Duration::from_secs(30),
             subprotocols: vec![DEFAULT_SUBPROTOCOL.to_string()],
             headers: HashMap::new(),
@@ -157,6 +146,16 @@ impl WebSocketConfig {
         })
     }
 
+    /// Returns this configuration with its URL replaced
+    ///
+    /// Used by `MqttClient`, which connects to the address it is given with
+    /// the rest of the configuration from
+    /// [`ConnectOptions::with_websocket_config`](crate::ConnectOptions::with_websocket_config).
+    pub(crate) fn with_url(mut self, url: &str) -> Result<Self> {
+        self.url = parse_websocket_url(url)?;
+        Ok(self)
+    }
+
     /// Sets the connection timeout
     #[must_use]
     pub fn with_timeout(mut self, timeout: Duration) -> Self {
@@ -510,6 +509,19 @@ fn header_value(name: &str, value: &str) -> Result<HeaderValue> {
     })
 }
 
+/// Parses a WebSocket URL, accepting only the ws and wss schemes
+fn parse_websocket_url(url: &str) -> Result<Url> {
+    let parsed_url = Url::parse(url)
+        .map_err(|e| MqttError::ProtocolError(format!("Invalid WebSocket URL: {e}")))?;
+
+    match parsed_url.scheme() {
+        "ws" | "wss" => Ok(parsed_url),
+        scheme => Err(MqttError::ProtocolError(format!(
+            "Unsupported WebSocket scheme: {scheme}. Use 'ws' or 'wss'"
+        ))),
+    }
+}
+
 /// Whether `s` is an HTTP token (RFC 9110 §5.6.2), as RFC 6455 §4.1 requires
 /// of each offered subprotocol.
 fn is_http_token(s: &str) -> bool {