1use std::{future::Future, time::Duration};
5
6use tokio::net::TcpStream;
7use tokio_tungstenite::{
8 client_async_tls, connect_async, tungstenite::client::IntoClientRequest, MaybeTlsStream,
9 WebSocketStream,
10};
11
12use crate::{ensure_crypto_provider, Kind, WsError};
13
14pub const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
17
18pub type Socket = WebSocketStream<MaybeTlsStream<TcpStream>>;
20
21pub async fn connect<R>(request: R, timeout: Duration) -> Result<Socket, WsError>
36where
37 R: IntoClientRequest + Unpin,
38{
39 ensure_crypto_provider();
40 let (socket, _response) = with_timeout(timeout, connect_async(request)).await??;
41 Ok(socket)
42}
43
44pub async fn connect_over<R>(request: R, tcp: TcpStream) -> Result<Socket, WsError>
54where
55 R: IntoClientRequest + Unpin,
56{
57 ensure_crypto_provider();
58 let (socket, _response) = client_async_tls(request, tcp).await?;
59 Ok(socket)
60}
61
62pub async fn with_timeout<F: Future>(timeout: Duration, connect: F) -> Result<F::Output, WsError> {
69 tokio::time::timeout(timeout, connect)
70 .await
71 .map_err(|_| WsError::connect_timeout(timeout))
72}
73
74pub fn retrying_can_fix(err: &WsError, retriable_status: impl Fn(u16) -> bool) -> bool {
97 match err.kind() {
98 Kind::Refused { status, .. } => retriable_status(status),
99 Kind::Url | Kind::HttpFormat | Kind::Tls | Kind::AttackAttempt | Kind::Capacity => false,
100 _ => true,
101 }
102}
103
104#[cfg(test)]
105mod tests {
106 use tokio_tungstenite::tungstenite::{
107 self,
108 error::{CapacityError, ProtocolError, TlsError, UrlError},
109 http,
110 };
111
112 use super::*;
113 use crate::test_server::{Script, ScriptedServer};
114
115 fn rule(status: u16) -> bool {
119 matches!(status, 408 | 425 | 429 | 500..=599)
120 }
121
122 fn refused(status: u16) -> WsError {
123 WsError::from(tungstenite::Error::Http(
124 http::Response::builder().status(status).body(None).unwrap(),
125 ))
126 }
127
128 #[tokio::test(start_paused = true)]
129 async fn a_hung_handshake_times_out_at_the_default() {
130 let server = ScriptedServer::start(vec![Script::stall()]).await;
131 let err = connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
132 .await
133 .expect_err("a server that never upgrades must not connect");
134 assert_eq!(
135 err.kind(),
136 Kind::ConnectTimeout {
137 after: Duration::from_secs(10)
138 }
139 );
140 }
141
142 #[tokio::test]
143 async fn a_refused_upgrade_keeps_its_status_and_retry_after() {
144 let server = ScriptedServer::start(vec![
145 Script::refuse_upgrade(429).with_header("Retry-After", "7")
146 ])
147 .await;
148 let err = connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
149 .await
150 .expect_err("a refused upgrade must not connect");
151 assert_eq!(
152 err.kind(),
153 Kind::Refused {
154 status: 429,
155 retry_after: Some("7")
156 }
157 );
158 }
159
160 #[tokio::test]
161 async fn a_503_upgrade_is_worth_retrying_and_a_404_is_not() {
162 let server = ScriptedServer::start(vec![
163 Script::refuse_upgrade(503),
164 Script::refuse_upgrade(404),
165 ])
166 .await;
167 let unwell = connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
168 .await
169 .expect_err("503");
170 assert_eq!(
171 unwell.kind(),
172 Kind::Refused {
173 status: 503,
174 retry_after: None
175 }
176 );
177 assert!(retrying_can_fix(&unwell, rule));
178
179 let missing = connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
180 .await
181 .expect_err("404");
182 assert!(!retrying_can_fix(&missing, rule));
183 }
184
185 #[tokio::test]
186 async fn a_server_that_upgrades_connects() {
187 let server = ScriptedServer::start(vec![Script::hold()]).await;
188 connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
189 .await
190 .expect("connect");
191 }
192
193 #[tokio::test]
194 async fn with_timeout_bounds_a_callers_own_future_and_keeps_its_output() {
195 let after = Duration::from_millis(20);
196 let err = with_timeout(after, std::future::pending::<()>())
197 .await
198 .expect_err("a future that never finishes must time out");
199 assert_eq!(err.kind(), Kind::ConnectTimeout { after });
200
201 let output: Result<u8, &str> = with_timeout(after, async { Err("its own error") })
202 .await
203 .expect("a finished future is not a timeout");
204 assert_eq!(output, Err("its own error"));
205 }
206
207 async fn dial(server: &ScriptedServer) -> TcpStream {
209 let addr = server.url.trim_start_matches("ws://");
210 TcpStream::connect(addr).await.expect("dial")
211 }
212
213 #[tokio::test]
214 async fn connect_over_upgrades_a_connection_the_caller_opened() {
215 let server = ScriptedServer::start(vec![Script::hold()]).await;
216 let tcp = dial(&server).await;
217 connect_over(server.url.as_str(), tcp)
218 .await
219 .expect("upgrade");
220 server
221 .wait_for("the handshake", |s| s.handshake_count() == 1)
222 .await;
223 }
224
225 #[tokio::test]
226 async fn connect_over_keeps_a_refusals_status_and_retry_after() {
227 let server = ScriptedServer::start(vec![
228 Script::refuse_upgrade(429).with_header("Retry-After", "7")
229 ])
230 .await;
231 let tcp = dial(&server).await;
232 let err = connect_over(server.url.as_str(), tcp)
233 .await
234 .expect_err("a refused upgrade must not connect");
235 assert_eq!(
236 err.kind(),
237 Kind::Refused {
238 status: 429,
239 retry_after: Some("7")
240 }
241 );
242 }
243
244 #[tokio::test]
245 async fn connect_over_is_bounded_by_with_timeout() {
246 let server = ScriptedServer::start(vec![Script::stall()]).await;
249 let tcp = dial(&server).await;
250 let after = Duration::from_millis(100);
251 let err = with_timeout(after, connect_over(server.url.as_str(), tcp))
252 .await
253 .expect_err("a stalled upgrade must time out");
254 assert_eq!(err.kind(), Kind::ConnectTimeout { after });
255 }
256
257 #[test]
258 fn the_status_rule_decides_a_refused_upgrade() {
259 for status in [301, 401, 403, 404, 418, 451, 600] {
260 assert!(!retrying_can_fix(&refused(status), rule), "{status}");
261 }
262 for status in [408, 425, 429, 500, 502, 503] {
263 assert!(retrying_can_fix(&refused(status), rule), "{status}");
264 }
265 }
266
267 #[test]
268 fn a_failure_that_replays_on_every_attempt_is_not_retried() {
269 let bad_header = http::header::HeaderName::from_bytes(b"in valid").unwrap_err();
270 for err in [
271 tungstenite::Error::Url(UrlError::NoHostName),
272 tungstenite::Error::HttpFormat(bad_header.into()),
273 tungstenite::Error::Tls(TlsError::InvalidDnsName),
274 tungstenite::Error::AttackAttempt,
275 tungstenite::Error::Capacity(CapacityError::TooManyHeaders),
276 ] {
277 let err = WsError::from(err);
278 assert!(!retrying_can_fix(&err, |_| true), "{err:?}");
279 }
280 }
281
282 #[test]
283 fn a_transport_failure_or_timeout_is_retried() {
284 let after = Duration::from_secs(10);
285 for err in [
286 WsError::from(tungstenite::Error::Io(
287 std::io::ErrorKind::ConnectionReset.into(),
288 )),
289 WsError::from(tungstenite::Error::Protocol(
290 ProtocolError::ResetWithoutClosingHandshake,
291 )),
292 WsError::from(tungstenite::Error::ConnectionClosed),
293 WsError::from(tungstenite::Error::Utf8),
294 WsError::connect_timeout(after),
295 WsError::no_answer(after),
296 WsError::from(tungstenite::Error::AlreadyClosed),
298 WsError::closed(Some(1008), "policy"),
299 ] {
300 assert!(retrying_can_fix(&err, |_| false), "{err:?}");
301 }
302 }
303}