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

wsio_client/
builder.rs
