1use std::{future::Future, time::Duration};
5
6use polyoxide_rt::{default_runtime, Runtime};
7#[cfg(not(target_arch = "wasm32"))]
8use tokio::net::TcpStream;
9#[cfg(not(target_arch = "wasm32"))]
10use tokio_tungstenite::client_async_tls;
11
12use crate::{
13 default_socket_transport, tungstenite::client::IntoClientRequest, BoxSocket, Kind, WsError,
14};
15
16pub const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
19
20pub async fn connect<R>(request: R, timeout: Duration) -> Result<BoxSocket, WsError>
30where
31 R: IntoClientRequest + Unpin,
32{
33 let request = request.into_client_request()?;
34 default_socket_transport()
35 .connect(request, timeout, default_runtime())
36 .await
37}
38
39#[cfg(not(target_arch = "wasm32"))]
49pub async fn connect_over<R>(request: R, tcp: TcpStream) -> Result<BoxSocket, WsError>
50where
51 R: IntoClientRequest + Unpin,
52{
53 crate::ensure_crypto_provider();
54 let (socket, _response) = client_async_tls(request, tcp).await?;
55 Ok(Box::new(socket))
56}
57
58pub async fn with_timeout<F: Future>(timeout: Duration, connect: F) -> Result<F::Output, WsError> {
67 let runtime = default_runtime();
68 bounded(runtime.as_ref(), timeout, connect).await
69}
70
71pub(crate) async fn bounded<F: Future>(
75 runtime: &dyn Runtime,
76 timeout: Duration,
77 future: F,
78) -> Result<F::Output, WsError> {
79 polyoxide_rt::timeout(runtime, timeout, future)
80 .await
81 .ok_or_else(|| WsError::connect_timeout(timeout))
82}
83
84pub fn retrying_can_fix(err: &WsError, retriable_status: impl Fn(u16) -> bool) -> bool {
108 match err.kind() {
109 Kind::Refused { status, .. } => retriable_status(status),
110 Kind::Url
111 | Kind::HttpFormat
112 | Kind::Tls
113 | Kind::AttackAttempt
114 | Kind::Capacity
115 | Kind::Unsupported => false,
116 _ => true,
117 }
118}
119
120#[cfg(test)]
121mod tests {
122 use tokio_tungstenite::tungstenite::{
123 self,
124 error::{CapacityError, ProtocolError, TlsError, UrlError},
125 http,
126 };
127
128 use super::*;
129 use crate::test_server::{Script, ScriptedServer};
130
131 fn rule(status: u16) -> bool {
135 matches!(status, 408 | 425 | 429 | 500..=599)
136 }
137
138 fn refused(status: u16) -> WsError {
139 WsError::from(tungstenite::Error::Http(
140 http::Response::builder().status(status).body(None).unwrap(),
141 ))
142 }
143
144 #[tokio::test(start_paused = true)]
145 async fn a_hung_handshake_times_out_at_the_default() {
146 let server = ScriptedServer::start(vec![Script::stall()]).await;
147 let err = connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
148 .await
149 .err()
150 .expect("a server that never upgrades must not connect");
151 assert_eq!(
152 err.kind(),
153 Kind::ConnectTimeout {
154 after: Duration::from_secs(10)
155 }
156 );
157 }
158
159 #[tokio::test]
160 async fn a_refused_upgrade_keeps_its_status_and_retry_after() {
161 let server = ScriptedServer::start(vec![
162 Script::refuse_upgrade(429).with_header("Retry-After", "7")
163 ])
164 .await;
165 let err = connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
166 .await
167 .err()
168 .expect("a refused upgrade must not connect");
169 assert_eq!(
170 err.kind(),
171 Kind::Refused {
172 status: 429,
173 retry_after: Some("7")
174 }
175 );
176 }
177
178 #[tokio::test]
179 async fn a_503_upgrade_is_worth_retrying_and_a_404_is_not() {
180 let server = ScriptedServer::start(vec![
181 Script::refuse_upgrade(503),
182 Script::refuse_upgrade(404),
183 ])
184 .await;
185 let unwell = connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
186 .await
187 .err()
188 .expect("503");
189 assert_eq!(
190 unwell.kind(),
191 Kind::Refused {
192 status: 503,
193 retry_after: None
194 }
195 );
196 assert!(retrying_can_fix(&unwell, rule));
197
198 let missing = connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
199 .await
200 .err()
201 .expect("404");
202 assert!(!retrying_can_fix(&missing, rule));
203 }
204
205 #[tokio::test]
206 async fn a_server_that_upgrades_connects() {
207 let server = ScriptedServer::start(vec![Script::hold()]).await;
208 connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
209 .await
210 .expect("connect");
211 }
212
213 #[tokio::test]
214 async fn with_timeout_bounds_a_callers_own_future_and_keeps_its_output() {
215 let after = Duration::from_millis(20);
216 let err = with_timeout(after, std::future::pending::<()>())
217 .await
218 .expect_err("a future that never finishes must time out");
219 assert_eq!(err.kind(), Kind::ConnectTimeout { after });
220
221 let output: Result<u8, &str> = with_timeout(after, async { Err("its own error") })
222 .await
223 .expect("a finished future is not a timeout");
224 assert_eq!(output, Err("its own error"));
225 }
226
227 async fn dial(server: &ScriptedServer) -> TcpStream {
229 let addr = server.url.trim_start_matches("ws://");
230 TcpStream::connect(addr).await.expect("dial")
231 }
232
233 #[tokio::test]
234 async fn connect_over_upgrades_a_connection_the_caller_opened() {
235 let server = ScriptedServer::start(vec![Script::hold()]).await;
236 let tcp = dial(&server).await;
237 connect_over(server.url.as_str(), tcp)
238 .await
239 .expect("upgrade");
240 server
241 .wait_for("the handshake", |s| s.handshake_count() == 1)
242 .await;
243 }
244
245 #[tokio::test]
246 async fn connect_over_keeps_a_refusals_status_and_retry_after() {
247 let server = ScriptedServer::start(vec![
248 Script::refuse_upgrade(429).with_header("Retry-After", "7")
249 ])
250 .await;
251 let tcp = dial(&server).await;
252 let err = connect_over(server.url.as_str(), tcp)
253 .await
254 .err()
255 .expect("a refused upgrade must not connect");
256 assert_eq!(
257 err.kind(),
258 Kind::Refused {
259 status: 429,
260 retry_after: Some("7")
261 }
262 );
263 }
264
265 #[tokio::test]
266 async fn connect_over_is_bounded_by_with_timeout() {
267 let server = ScriptedServer::start(vec![Script::stall()]).await;
270 let tcp = dial(&server).await;
271 let after = Duration::from_millis(100);
272 let err = with_timeout(after, connect_over(server.url.as_str(), tcp))
273 .await
274 .err()
275 .expect("a stalled upgrade must time out");
276 assert_eq!(err.kind(), Kind::ConnectTimeout { after });
277 }
278
279 #[test]
280 fn the_status_rule_decides_a_refused_upgrade() {
281 for status in [301, 401, 403, 404, 418, 451, 600] {
282 assert!(!retrying_can_fix(&refused(status), rule), "{status}");
283 }
284 for status in [408, 425, 429, 500, 502, 503] {
285 assert!(retrying_can_fix(&refused(status), rule), "{status}");
286 }
287 }
288
289 #[test]
290 fn a_failure_that_replays_on_every_attempt_is_not_retried() {
291 let bad_header = http::header::HeaderName::from_bytes(b"in valid").unwrap_err();
292 for err in [
293 tungstenite::Error::Url(UrlError::NoHostName),
294 tungstenite::Error::HttpFormat(bad_header.into()),
295 tungstenite::Error::Tls(TlsError::InvalidDnsName),
296 tungstenite::Error::AttackAttempt,
297 tungstenite::Error::Capacity(CapacityError::TooManyHeaders),
298 tungstenite::Error::Io(std::io::ErrorKind::Unsupported.into()),
299 ] {
300 let err = WsError::from(err);
301 assert!(!retrying_can_fix(&err, |_| true), "{err:?}");
302 }
303 }
304
305 #[test]
306 fn a_transport_failure_or_timeout_is_retried() {
307 let after = Duration::from_secs(10);
308 for err in [
309 WsError::from(tungstenite::Error::Io(
310 std::io::ErrorKind::ConnectionReset.into(),
311 )),
312 WsError::from(tungstenite::Error::Protocol(
313 ProtocolError::ResetWithoutClosingHandshake,
314 )),
315 WsError::from(tungstenite::Error::ConnectionClosed),
316 WsError::from(tungstenite::Error::Utf8),
317 WsError::connect_timeout(after),
318 WsError::no_answer(after),
319 WsError::from(tungstenite::Error::AlreadyClosed),
321 WsError::closed(Some(1008), "policy"),
322 ] {
323 assert!(retrying_can_fix(&err, |_| false), "{err:?}");
324 }
325 }
326}