apollo-http-client 0.2.0

//! Proxy connector and connection pools for routing requests through an HTTP(S) proxy.

use std::future::Future;
use std::io;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;

use apollo_opentelemetry::metrics::UpDownCounterExt as _;
use base64::Engine as _;
use bytes::Bytes;
use http::header::HeaderValue;
use http::uri::{Authority, Scheme};
use http_body_util::Empty;
use hyper::client::conn::http1;
use hyper_util::rt::TokioIo;
use opentelemetry::KeyValue;
use opentelemetry_semantic_conventions::attribute as semconv;
use rustls::ClientConfig;
use tokio::net::TcpStream;
use tokio_rustls::TlsConnector;
use tower::BoxError;
use tower::Service;

use crate::config::Protocol;
use crate::error::HttpClientError;
use crate::metrics::{ConnectionMetrics, STATE_CUSTOM_CONNECTING};
use crate::protocol::handshake::{FixedOrigin, Http1Handshake, Http2Handshake};
use crate::protocol::{HostSender, Origin};

// ---- ProxySenderConfig -------------------------------------------------------

/// All proxy-related configuration extracted from [`crate::config::HttpClientConfig`].
pub(crate) struct ProxySenderConfig {
    pub(crate) proxy_host: String,
    pub(crate) proxy_port: u16,
    pub(crate) proxy_auth: Option<HeaderValue>,
    pub(crate) tls_config: Arc<ClientConfig>,
    pub(crate) protocol: Protocol,
    pub(crate) connect_timeout: Duration,
    pub(crate) h1_max_idle: usize,
    pub(crate) h1_idle_timeout: Duration,
    pub(crate) h2_idle_timeout: Duration,
    pub(crate) keep_alive_interval: Duration,
    pub(crate) keep_alive_timeout: Duration,
    pub(crate) keep_alive_while_idle: bool,
    pub(crate) conn_metrics: ConnectionMetrics,
}

// ---- ProxyStream -------------------------------------------------------------

/// IO stream returned by [`ProxyConnector`]: either a plain TCP connection to
/// the proxy (for HTTP targets) or a TLS connection through a CONNECT tunnel
/// (for HTTPS targets).
enum ProxyStream {
    Plain(TokioIo<TcpStream>),
    // Boxed to reduce enum size: TlsStream is large (~1 KiB) while TcpStream is small.
    Tunnel(Box<TokioIo<tokio_rustls::client::TlsStream<TokioIo<hyper::upgrade::Upgraded>>>>),
}

impl ProxyStream {
    /// Returns the ALPN protocol negotiated during the TLS handshake, or `None`
    /// for plain connections. Used to route ALPN HTTPS connections to H1 or H2.
    fn alpn_protocol(&self) -> Option<&[u8]> {
        match self {
            ProxyStream::Plain(_) => None,
            ProxyStream::Tunnel(tls_io) => {
                let (_, conn) = tls_io.inner().get_ref();
                conn.alpn_protocol()
            }
        }
    }
}

impl hyper::rt::Read for ProxyStream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: hyper::rt::ReadBufCursor<'_>,
    ) -> Poll<io::Result<()>> {
        match &mut *self {
            ProxyStream::Plain(s) => Pin::new(s).poll_read(cx, buf),
            ProxyStream::Tunnel(s) => Pin::new(&mut **s).poll_read(cx, buf),
        }
    }
}

impl hyper::rt::Write for ProxyStream {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        match &mut *self {
            ProxyStream::Plain(s) => Pin::new(s).poll_write(cx, buf),
            ProxyStream::Tunnel(s) => Pin::new(&mut **s).poll_write(cx, buf),
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        match &mut *self {
            ProxyStream::Plain(s) => Pin::new(s).poll_flush(cx),
            ProxyStream::Tunnel(s) => Pin::new(&mut **s).poll_flush(cx),
        }
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        match &mut *self {
            ProxyStream::Plain(s) => Pin::new(s).poll_shutdown(cx),
            ProxyStream::Tunnel(s) => Pin::new(&mut **s).poll_shutdown(cx),
        }
    }
}

// ---- ProxyConnector ----------------------------------------------------------

/// Connects to the proxy and returns a [`ProxyStream`].
///
/// For HTTP targets, opens a plain TCP connection to the proxy.
/// For HTTPS targets, opens a CONNECT tunnel and performs the TLS handshake
/// inside it, so higher layers see a direct TLS stream to the target.
#[derive(Clone)]
struct ProxyConnector {
    proxy_host: String,
    proxy_port: u16,
    proxy_auth: Option<HeaderValue>,
    tls_config: Arc<ClientConfig>,
    connect_timeout: Duration,
    metrics: ConnectionMetrics,
}

impl ProxyConnector {
    fn from_config(config: &ProxySenderConfig) -> Self {
        Self {
            proxy_host: config.proxy_host.clone(),
            proxy_port: config.proxy_port,
            proxy_auth: config.proxy_auth.clone(),
            tls_config: config.tls_config.clone(),
            connect_timeout: config.connect_timeout,
            metrics: config.conn_metrics.clone(),
        }
    }
}

impl Service<Origin> for ProxyConnector {
    type Response = ProxyStream;
    type Error = BoxError;
    type Future = Pin<Box<dyn Future<Output = Result<ProxyStream, BoxError>> + Send>>;

    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, origin: Origin) -> Self::Future {
        let proxy_host = self.proxy_host.clone();
        let proxy_port = self.proxy_port;
        let proxy_auth = self.proxy_auth.clone();
        let tls_config = self.tls_config.clone();
        let connect_timeout = self.connect_timeout;
        let metrics = self.metrics.clone();

        Box::pin(async move {
            let addr = format!("{proxy_host}:{proxy_port}");

            let mut attrs = metrics.connection_attrs(None, &origin);
            attrs.push(KeyValue::new(
                semconv::HTTP_CONNECTION_STATE,
                STATE_CUSTOM_CONNECTING,
            ));
            let _connecting = metrics.counter().track(attrs);

            // The full proxy setup — TCP connect, CONNECT handshake, and tunnel TLS
            // handshake — must complete within `connect_timeout` so a stalled or
            // misbehaving proxy cannot delay request dispatch indefinitely. This
            // mirrors the non-proxy path, which applies the same timeout to TCP+TLS.
            let setup = async {
                let tcp = TcpStream::connect(&addr).await.map_err(|source| {
                    HttpClientError::Transport {
                        source: source.into(),
                    }
                })?;
                if origin.scheme == Scheme::HTTPS {
                    let tls =
                        connect_tunnel(tcp, &origin.authority, proxy_auth.as_ref(), tls_config)
                            .await?;
                    Ok::<_, BoxError>(ProxyStream::Tunnel(Box::new(TokioIo::new(tls))))
                } else {
                    Ok(ProxyStream::Plain(TokioIo::new(tcp)))
                }
            };

            match tokio::time::timeout(connect_timeout, setup).await {
                Ok(result) => result,
                Err(_) => Err(HttpClientError::ConnectionTimeout.into()),
            }
        })
    }
}

// ---- CONNECT tunnel ----------------------------------------------------------

/// Opens a CONNECT tunnel through the proxy to `authority` and performs the TLS
/// handshake inside it, returning a TLS stream over the upgraded connection.
async fn connect_tunnel(
    tcp: TcpStream,
    authority: &Authority,
    proxy_auth: Option<&HeaderValue>,
    tls_config: Arc<ClientConfig>,
) -> Result<tokio_rustls::client::TlsStream<TokioIo<hyper::upgrade::Upgraded>>, BoxError> {
    let io = TokioIo::new(tcp);
    let (mut sender, conn): (http1::SendRequest<Empty<Bytes>>, _) = http1::handshake(io).await?;

    // Spawn the connection driver; `with_upgrades` is required so hyper can hand
    // off the raw TCP stream once the CONNECT response is received. We retain the
    // `JoinHandle` so the upgrade wait below can detect a driver that died (IO
    // error, panic) before it had a chance to signal upgrade — otherwise
    // `hyper::upgrade::on` would block forever waiting for a signal that never
    // arrives, until the outer connect timeout fires.
    let driver = tokio::spawn(conn.with_upgrades());

    // Build the CONNECT request using authority-form (host:port).
    let connect_target = format!(
        "{}:{}",
        authority.host(),
        authority.port_u16().unwrap_or(443)
    );
    let mut req = http::Request::builder()
        .method(http::Method::CONNECT)
        .uri(
            http::Uri::builder()
                .authority(connect_target.as_str())
                .build()?,
        )
        .version(http::Version::HTTP_11)
        .header(http::header::HOST, authority.as_str())
        .body(Empty::<Bytes>::new())?;

    if let Some(auth) = proxy_auth {
        req.headers_mut()
            .insert(http::header::PROXY_AUTHORIZATION, auth.clone());
    }

    let resp = sender.send_request(req).await?;
    let status = resp.status();
    if status != http::StatusCode::OK {
        return Err(HttpClientError::ProxyTunnel { status }.into());
    }

    // On the success path, `with_upgrades` resolves `Ok(())` shortly after handing
    // off the socket to `hyper::upgrade::on`, so both futures become ready around
    // the same time. To avoid racing them against each other, fold the driver
    // result into a future that resolves *only on failure* — a clean exit becomes
    // `pending`, so we always wait for the upgrade signal in the normal path.
    let driver_outcome = driver_failure(driver);
    tokio::pin!(driver_outcome);
    let upgraded = tokio::select! {
        upgraded = hyper::upgrade::on(resp) => upgraded?,
        err = &mut driver_outcome => return Err(err),
    };

    // Perform TLS inside the tunnel. The TLS config already has the correct ALPN
    // protocols set (e.g. ["h2", "http/1.1"] for ALPN mode).
    let server_name = rustls::pki_types::ServerName::try_from(authority.host().to_owned())
        .map_err(|e| format!("invalid proxy target server name: {e}"))?;
    let tls = TlsConnector::from(tls_config)
        .connect(server_name, TokioIo::new(upgraded))
        .await?;

    Ok(tls)
}

/// Adapts a spawned connection-driver `JoinHandle` into a future that resolves
/// **only on failure** — a clean `Ok(Ok(()))` exit becomes `pending`. This lets
/// callers race the driver against `hyper::upgrade::on` without the success path
/// (where both complete around the same time) ever picking the driver branch.
async fn driver_failure<E>(driver: tokio::task::JoinHandle<Result<(), E>>) -> BoxError
where
    E: Into<BoxError>,
{
    match driver.await {
        Ok(Ok(())) => std::future::pending::<BoxError>().await,
        Ok(Err(err)) => err.into(),
        Err(join_err) => join_err.into(),
    }
}

// ---- Auth helper -------------------------------------------------------------

pub(crate) fn proxy_authorization(url: &url::Url) -> Option<HeaderValue> {
    let username = url.username();
    if username.is_empty() {
        return None;
    }
    // `url::Url` returns the userinfo in percent-encoded form. RFC 7617 Basic auth
    // requires the decoded `userid:password` pair, so percent-decode both halves
    // before base64-encoding.
    let username = percent_encoding::percent_decode_str(username).decode_utf8_lossy();
    let password =
        percent_encoding::percent_decode_str(url.password().unwrap_or("")).decode_utf8_lossy();
    let encoded =
        base64::engine::general_purpose::STANDARD.encode(format!("{username}:{password}"));
    // base64 output is always ASCII, so HeaderValue construction cannot fail.
    Some(
        HeaderValue::from_str(&format!("Basic {encoded}"))
            .expect("base64 output is always a valid HeaderValue"),
    )
}

// ---- Pool builders -----------------------------------------------------------

/// Returns a [`HostSender`] that routes all requests through the configured proxy.
///
/// - ALPN HTTPS: negotiates H1 or H2 inside the CONNECT tunnel.
/// - HTTP/2 HTTPS: uses a single H2 connection through the CONNECT tunnel.
/// - Everything else (HTTP targets, or HTTP/1.1): uses an H1 connection pool
///   with absolute-form request forwarding.
pub(crate) fn new_sender(config: &Arc<ProxySenderConfig>, origin: Origin) -> HostSender {
    match (config.protocol, origin.scheme.as_str()) {
        (Protocol::Alpn, "https") => new_alpn_https_sender(config, origin),
        (Protocol::Http2, "https") => new_h2_sender(config, origin),
        _ => new_h1_sender(config, origin),
    }
}

fn new_h1_sender(config: &Arc<ProxySenderConfig>, origin: Origin) -> HostSender {
    let attrs = config
        .conn_metrics
        .connection_attrs(Some(http::Version::HTTP_11), &origin);
    let connector = Http1Handshake::new(
        ProxyConnector::from_config(config),
        config.conn_metrics.clone(),
        attrs,
    );
    // Proxy-Authorization is added per-request only for plain HTTP forwarding.
    // For HTTPS targets the connection goes through a CONNECT tunnel; the auth
    // header is sent once during tunnel setup and must not appear in the
    // application request (which is encrypted and goes directly to the backend,
    // never seen by the proxy).
    let proxy_auth = if origin.scheme == Scheme::HTTPS {
        None
    } else {
        config.proxy_auth.clone()
    };
    crate::protocol::h1::new_sender(
        connector,
        config.h1_max_idle,
        config.h1_idle_timeout,
        origin,
        move |req| {
            if let Some(auth) = &proxy_auth {
                req.headers_mut()
                    .entry(http::header::PROXY_AUTHORIZATION)
                    .or_insert_with(|| auth.clone());
            }
        },
    )
}

fn new_h2_sender(config: &Arc<ProxySenderConfig>, origin: Origin) -> HostSender {
    let attrs = config
        .conn_metrics
        .connection_attrs(Some(http::Version::HTTP_2), &origin);
    let handshake = Http2Handshake {
        inner: FixedOrigin {
            inner: ProxyConnector::from_config(config),
            origin,
        },
        metrics: config.conn_metrics.clone(),
        attrs,
        keep_alive_interval: config.keep_alive_interval,
        keep_alive_timeout: config.keep_alive_timeout,
        keep_alive_while_idle: config.keep_alive_while_idle,
    };
    crate::protocol::h2::new_sender(handshake, config.h2_idle_timeout, (), |_| {})
}

fn new_alpn_https_sender(config: &Arc<ProxySenderConfig>, origin: Origin) -> HostSender {
    crate::protocol::alpn::new_sender_with_connector(
        ProxyConnector::from_config(config),
        |io: &ProxyStream| io.alpn_protocol() == Some(b"h2"),
        config.conn_metrics.clone(),
        crate::protocol::alpn::AlpnPoolConfig {
            h1_max_idle: config.h1_max_idle,
            h1_idle_timeout: config.h1_idle_timeout,
            h2_keep_alive_interval: config.keep_alive_interval,
            h2_keep_alive_timeout: config.keep_alive_timeout,
            h2_keep_alive_while_idle: config.keep_alive_while_idle,
            h2_idle_timeout: config.h2_idle_timeout,
        },
        origin,
    )
}

// ---- Unit tests --------------------------------------------------------------

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use tower::BoxError;

    use super::{driver_failure, proxy_authorization};

    #[test]
    fn proxy_authorization_returns_none_when_no_credentials() {
        let url: url::Url = "http://proxy.example.com:3128".parse().unwrap();
        assert!(proxy_authorization(&url).is_none());
    }

    #[test]
    fn proxy_authorization_encodes_credentials_as_basic_auth() {
        let url: url::Url = "http://user:pass@proxy.example.com:3128".parse().unwrap();
        let header = proxy_authorization(&url).expect("credentials present");
        // base64("user:pass") = "dXNlcjpwYXNz"
        assert_eq!(header.to_str().unwrap(), "Basic dXNlcjpwYXNz");
    }

    #[test]
    fn proxy_authorization_handles_empty_password() {
        let url: url::Url = "http://user:@proxy.example.com:3128".parse().unwrap();
        let header = proxy_authorization(&url).expect("credentials present");
        // base64("user:") = "dXNlcjo="
        assert_eq!(header.to_str().unwrap(), "Basic dXNlcjo=");
    }

    #[test]
    fn proxy_authorization_percent_decodes_credentials() {
        // `pass%3Aword` decodes to `pass:word`; per RFC 7617 the colon is part of the
        // password (the userid is everything before the *first* colon in the decoded
        // pair). The encoded form must not be base64'd verbatim.
        let url: url::Url = "http://user:pass%3Aword@proxy.example.com:3128"
            .parse()
            .unwrap();
        let header = proxy_authorization(&url).expect("credentials present");
        // base64("user:pass:word") = "dXNlcjpwYXNzOndvcmQ="
        assert_eq!(header.to_str().unwrap(), "Basic dXNlcjpwYXNzOndvcmQ=");
    }

    /// A clean `Ok(Ok(()))` driver exit must never resolve `driver_failure` — the
    /// success path of `connect_tunnel` relies on the upgrade branch winning, so
    /// the failure future has to stay pending forever for an Ok-exiting driver.
    #[tokio::test]
    async fn driver_failure_stays_pending_when_driver_exits_ok() {
        let driver: tokio::task::JoinHandle<Result<(), BoxError>> = tokio::spawn(async { Ok(()) });
        // Give the spawned task a chance to actually complete before we poll the
        // adapter, so we exercise the "JoinHandle already resolved Ok" branch.
        tokio::task::yield_now().await;

        let timed_out = tokio::time::timeout(Duration::from_millis(50), driver_failure(driver))
            .await
            .is_err();
        assert!(
            timed_out,
            "driver_failure should remain pending when driver exits Ok"
        );
    }

    /// A driver that exits `Err(_)` resolves `driver_failure` with that error,
    /// preserving it as a `BoxError` so `connect_tunnel` can propagate it through
    /// the existing `into_http_client_error` mapper.
    #[tokio::test]
    async fn driver_failure_propagates_driver_error() {
        let driver: tokio::task::JoinHandle<Result<(), BoxError>> =
            tokio::spawn(async { Err("simulated driver IO error".into()) });

        let err = driver_failure(driver).await;
        assert_eq!(err.to_string(), "simulated driver IO error");
    }

    /// A driver task that is cancelled (or panics) resolves `driver_failure`
    /// with the `JoinError`, so the caller learns about the death rather than
    /// hanging on `upgrade::on` waiting for a signal that will never arrive.
    #[tokio::test]
    async fn driver_failure_propagates_join_error_on_cancel() {
        let driver: tokio::task::JoinHandle<Result<(), BoxError>> =
            tokio::spawn(async { std::future::pending::<Result<(), BoxError>>().await });
        driver.abort();

        // We only need to confirm the adapter resolves at all — the exact
        // message is `JoinError`'s and not something this crate owns.
        let err = driver_failure(driver).await;
        assert!(
            !err.to_string().is_empty(),
            "JoinError should produce a non-empty BoxError"
        );
    }
}