remora/crates/remora-outbound/src/transport.rs

use std::{fmt, net::SocketAddr};

use http::{HeaderMap, HeaderName, Method, StatusCode};
use reqwest::{Certificate, redirect::Policy};

use crate::{CanonicalUrl, OutboundError, OutboundPolicy, Resolver};

const REDIRECTS: [StatusCode; 5] = [
    StatusCode::MOVED_PERMANENTLY,
    StatusCode::FOUND,
    StatusCode::SEE_OTHER,
    StatusCode::TEMPORARY_REDIRECT,
    StatusCode::PERMANENT_REDIRECT,
];

#[derive(Clone, Debug)]
pub struct OutboundClient<R> {
    resolver: R,
    policy: OutboundPolicy,
    roots: Vec<Certificate>,
}

#[derive(Clone, Debug)]
pub struct OutboundClientBuilder<R> {
    client: OutboundClient<R>,
}

pub struct OutboundRequest {
    method: Method,
    url: CanonicalUrl,
    headers: HeaderMap,
    credential_headers: HeaderMap,
}

pub struct OutboundResponse {
    final_url: CanonicalUrl,
    response: reqwest::Response,
}

impl fmt::Debug for OutboundResponse {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("OutboundResponse")
            .field("status", &self.response.status())
            .finish_non_exhaustive()
    }
}

impl<R: Resolver> OutboundClient<R> {
    #[must_use]
    pub fn builder(resolver: R, policy: OutboundPolicy) -> OutboundClientBuilder<R> {
        OutboundClientBuilder {
            client: Self {
                resolver,
                policy,
                roots: Vec::new(),
            },
        }
    }

    pub async fn execute(
        &self,
        mut request: OutboundRequest,
    ) -> Result<OutboundResponse, OutboundError> {
        let mut redirects = 0;
        loop {
            let endpoint = self.policy.authorize(&request.url, &self.resolver).await?;
            let client = self.client_for(&request.url, &endpoint)?;
            let mut headers = request.headers.clone();
            headers.extend(request.credential_headers.clone());
            let response = client
                .request(request.method.clone(), request.url.as_url().clone())
                .headers(headers)
                .send()
                .await
                .map_err(|_| OutboundError::Transport)?;
            if !REDIRECTS.contains(&response.status()) {
                return Ok(OutboundResponse {
                    final_url: request.url,
                    response,
                });
            }
            if redirects == self.policy.max_redirects() {
                return Err(OutboundError::RedirectLimit);
            }
            let location = response
                .headers()
                .get(http::header::LOCATION)
                .ok_or(OutboundError::InvalidRedirect)?
                .to_str()
                .map_err(|_| OutboundError::InvalidRedirect)?;
            let next = request.url.join(location)?;
            if request.url.scheme() == "https" && next.scheme() == "http" {
                return Err(OutboundError::HttpsDowngrade);
            }
            if request.url.origin()? != next.origin()? {
                request.credential_headers.clear();
            }
            request.url = next;
            redirects += 1;
        }
    }

    fn client_for(
        &self,
        url: &CanonicalUrl,
        endpoint: &crate::AllowedEndpoint,
    ) -> Result<reqwest::Client, OutboundError> {
        let mut builder = reqwest::Client::builder()
            .redirect(Policy::none())
            .no_proxy();
        if url.literal_address().is_none() {
            builder = builder.resolve(
                url.host_str(),
                SocketAddr::new(endpoint.address(), endpoint.port()),
            );
        }
        for root in &self.roots {
            builder = builder.add_root_certificate(root.clone());
        }
        builder.build().map_err(|_| OutboundError::Transport)
    }
}

impl<R> OutboundClientBuilder<R> {
    #[must_use]
    pub fn add_root_certificate(mut self, certificate: Certificate) -> Self {
        self.client.roots.push(certificate);
        self
    }

    #[must_use]
    pub fn build(self) -> OutboundClient<R> {
        self.client
    }
}

impl OutboundRequest {
    pub fn get(url: &str) -> Result<Self, OutboundError> {
        Ok(Self {
            method: Method::GET,
            url: CanonicalUrl::parse(url)?,
            headers: HeaderMap::new(),
            credential_headers: HeaderMap::new(),
        })
    }

    pub fn head(url: &str) -> Result<Self, OutboundError> {
        Ok(Self {
            method: Method::HEAD,
            url: CanonicalUrl::parse(url)?,
            headers: HeaderMap::new(),
            credential_headers: HeaderMap::new(),
        })
    }

    pub fn with_public_headers(mut self, headers: HeaderMap) -> Result<Self, OutboundError> {
        if headers
            .keys()
            .any(|name| !is_public_header(name) || self.credential_headers.contains_key(name))
        {
            return Err(OutboundError::InvalidHeaders);
        }
        self.headers = headers;
        Ok(self)
    }

    pub fn with_credential_headers(mut self, headers: HeaderMap) -> Result<Self, OutboundError> {
        if headers
            .keys()
            .any(|name| is_connection_controlled_header(name) || self.headers.contains_key(name))
        {
            return Err(OutboundError::InvalidHeaders);
        }
        self.credential_headers = headers;
        Ok(self)
    }
}

impl fmt::Debug for OutboundRequest {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("OutboundRequest")
            .field("method", &self.method)
            .field("header_count", &self.headers.len())
            .field("credential_header_count", &self.credential_headers.len())
            .finish_non_exhaustive()
    }
}

impl OutboundResponse {
    #[must_use]
    pub fn final_url(&self) -> &CanonicalUrl {
        &self.final_url
    }

    #[must_use]
    pub fn response(&self) -> &reqwest::Response {
        &self.response
    }

    #[must_use]
    pub fn into_response(self) -> reqwest::Response {
        self.response
    }
}

fn is_public_header(name: &HeaderName) -> bool {
    matches!(
        name.as_str(),
        "accept"
            | "accept-charset"
            | "accept-encoding"
            | "accept-language"
            | "cache-control"
            | "if-match"
            | "if-modified-since"
            | "if-none-match"
            | "if-range"
            | "if-unmodified-since"
            | "pragma"
            | "range"
            | "user-agent"
    )
}

fn is_connection_controlled_header(name: &HeaderName) -> bool {
    matches!(
        name.as_str(),
        "connection"
            | "content-length"
            | "host"
            | "keep-alive"
            | "proxy-authenticate"
            | "proxy-authorization"
            | "te"
            | "trailer"
            | "transfer-encoding"
            | "upgrade"
    )
}