Enforce outbound policy at each HTTP connection

crates/remora-outbound/src/transport.rs

@@ -0,0 +1,182 @@
+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);
+    }
+}