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::protocol::WebSocketConfig;
15use url::Url;
16
17use crate::{
18 WsIoClient,
19 config::WsIoClientConfig,
20 core::packet::codecs::WsIoPacketCodec,
21 runtime::WsIoClientRuntime,
22 session::WsIoClientSession,
23};
24
25pub struct WsIoClientBuilder {
27 config: WsIoClientConfig,
28 connect_url: Url,
29}
30
31impl WsIoClientBuilder {
32 pub(crate) fn new(mut url: Url) -> Result<Self> {
33 if !matches!(url.scheme(), "ws" | "wss") {
34 bail!("Invalid URL scheme: {}", url.scheme());
35 }
36
37 let mut query_pairs = url.query_pairs().collect::<Vec<_>>();
38 query_pairs.retain(|(k, _)| k != "namespace");
39 query_pairs.push(("namespace".into(), Self::normalize_url_path(url.path()).into()));
40 let query = query_pairs
41 .iter()
42 .map(|(k, v)| format!("{k}={v}"))
43 .collect::<Vec<_>>()
44 .join("&");
45
46 url.set_query(Some(&query));
47 url.set_path("ws.io");
48 Ok(Self {
49 config: WsIoClientConfig {
50 init_handler: None,
51 init_handler_timeout: Duration::from_secs(3),
52 init_packet_timeout: Duration::from_secs(5),
53 on_session_close_handler: None,
54 on_session_close_handler_timeout: Duration::from_secs(2),
55 on_session_ready_handler: None,
56 packet_codec: WsIoPacketCodec::SerdeJson,
57 ping_interval: Duration::from_secs(25),
58 ready_packet_timeout: Duration::from_secs(5),
59 reconnect_delay: Duration::from_secs(1),
60 websocket_config: WebSocketConfig::default()
61 .max_frame_size(Some(8 * 1024 * 1024))
62 .max_message_size(Some(16 * 1024 * 1024))
63 .max_write_buffer_size(2 * 1024 * 1024)
64 .read_buffer_size(8 * 1024)
65 .write_buffer_size(8 * 1024),
66 },
67 connect_url: url,
68 })
69 }
70
71 fn normalize_url_path(path: &str) -> String {
73 format!(
74 "/{}",
75 path.split('/').filter(|s| !s.is_empty()).collect::<Vec<_>>().join("/")
76 )
77 }
78
79 pub fn build(self) -> WsIoClient {
81 WsIoClient(WsIoClientRuntime::new(self.config, self.connect_url))
82 }
83
84 pub fn init_handler_timeout(mut self, duration: Duration) -> Self {
85 self.config.init_handler_timeout = duration;
86 self
87 }
88
89 pub fn init_packet_timeout(mut self, duration: Duration) -> Self {
90 self.config.init_packet_timeout = duration;
91 self
92 }
93
94 pub fn on_session_close<H, Fut>(mut self, handler: H) -> Self
95 where
96 H: Fn(Arc<WsIoClientSession>) -> Fut + Send + Sync + 'static,
97 Fut: Future<Output = Result<()>> + Send + 'static,
98 {
99 self.config.on_session_close_handler = Some(Box::new(move |session| Box::pin(handler(session))));
100 self
101 }
102
103 pub fn on_session_close_handler_timeout(mut self, duration: Duration) -> Self {
104 self.config.on_session_close_handler_timeout = duration;
105 self
106 }
107
108 pub fn on_session_ready<H, Fut>(mut self, handler: H) -> Self
109 where
110 H: Fn(Arc<WsIoClientSession>) -> Fut + Send + Sync + 'static,
111 Fut: Future<Output = Result<()>> + Send + 'static,
112 {
113 self.config.on_session_ready_handler = Some(Arc::new(move |session| Box::pin(handler(session))));
114 self
115 }
116
117 pub fn packet_codec(mut self, packet_codec: WsIoPacketCodec) -> Self {
118 self.config.packet_codec = packet_codec;
119 self
120 }
121
122 pub fn ping_interval(mut self, duration: Duration) -> Self {
123 self.config.ping_interval = duration;
124 self
125 }
126
127 pub fn ready_packet_timeout(mut self, duration: Duration) -> Self {
128 self.config.ready_packet_timeout = duration;
129 self
130 }
131
132 pub fn reconnect_delay(mut self, delay: Duration) -> Self {
133 self.config.reconnect_delay = delay;
134 self
135 }
136
137 pub fn request_path(mut self, request_path: impl AsRef<str>) -> Self {
138 self.connect_url
139 .set_path(&Self::normalize_url_path(request_path.as_ref()));
140
141 self
142 }
143
144 pub fn websocket_config(mut self, websocket_config: WebSocketConfig) -> Self {
145 self.config.websocket_config = websocket_config;
146 self
147 }
148
149 pub fn websocket_config_mut<F: FnOnce(&mut WebSocketConfig)>(mut self, f: F) -> Self {
150 f(&mut self.config.websocket_config);
151 self
152 }
153
154 pub fn with_init_handler<H, Fut, D, R>(mut self, handler: H) -> WsIoClientBuilder
155 where
156 H: Fn(Arc<WsIoClientSession>, Option<D>) -> Fut + Send + Sync + 'static,
157 Fut: Future<Output = Result<Option<R>>> + Send + 'static,
158 D: DeserializeOwned + Send + 'static,
159 R: Serialize + Send + 'static,
160 {
161 let handler = Arc::new(handler);
162 self.config.init_handler = Some(Box::new(move |session, bytes, packet_codec| {
163 let handler = handler.clone();
164 Box::pin(async move {
165 handler(session, bytes.map(|bytes| packet_codec.decode_data(bytes)).transpose()?)
166 .await?
167 .map(|data| packet_codec.encode_data(&data))
168 .transpose()
169 })
170 }));
171
172 self
173 }
174}
175
176#[cfg(test)]
177mod tests {
178 use super::*;
179
180 #[test]
181 fn test_builder_new_valid_ws_url() {
182 let result = WsIoClientBuilder::new(Url::parse("ws://localhost:8080/socket").unwrap());
183 assert!(result.is_ok());
184 }
185
186 #[test]
187 fn test_builder_new_valid_wss_url() {
188 let result = WsIoClientBuilder::new(Url::parse("wss://localhost:8080/socket").unwrap());
189 assert!(result.is_ok());
190 }
191
192 #[test]
193 fn test_builder_new_invalid_scheme() {
194 let result = WsIoClientBuilder::new(Url::parse("http://localhost:8080/socket").unwrap());
195 assert!(result.is_err());
196 if let Err(e) = result {
197 let err_msg = format!("{e}");
198 assert!(err_msg.contains("Invalid URL scheme"));
199 }
200 }
201
202 #[test]
203 fn test_builder_configuration_chaining() {
204 WsIoClientBuilder::new(Url::parse("ws://localhost:8080/socket").unwrap())
205 .unwrap()
206 .init_handler_timeout(Duration::from_secs(10))
207 .init_packet_timeout(Duration::from_secs(15))
208 .on_session_close_handler_timeout(Duration::from_secs(5))
209 .packet_codec(WsIoPacketCodec::SerdeJson)
210 .ping_interval(Duration::from_secs(30))
211 .ready_packet_timeout(Duration::from_secs(10))
212 .reconnect_delay(Duration::from_secs(5))
213 .request_path("/custom/path")
214 .build();
215 }
216
217 #[test]
218 fn test_builder_request_path_normalizes() {
219 WsIoClientBuilder::new(Url::parse("ws://localhost:8080/socket").unwrap())
221 .unwrap()
222 .request_path("/multiple//slashes///path/")
223 .build();
224 }
225
226 #[test]
227 fn test_builder_websocket_config_override() {
228 WsIoClientBuilder::new(Url::parse("ws://localhost:8080/socket").unwrap())
229 .unwrap()
230 .websocket_config_mut(|config| {
231 *config = config.max_frame_size(Some(1024 * 1024));
232 })
233 .build();
234 }
235
236 #[test]
237 fn test_builder_all_timeout_configurations() {
238 WsIoClientBuilder::new(Url::parse("ws://localhost:8080/socket").unwrap())
239 .unwrap()
240 .init_handler_timeout(Duration::from_secs(1))
241 .init_packet_timeout(Duration::from_secs(2))
242 .on_session_close_handler_timeout(Duration::from_secs(3))
243 .ready_packet_timeout(Duration::from_secs(4))
244 .build();
245 }
246
247 #[test]
248 fn test_builder_reconnect_delay_configuration() {
249 WsIoClientBuilder::new(Url::parse("ws://localhost:8080/socket").unwrap())
250 .unwrap()
251 .reconnect_delay(Duration::from_millis(500))
252 .build();
253 }
254}