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

use std::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>,
}

#[derive(Clone, Debug)]
pub struct OutboundRequest {
    method: Method,
    url: CanonicalUrl,
    headers: HeaderMap,
}

#[derive(Debug)]
pub struct OutboundResponse {
    final_url: CanonicalUrl,
    response: reqwest::Response,
}

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 response = client
                .request(request.method.clone(), request.url.as_url().clone())
                .headers(request.headers.clone())
                .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()? {
                strip_cross_origin_credentials(&mut request.headers);
            }
            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(),
        })
    }

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

    pub fn with_headers(mut self, headers: HeaderMap) -> Result<Self, OutboundError> {
        if headers.contains_key(http::header::HOST)
            || headers.contains_key(http::header::CONNECTION)
            || headers.contains_key(http::header::PROXY_AUTHORIZATION)
        {
            return Err(OutboundError::InvalidHeaders);
        }
        self.headers = headers;
        Ok(self)
    }
}

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 strip_cross_origin_credentials(headers: &mut HeaderMap) {
    for name in [
        http::header::AUTHORIZATION,
        http::header::COOKIE,
        HeaderName::from_static("cookie2"),
    ] {
        headers.remove(name);
    }
}