wsio_client/
builder.rs

1use std::{
2    sync::Arc,
3    time::Duration,
4};
5
6use anyhow::{
7    Result,
8    bail,
9};
10use serde::{
11    Serialize,
12    de::DeserializeOwned,
13};
14use tokio_tungstenite::tungstenite::{
15    http::Request,
16    protocol::WebSocketConfig,
17};
18use url::Url;
19
20use crate::{
21    WsIoClient,
22    config::WsIoClientConfig,
23    core::packet::codecs::WsIoPacketCodec,
24    runtime::WsIoClientRuntime,
25    session::WsIoClientSession,
26};
27
28// Structs
29
30/// Builder for configuring and creating a [`WsIoClient`].
31///
32/// The URL passed to the client constructor selects the namespace from its path,
33/// while the actual WebSocket request path defaults to `/ws.io`.
34#[derive(Debug)]
35pub struct WsIoClientBuilder {
36    config: WsIoClientConfig,
37    connect_url: Url,
38}
39
40impl WsIoClientBuilder {
41    pub(crate) fn new(mut url: Url) -> Result<Self> {
42        if !matches!(url.scheme(), "ws" | "wss") {
43            bail!("Invalid URL scheme: {}", url.scheme());
44        }
45
46        let mut query_pairs = url.query_pairs().collect::<Vec<_>>();
47        query_pairs.retain(|(k, _)| k != "namespace");
48        query_pairs.push(("namespace".into(), Self::normalize_url_path(url.path()).into()));
49        let query = query_pairs
50            .iter()
51            .map(|(k, v)| format!("{k}={v}"))
52            .collect::<Vec<_>>()
53            .join("&");
54
55        url.set_query(Some(&query));
56        url.set_path("ws.io");
57        Ok(Self {
58            config: WsIoClientConfig {
59                connect_timeout: Some(Duration::from_secs(10)),
60                disconnect_timeout: Duration::from_secs(5),
61                init_handler: None,
62                init_handler_timeout: Duration::from_secs(3),
63                init_packet_timeout: Duration::from_secs(5),
64                on_session_close_handler: None,
65                on_session_close_handler_timeout: Duration::from_secs(2),
66                on_session_ready_handler: None,
67                packet_codec: WsIoPacketCodec::SerdeJson,
68                ping_interval: Duration::from_secs(25),
69                ready_packet_timeout: Duration::from_secs(5),
70                reconnect_delay: Duration::from_secs(1),
71                request_modifier: None,
72                websocket_config: WebSocketConfig::default()
73                    .max_frame_size(Some(8 * 1024 * 1024))
74                    .max_message_size(Some(16 * 1024 * 1024))
75                    .max_write_buffer_size(2 * 1024 * 1024)
76                    .read_buffer_size(8 * 1024)
77                    .write_buffer_size(8 * 1024),
78            },
79            connect_url: url,
80        })
81    }
82
83    // Private methods
84    fn normalize_url_path(path: &str) -> String {
85        format!(
86            "/{}",
87            path.split('/').filter(|s| !s.is_empty()).collect::<Vec<_>>().join("/")
88        )
89    }
90
91    // Public methods
92
93    /// Builds a [`WsIoClient`] with the accumulated configuration.
94    pub fn build(self) -> WsIoClient {
95        WsIoClient(WsIoClientRuntime::new(self.config, self.connect_url))
96    }
97
98    /// Sets how long the client waits for the WebSocket connection attempt.
99    ///
100    /// This timeout covers the transport connection and WebSocket HTTP upgrade
101    /// handshake. It does not cover ws.io protocol initialization after the
102    /// WebSocket connection is established; use [`Self::init_packet_timeout`]
103    /// and [`Self::ready_packet_timeout`] for that phase. Pass `None` to disable
104    /// the connection-attempt timeout.
105    pub fn connect_timeout(mut self, duration: impl Into<Option<Duration>>) -> Self {
106        self.config.connect_timeout = duration.into();
107        self
108    }
109
110    /// Sets how long `disconnect().await` waits for graceful WebSocket
111    /// shutdown before aborting the connection read/write tasks.
112    pub fn disconnect_timeout(mut self, duration: Duration) -> Self {
113        self.config.disconnect_timeout = duration;
114        self
115    }
116
117    /// Sets the maximum duration allowed for the init handler to run.
118    ///
119    /// The init handler is registered with [`Self::with_init_handler`] and is
120    /// invoked after the server sends the init packet.
121    pub fn init_handler_timeout(mut self, duration: Duration) -> Self {
122        self.config.init_handler_timeout = duration;
123        self
124    }
125
126    /// Sets how long the client waits for the server init packet after the
127    /// WebSocket connection is established.
128    ///
129    /// If the init packet is not received before this timeout, the session is
130    /// closed and the runtime may reconnect according to [`Self::reconnect_delay`].
131    pub fn init_packet_timeout(mut self, duration: Duration) -> Self {
132        self.config.init_packet_timeout = duration;
133        self
134    }
135
136    /// Registers a handler that runs when a session closes.
137    ///
138    /// The handler is awaited during session cleanup and is bounded by
139    /// [`Self::on_session_close_handler_timeout`].
140    pub fn on_session_close<H, Fut>(mut self, handler: H) -> Self
141    where
142        H: Fn(Arc<WsIoClientSession>) -> Fut + Send + Sync + 'static,
143        Fut: Future<Output = Result<()>> + Send + 'static,
144    {
145        self.config.on_session_close_handler = Some(Box::new(move |session| Box::pin(handler(session))));
146        self
147    }
148
149    /// Sets the maximum duration allowed for the session-close handler to run.
150    pub fn on_session_close_handler_timeout(mut self, duration: Duration) -> Self {
151        self.config.on_session_close_handler_timeout = duration;
152        self
153    }
154
155    /// Registers a handler that runs after a session becomes ready.
156    ///
157    /// The handler is spawned asynchronously after the ready packet is received,
158    /// so it does not block the connection handshake.
159    pub fn on_session_ready<H, Fut>(mut self, handler: H) -> Self
160    where
161        H: Fn(Arc<WsIoClientSession>) -> Fut + Send + Sync + 'static,
162        Fut: Future<Output = Result<()>> + Send + 'static,
163    {
164        self.config.on_session_ready_handler = Some(Arc::new(move |session| Box::pin(handler(session))));
165        self
166    }
167
168    /// Sets the packet codec used to encode and decode ws.io protocol packets.
169    ///
170    /// This must match the server namespace codec.
171    pub fn packet_codec(mut self, packet_codec: WsIoPacketCodec) -> Self {
172        self.config.packet_codec = packet_codec;
173        self
174    }
175
176    /// Sets the interval for client heartbeat frames.
177    ///
178    /// After session initialization starts, the client periodically sends a
179    /// one-byte binary WebSocket frame. The server treats single-byte binary
180    /// frames as heartbeats and ignores them before packet decoding.
181    pub fn ping_interval(mut self, duration: Duration) -> Self {
182        self.config.ping_interval = duration;
183        self
184    }
185
186    /// Sets how long the client waits for the server ready packet.
187    ///
188    /// The ready timeout starts after the client handles the server init packet
189    /// and sends its init response.
190    pub fn ready_packet_timeout(mut self, duration: Duration) -> Self {
191        self.config.ready_packet_timeout = duration;
192        self
193    }
194
195    /// Sets the delay before the runtime attempts another connection.
196    ///
197    /// This delay is used after a connection attempt/session ends while the client
198    /// runtime is still running.
199    pub fn reconnect_delay(mut self, delay: Duration) -> Self {
200        self.config.reconnect_delay = delay;
201        self
202    }
203
204    /// Registers an async modifier for the WebSocket HTTP request.
205    ///
206    /// Use this to add headers or adjust request metadata before
207    /// `connect_async_with_config` is called.
208    pub fn request_modifier<M, Fut>(mut self, modifier: M) -> Self
209    where
210        M: Fn(Request<()>) -> Fut + Send + Sync + 'static,
211        Fut: Future<Output = Result<Request<()>>> + Send + 'static,
212    {
213        self.config.request_modifier = Some(Box::new(move |request| Box::pin(modifier(request))));
214        self
215    }
216
217    /// Sets the WebSocket HTTP request path.
218    ///
219    /// Paths are normalized to a single leading slash with empty path segments
220    /// removed. This controls the request URI path, not the namespace query value
221    /// inferred from the original URL passed to the builder.
222    pub fn request_path(mut self, request_path: impl AsRef<str>) -> Self {
223        self.connect_url
224            .set_path(&Self::normalize_url_path(request_path.as_ref()));
225
226        self
227    }
228
229    /// Replaces the full Tungstenite WebSocket configuration.
230    ///
231    /// This controls transport limits and buffer sizes passed to the WebSocket
232    /// connection. It is also used to derive internal channel capacity from the
233    /// configured max-write/write-buffer ratio.
234    pub fn websocket_config(mut self, websocket_config: WebSocketConfig) -> Self {
235        self.config.websocket_config = websocket_config;
236        self
237    }
238
239    /// Mutates the current Tungstenite WebSocket configuration in place.
240    ///
241    /// Prefer this when you want to adjust one or two fields while keeping the
242    /// builder defaults for the rest.
243    pub fn websocket_config_mut<F: FnOnce(&mut WebSocketConfig)>(mut self, f: F) -> Self {
244        f(&mut self.config.websocket_config);
245        self
246    }
247
248    /// Registers the client-side init handler.
249    ///
250    /// The handler receives the session and the optional server init payload
251    /// decoded as `D`. Its optional return value is encoded as `R` and sent back
252    /// to the server as the client init response.
253    pub fn with_init_handler<H, Fut, D, R>(mut self, handler: H) -> Self
254    where
255        H: Fn(Arc<WsIoClientSession>, Option<D>) -> Fut + Send + Sync + 'static,
256        Fut: Future<Output = Result<Option<R>>> + Send + 'static,
257        D: DeserializeOwned + Send + 'static,
258        R: Serialize + Send + 'static,
259    {
260        let handler = Arc::new(handler);
261        self.config.init_handler = Some(Box::new(move |session, bytes, packet_codec| {
262            let handler = handler.clone();
263            Box::pin(async move {
264                handler(session, bytes.map(|bytes| packet_codec.decode_data(bytes)).transpose()?)
265                    .await?
266                    .map(|data| packet_codec.encode_data(&data))
267                    .transpose()
268            })
269        }));
270
271        self
272    }
273}
274
275#[cfg(test)]
276mod tests {
277    use tokio_tungstenite::tungstenite::http::HeaderValue;
278
279    use super::*;
280
281    const TEST_URL: &str = "ws://localhost:8080/socket";
282
283    fn test_builder() -> WsIoClientBuilder {
284        WsIoClientBuilder::new(Url::parse(TEST_URL).unwrap()).unwrap()
285    }
286
287    #[test]
288    fn test_builder_new_valid_ws_url_sets_default_request_path_and_namespace_query() {
289        let builder = test_builder();
290
291        assert_eq!(builder.connect_url.path(), "/ws.io");
292        let namespace = builder
293            .connect_url
294            .query_pairs()
295            .find(|(key, _)| key == "namespace")
296            .unwrap()
297            .1;
298
299        assert_eq!(namespace, "/socket");
300    }
301
302    #[test]
303    fn test_builder_new_valid_wss_url() {
304        WsIoClientBuilder::new(Url::parse("wss://localhost:8080/socket").unwrap()).unwrap();
305    }
306
307    #[test]
308    fn test_builder_new_invalid_scheme() {
309        let error = WsIoClientBuilder::new(Url::parse("http://localhost:8080/socket").unwrap()).unwrap_err();
310        assert!(error.to_string().contains("Invalid URL scheme"));
311    }
312
313    #[test]
314    fn test_builder_configuration_chaining_updates_runtime_config() {
315        let builder = test_builder()
316            .connect_timeout(Duration::from_secs(25))
317            .disconnect_timeout(Duration::from_secs(20))
318            .init_handler_timeout(Duration::from_secs(10))
319            .init_packet_timeout(Duration::from_secs(15))
320            .on_session_close_handler_timeout(Duration::from_secs(5))
321            .packet_codec(WsIoPacketCodec::SerdeJson)
322            .ping_interval(Duration::from_secs(30))
323            .ready_packet_timeout(Duration::from_secs(10))
324            .reconnect_delay(Duration::from_secs(5))
325            .request_path("/custom/path");
326
327        assert_eq!(builder.connect_url.path(), "/custom/path");
328
329        let client = builder.build();
330
331        let config = &client.0.config;
332        assert_eq!(config.connect_timeout, Some(Duration::from_secs(25)));
333        assert_eq!(config.disconnect_timeout, Duration::from_secs(20));
334        assert_eq!(config.init_handler_timeout, Duration::from_secs(10));
335        assert_eq!(config.init_packet_timeout, Duration::from_secs(15));
336        assert_eq!(config.on_session_close_handler_timeout, Duration::from_secs(5));
337        assert!(matches!(config.packet_codec, WsIoPacketCodec::SerdeJson));
338        assert_eq!(config.ping_interval, Duration::from_secs(30));
339        assert_eq!(config.ready_packet_timeout, Duration::from_secs(10));
340        assert_eq!(config.reconnect_delay, Duration::from_secs(5));
341    }
342
343    #[test]
344    fn test_builder_request_path_normalizes() {
345        let builder = test_builder().request_path("/multiple//slashes///path/");
346
347        assert_eq!(builder.connect_url.path(), "/multiple/slashes/path");
348    }
349
350    #[test]
351    fn test_builder_websocket_config_override() {
352        let client = test_builder()
353            .websocket_config_mut(|config| {
354                *config = config.max_frame_size(Some(1024 * 1024));
355            })
356            .build();
357
358        assert_eq!(client.0.config.websocket_config.max_frame_size, Some(1024 * 1024));
359    }
360
361    #[test]
362    fn test_builder_websocket_config_replaces_defaults() {
363        let config = WebSocketConfig::default().max_frame_size(Some(42));
364        let client = test_builder().websocket_config(config).build();
365
366        assert_eq!(client.0.config.websocket_config.max_frame_size, Some(42));
367    }
368
369    #[test]
370    fn test_builder_with_init_and_session_handlers_registers_callbacks() {
371        let client = test_builder()
372            .with_init_handler(|_session, _data: Option<String>| async { Ok(Some("response".to_string())) })
373            .on_session_ready(|_session| async { Ok(()) })
374            .on_session_close(|_session| async { Ok(()) })
375            .build();
376
377        assert!(client.0.config.init_handler.is_some());
378        assert!(client.0.config.on_session_ready_handler.is_some());
379        assert!(client.0.config.on_session_close_handler.is_some());
380    }
381
382    #[test]
383    fn test_builder_request_modifier_registers_async_callback() {
384        let client = test_builder()
385            .request_modifier(|mut request| async move {
386                request
387                    .headers_mut()
388                    .insert("x-wsio-test", HeaderValue::from_static("enabled"));
389
390                Ok(request)
391            })
392            .build();
393
394        assert!(client.0.config.request_modifier.is_some());
395    }
396}
wsio_client/builder.rs

wsio_client/
builder.rs
