polyoxide-ws 0.48.0

//! Opening a connection within a timeout (W8), and whether retrying a failed
//! one can help (W7).

use std::{future::Future, time::Duration};

use polyoxide_rt::{default_runtime, Runtime};
#[cfg(not(target_arch = "wasm32"))]
use tokio::net::TcpStream;
#[cfg(not(target_arch = "wasm32"))]
use tokio_tungstenite::client_async_tls;

use crate::{
    default_socket_transport, tungstenite::client::IntoClientRequest, BoxSocket, Kind, WsError,
};

/// How long one connection attempt may take unless a venue says otherwise:
/// the TCP connect, the TLS handshake and the WebSocket upgrade together.
pub const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);

/// Open a connection through the target's default [`SocketTransport`](crate::SocketTransport)
/// and runtime, giving up after `timeout`.
///
/// `request` is a URL, or a handshake request with its own headers, built
/// afresh for each attempt so any authentication in it is recomputed. A
/// connection not open within `timeout` fails as [`Kind::ConnectTimeout`];
/// every other failure is the transport's, such as [`Kind::Refused`] for an
/// upgrade the server answered with an HTTP status. A caller that injects its
/// own transport or runtime opens its connections through them instead.
pub async fn connect<R>(request: R, timeout: Duration) -> Result<BoxSocket, WsError>
where
    R: IntoClientRequest + Unpin,
{
    let request = request.into_client_request()?;
    default_socket_transport()
        .connect(request, timeout, default_runtime())
        .await
}

/// Upgrade a TCP connection the caller opened itself: the TLS handshake when
/// `request` is a `wss://` URL, then the WebSocket upgrade.
///
/// For a transport whose connect is its own: one that resolves the host and
/// dials each address itself, and hands the first TCP connection that opens
/// to this. The rustls provider is installed first
/// ([`ensure_crypto_provider`](crate::ensure_crypto_provider)). It has no
/// timeout of its own, so bound it with [`with_timeout`]. A refused upgrade
/// fails as [`Kind::Refused`], as it does through [`connect`]. Native only.
#[cfg(not(target_arch = "wasm32"))]
pub async fn connect_over<R>(request: R, tcp: TcpStream) -> Result<BoxSocket, WsError>
where
    R: IntoClientRequest + Unpin,
{
    crate::ensure_crypto_provider();
    let (socket, _response) = client_async_tls(request, tcp).await?;
    Ok(Box::new(socket))
}

/// Run a caller's own connect future, giving up after `timeout` with
/// [`Kind::ConnectTimeout`].
///
/// The future's own output comes back as it is, so its errors stay the
/// caller's. It runs on the target's default runtime; a transport bounds its
/// own steps on the runtime it is given instead. Install the rustls provider
/// before a TLS connect of its own
/// ([`ensure_crypto_provider`](crate::ensure_crypto_provider)).
pub async fn with_timeout<F: Future>(timeout: Duration, connect: F) -> Result<F::Output, WsError> {
    let runtime = default_runtime();
    bounded(runtime.as_ref(), timeout, connect).await
}

/// Run `future` within `timeout` on `runtime`, giving up with
/// [`Kind::ConnectTimeout`]. A transport bounds each of its own steps with the
/// runtime it is given.
pub(crate) async fn bounded<F: Future>(
    runtime: &dyn Runtime,
    timeout: Duration,
    future: F,
) -> Result<F::Output, WsError> {
    polyoxide_rt::timeout(runtime, timeout, future)
        .await
        .ok_or_else(|| WsError::connect_timeout(timeout))
}

/// Whether reconnecting can help after a connect failed this way.
///
/// A refused upgrade is decided by `retriable_status`, the status rule. Pass
/// `polyoxide-venue`'s, which `Classify::is_retriable` reads too, so the two
/// can never disagree on a status:
/// `|status| class_for_handshake_status(status).is_retriable()`.
///
/// A bad URL, an upgrade request that could not be built, a refused TLS
/// configuration, an attack detection, a capacity limit and an operation the
/// transport cannot perform replay the same failure on every attempt, so they
/// are `false`. Everything else, an I/O failure, a timeout or a lost
/// connection among them, is `true`.
///
/// This decides only whether a reconnect can help, not what a failure is:
/// that is the venue's classification. What to do after
/// [`Kind::AlreadyClosed`] (D19) and what a close code means stay with each
/// venue's recovery, so both are `true` here.
///
/// The basis is one socket crate's former `retrying_can_fix` on its error
/// (`src/error.rs`), with the status rule injected, since this crate cannot
/// see `polyoxide-venue`. That crate's recovery calls this since Story 4.10,
/// with its own row ahead of it: it retries a `Capacity` failure, which this
/// gives up on.
pub fn retrying_can_fix(err: &WsError, retriable_status: impl Fn(u16) -> bool) -> bool {
    match err.kind() {
        Kind::Refused { status, .. } => retriable_status(status),
        Kind::Url
        | Kind::HttpFormat
        | Kind::Tls
        | Kind::AttackAttempt
        | Kind::Capacity
        | Kind::Unsupported => false,
        _ => true,
    }
}

#[cfg(test)]
mod tests {
    use tokio_tungstenite::tungstenite::{
        self,
        error::{CapacityError, ProtocolError, TlsError, UrlError},
        http,
    };

    use super::*;
    use crate::test_server::{Script, ScriptedServer};

    /// A stand-in for `polyoxide-venue`'s status rule, which this crate
    /// cannot see; venue's table test checks the real one against the
    /// classification.
    fn rule(status: u16) -> bool {
        matches!(status, 408 | 425 | 429 | 500..=599)
    }

    fn refused(status: u16) -> WsError {
        WsError::from(tungstenite::Error::Http(
            http::Response::builder().status(status).body(None).unwrap(),
        ))
    }

    #[tokio::test(start_paused = true)]
    async fn a_hung_handshake_times_out_at_the_default() {
        let server = ScriptedServer::start(vec![Script::stall()]).await;
        let err = connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
            .await
            .err()
            .expect("a server that never upgrades must not connect");
        assert_eq!(
            err.kind(),
            Kind::ConnectTimeout {
                after: Duration::from_secs(10)
            }
        );
    }

    #[tokio::test]
    async fn a_refused_upgrade_keeps_its_status_and_retry_after() {
        let server = ScriptedServer::start(vec![
            Script::refuse_upgrade(429).with_header("Retry-After", "7")
        ])
        .await;
        let err = connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
            .await
            .err()
            .expect("a refused upgrade must not connect");
        assert_eq!(
            err.kind(),
            Kind::Refused {
                status: 429,
                retry_after: Some("7")
            }
        );
    }

    #[tokio::test]
    async fn a_503_upgrade_is_worth_retrying_and_a_404_is_not() {
        let server = ScriptedServer::start(vec![
            Script::refuse_upgrade(503),
            Script::refuse_upgrade(404),
        ])
        .await;
        let unwell = connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
            .await
            .err()
            .expect("503");
        assert_eq!(
            unwell.kind(),
            Kind::Refused {
                status: 503,
                retry_after: None
            }
        );
        assert!(retrying_can_fix(&unwell, rule));

        let missing = connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
            .await
            .err()
            .expect("404");
        assert!(!retrying_can_fix(&missing, rule));
    }

    #[tokio::test]
    async fn a_server_that_upgrades_connects() {
        let server = ScriptedServer::start(vec![Script::hold()]).await;
        connect(server.url.as_str(), DEFAULT_CONNECT_TIMEOUT)
            .await
            .expect("connect");
    }

    #[tokio::test]
    async fn with_timeout_bounds_a_callers_own_future_and_keeps_its_output() {
        let after = Duration::from_millis(20);
        let err = with_timeout(after, std::future::pending::<()>())
            .await
            .expect_err("a future that never finishes must time out");
        assert_eq!(err.kind(), Kind::ConnectTimeout { after });

        let output: Result<u8, &str> = with_timeout(after, async { Err("its own error") })
            .await
            .expect("a finished future is not a timeout");
        assert_eq!(output, Err("its own error"));
    }

    /// A TCP connection to `server`, opened as a venue's own connect would.
    async fn dial(server: &ScriptedServer) -> TcpStream {
        let addr = server.url.trim_start_matches("ws://");
        TcpStream::connect(addr).await.expect("dial")
    }

    #[tokio::test]
    async fn connect_over_upgrades_a_connection_the_caller_opened() {
        let server = ScriptedServer::start(vec![Script::hold()]).await;
        let tcp = dial(&server).await;
        connect_over(server.url.as_str(), tcp)
            .await
            .expect("upgrade");
        server
            .wait_for("the handshake", |s| s.handshake_count() == 1)
            .await;
    }

    #[tokio::test]
    async fn connect_over_keeps_a_refusals_status_and_retry_after() {
        let server = ScriptedServer::start(vec![
            Script::refuse_upgrade(429).with_header("Retry-After", "7")
        ])
        .await;
        let tcp = dial(&server).await;
        let err = connect_over(server.url.as_str(), tcp)
            .await
            .err()
            .expect("a refused upgrade must not connect");
        assert_eq!(
            err.kind(),
            Kind::Refused {
                status: 429,
                retry_after: Some("7")
            }
        );
    }

    #[tokio::test]
    async fn connect_over_is_bounded_by_with_timeout() {
        // It has no timeout of its own; a server that never answers the
        // upgrade holds it until the caller's bound.
        let server = ScriptedServer::start(vec![Script::stall()]).await;
        let tcp = dial(&server).await;
        let after = Duration::from_millis(100);
        let err = with_timeout(after, connect_over(server.url.as_str(), tcp))
            .await
            .err()
            .expect("a stalled upgrade must time out");
        assert_eq!(err.kind(), Kind::ConnectTimeout { after });
    }

    #[test]
    fn the_status_rule_decides_a_refused_upgrade() {
        for status in [301, 401, 403, 404, 418, 451, 600] {
            assert!(!retrying_can_fix(&refused(status), rule), "{status}");
        }
        for status in [408, 425, 429, 500, 502, 503] {
            assert!(retrying_can_fix(&refused(status), rule), "{status}");
        }
    }

    #[test]
    fn a_failure_that_replays_on_every_attempt_is_not_retried() {
        let bad_header = http::header::HeaderName::from_bytes(b"in valid").unwrap_err();
        for err in [
            tungstenite::Error::Url(UrlError::NoHostName),
            tungstenite::Error::HttpFormat(bad_header.into()),
            tungstenite::Error::Tls(TlsError::InvalidDnsName),
            tungstenite::Error::AttackAttempt,
            tungstenite::Error::Capacity(CapacityError::TooManyHeaders),
            tungstenite::Error::Io(std::io::ErrorKind::Unsupported.into()),
        ] {
            let err = WsError::from(err);
            assert!(!retrying_can_fix(&err, |_| true), "{err:?}");
        }
    }

    #[test]
    fn a_transport_failure_or_timeout_is_retried() {
        let after = Duration::from_secs(10);
        for err in [
            WsError::from(tungstenite::Error::Io(
                std::io::ErrorKind::ConnectionReset.into(),
            )),
            WsError::from(tungstenite::Error::Protocol(
                ProtocolError::ResetWithoutClosingHandshake,
            )),
            WsError::from(tungstenite::Error::ConnectionClosed),
            WsError::from(tungstenite::Error::Utf8),
            WsError::connect_timeout(after),
            WsError::no_answer(after),
            // D19 and a close code's meaning are the venue's to decide.
            WsError::from(tungstenite::Error::AlreadyClosed),
            WsError::closed(Some(1008), "policy"),
        ] {
            assert!(retrying_can_fix(&err, |_| false), "{err:?}");
        }
    }
}