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 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
14/// How long one connection attempt may take unless a venue says otherwise:
15/// the TCP connect, the TLS handshake and the WebSocket upgrade together.
16pub const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
17
18/// An open connection, over TLS or plain TCP.
19pub type Socket = WebSocketStream<MaybeTlsStream<TcpStream>>;
20
21/// Open a connection, 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. The
25/// rustls provider is installed first ([`ensure_crypto_provider`]). A
26/// connection not open within `timeout` fails as
27/// [`Kind::ConnectTimeout`]; every other failure is the transport's, such as
28/// [`Kind::Refused`] for an upgrade the server answered with an HTTP status.
29///
30/// The basis is two socket crates' former connects, one crate's `open`
31/// (`src/client.rs`) and the other's bare-tier `open`, each on this function
32/// since Stories 4.7 and 4.10. A venue whose connect is its own bounds it
33/// with [`with_timeout`] instead, as one venue's per-address connect does
34/// (D20), each step bounded and the upgrade through [`connect_over`].
35pub 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
44/// Upgrade a TCP connection the caller opened itself: the TLS handshake when
45/// `request` is a `wss://` URL, then the WebSocket upgrade.
46///
47/// For a venue whose connect is its own: one that resolves the host and
48/// dials each address itself, and hands the first TCP connection that opens
49/// to this (D20). The rustls provider is installed first
50/// ([`ensure_crypto_provider`]). It has no timeout of its own, so bound it
51/// with [`with_timeout`]. A refused upgrade fails as [`Kind::Refused`], as it
52/// does through [`connect`].
53pub 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
62/// Run a caller's own connect future, giving up after `timeout` with
63/// [`Kind::ConnectTimeout`].
64///
65/// The future's own output comes back as it is, so its errors stay the
66/// caller's. Install the rustls provider before a TLS connect of its own
67/// ([`ensure_crypto_provider`]).
68pub 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
74/// Whether reconnecting can help after a connect failed this way.
75///
76/// A refused upgrade is decided by `retriable_status`, the status rule. Pass
77/// `polyoxide-venue`'s, which `Classify::is_retriable` reads too, so the two
78/// can never disagree on a status:
79/// `|status| class_for_handshake_status(status).is_retriable()`.
80///
81/// A bad URL, an upgrade request that could not be built, a refused TLS
82/// configuration, an attack detection and a capacity limit replay the same
83/// failure on every attempt, so they are `false`. Everything else, an I/O
84/// failure, a timeout or a lost connection among them, is `true`.
85///
86/// This decides only whether a reconnect can help, not what a failure is:
87/// that is the venue's classification. What to do after
88/// [`Kind::AlreadyClosed`] (D19) and what a close code means stay with each
89/// venue's recovery, so both are `true` here.
90///
91/// The basis is one socket crate's former `retrying_can_fix` on its error
92/// (`src/error.rs`), with the status rule injected, since this crate cannot
93/// see `polyoxide-venue`. That crate's recovery calls this since Story 4.10,
94/// with its own row ahead of it: it retries a `Capacity` failure, which this
95/// gives up on.
96pub 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    /// A stand-in for `polyoxide-venue`'s status rule, which this crate
116    /// cannot see; venue's table test checks the real one against the
117    /// classification.
118    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    /// A TCP connection to `server`, opened as a venue's own connect would.
208    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        // It has no timeout of its own; a server that never answers the
247        // upgrade holds it until the caller's bound.
248        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            // D19 and a close code's meaning are the venue's to decide.
297            WsError::from(tungstenite::Error::AlreadyClosed),
298            WsError::closed(Some(1008), "policy"),
299        ] {
300            assert!(retrying_can_fix(&err, |_| false), "{err:?}");
301        }
302    }
303}