Skip to main content

polyoxide_ws/
connect.rs

1//! Opening a connection within a timeout (W8), and whether retrying a failed
2//! one can help (W7).
3
4use 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
16/// How long one connection attempt may take unless a venue says otherwise:
17/// the TCP connect, the TLS handshake and the WebSocket upgrade together.
18pub const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
19
20/// Open a connection through the target's default [`SocketTransport`](crate::SocketTransport)
21/// and runtime, giving up after `timeout`.
22///
23/// `request` is a URL, or a handshake request with its own headers, built
24/// afresh for each attempt so any authentication in it is recomputed. A
25/// connection not open within `timeout` fails as [`Kind::ConnectTimeout`];
26/// every other failure is the transport's, such as [`Kind::Refused`] for an
27/// upgrade the server answered with an HTTP status. A caller that injects its
28/// own transport or runtime opens its connections through them instead.
29pub 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/// Upgrade a TCP connection the caller opened itself: the TLS handshake when
40/// `request` is a `wss://` URL, then the WebSocket upgrade.
41///
42/// For a transport whose connect is its own: one that resolves the host and
43/// dials each address itself, and hands the first TCP connection that opens
44/// to this. The rustls provider is installed first
45/// ([`ensure_crypto_provider`](crate::ensure_crypto_provider)). It has no
46/// timeout of its own, so bound it with [`with_timeout`]. A refused upgrade
47/// fails as [`Kind::Refused`], as it does through [`connect`]. Native only.
48#[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
58/// Run a caller's own connect future, giving up after `timeout` with
59/// [`Kind::ConnectTimeout`].
60///
61/// The future's own output comes back as it is, so its errors stay the
62/// caller's. It runs on the target's default runtime; a transport bounds its
63/// own steps on the runtime it is given instead. Install the rustls provider
64/// before a TLS connect of its own
65/// ([`ensure_crypto_provider`](crate::ensure_crypto_provider)).
66pub 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
71/// Run `future` within `timeout` on `runtime`, giving up with
72/// [`Kind::ConnectTimeout`]. A transport bounds each of its own steps with the
73/// runtime it is given.
74pub(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
84/// Whether reconnecting can help after a connect failed this way.
85///
86/// A refused upgrade is decided by `retriable_status`, the status rule. Pass
87/// `polyoxide-venue`'s, which `Classify::is_retriable` reads too, so the two
88/// can never disagree on a status:
89/// `|status| class_for_handshake_status(status).is_retriable()`.
90///
91/// A bad URL, an upgrade request that could not be built, a refused TLS
92/// configuration, an attack detection, a capacity limit and an operation the
93/// transport cannot perform replay the same failure on every attempt, so they
94/// are `false`. Everything else, an I/O failure, a timeout or a lost
95/// connection among them, is `true`.
96///
97/// This decides only whether a reconnect can help, not what a failure is:
98/// that is the venue's classification. What to do after
99/// [`Kind::AlreadyClosed`] (D19) and what a close code means stay with each
100/// venue's recovery, so both are `true` here.
101///
102/// The basis is one socket crate's former `retrying_can_fix` on its error
103/// (`src/error.rs`), with the status rule injected, since this crate cannot
104/// see `polyoxide-venue`. That crate's recovery calls this since Story 4.10,
105/// with its own row ahead of it: it retries a `Capacity` failure, which this
106/// gives up on.
107pub 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    /// A stand-in for `polyoxide-venue`'s status rule, which this crate
132    /// cannot see; venue's table test checks the real one against the
133    /// classification.
134    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    /// A TCP connection to `server`, opened as a venue's own connect would.
228    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        // It has no timeout of its own; a server that never answers the
268        // upgrade holds it until the caller's bound.
269        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            // D19 and a close code's meaning are the venue's to decide.
320            WsError::from(tungstenite::Error::AlreadyClosed),
321            WsError::closed(Some(1008), "policy"),
322        ] {
323            assert!(retrying_can_fix(&err, |_| false), "{err:?}");
324        }
325    }
326}