Enforce outbound policy at each HTTP connection

crates/remora-outbound/src/policy.rs

@@ -0,0 +1,185 @@
+use std::{
+    collections::HashSet,
+    net::{IpAddr, Ipv4Addr, Ipv6Addr},
+};
+
+use ipnet::IpNet;
+
+use crate::{CanonicalUrl, OutboundError, Resolver};
+
+const METADATA_V4: &str = "169.254.169.254";
+const METADATA_V6: &str = "fd00:ec2::254";
+
+#[derive(Clone, Debug)]
+pub struct OutboundPolicy {
+    allowed_scheme_ports: HashSet<(String, u16)>,
+    allowed_non_global: Vec<IpNet>,
+    max_redirects: usize,
+}
+
+#[derive(Clone, Debug)]
+pub struct OutboundPolicyBuilder {
+    policy: OutboundPolicy,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct AllowedEndpoint {
+    address: IpAddr,
+    port: u16,
+}
+
+impl OutboundPolicy {
+    #[must_use]
+    pub fn https_only() -> Self {
+        OutboundPolicyBuilder::new().build()
+    }
+
+    pub async fn authorize<R: Resolver>(
+        &self,
+        url: &CanonicalUrl,
+        resolver: &R,
+    ) -> Result<AllowedEndpoint, OutboundError> {
+        let port = url.port().ok_or(OutboundError::SchemePortDenied)?;
+        if !self
+            .allowed_scheme_ports
+            .iter()
+            .any(|(scheme, allowed_port)| scheme == url.scheme() && *allowed_port == port)
+        {
+            return Err(OutboundError::SchemePortDenied);
+        }
+        let addresses = if let Some(address) = url.literal_address() {
+            vec![address]
+        } else {
+            resolver
+                .resolve(url.host_str(), port)
+                .await
+                .map_err(|_| OutboundError::ResolutionFailed)?
+        };
+        let address = *addresses.first().ok_or(OutboundError::EmptyResolution)?;
+        if addresses
+            .iter()
+            .any(|address| !self.address_allowed(*address))
+        {
+            return Err(OutboundError::AddressDenied);
+        }
+        Ok(AllowedEndpoint { address, port })
+    }
+
+    #[must_use]
+    pub fn max_redirects(&self) -> usize {
+        self.max_redirects
+    }
+
+    fn address_allowed(&self, address: IpAddr) -> bool {
+        let effective = match address {
+            IpAddr::V6(address) => address.to_ipv4_mapped().map_or(address.into(), IpAddr::V4),
+            address => address,
+        };
+        if effective.to_string() == METADATA_V4 || effective.to_string() == METADATA_V6 {
+            return false;
+        }
+        if self
+            .allowed_non_global
+            .iter()
+            .any(|network| network.contains(&effective))
+        {
+            return true;
+        }
+        is_global_unicast(effective)
+    }
+}
+
+impl Default for OutboundPolicy {
+    fn default() -> Self {
+        Self::https_only()
+    }
+}
+
+impl OutboundPolicyBuilder {
+    #[must_use]
+    pub fn new() -> Self {
+        Self {
+            policy: OutboundPolicy {
+                allowed_scheme_ports: HashSet::from([("https".to_owned(), 443)]),
+                allowed_non_global: Vec::new(),
+                max_redirects: 5,
+            },
+        }
+    }
+
+    #[must_use]
+    pub fn allow_scheme_port(mut self, scheme: impl Into<String>, port: u16) -> Self {
+        self.policy
+            .allowed_scheme_ports
+            .insert((scheme.into().to_ascii_lowercase(), port));
+        self
+    }
+
+    #[must_use]
+    pub fn allow_non_global(mut self, network: IpNet) -> Self {
+        self.policy.allowed_non_global.push(network);
+        self
+    }
+
+    #[must_use]
+    pub fn max_redirects(mut self, max_redirects: usize) -> Self {
+        self.policy.max_redirects = max_redirects;
+        self
+    }
+
+    #[must_use]
+    pub fn build(self) -> OutboundPolicy {
+        self.policy
+    }
+}
+
+impl Default for OutboundPolicyBuilder {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl AllowedEndpoint {
+    #[must_use]
+    pub fn address(&self) -> IpAddr {
+        self.address
+    }
+
+    #[must_use]
+    pub fn port(&self) -> u16 {
+        self.port
+    }
+}
+
+fn is_global_unicast(address: IpAddr) -> bool {
+    match address {
+        IpAddr::V4(address) => is_global_ipv4(address),
+        IpAddr::V6(address) => is_global_ipv6(address),
+    }
+}
+
+fn is_global_ipv4(address: Ipv4Addr) -> bool {
+    if address.is_unspecified()
+        || address.is_loopback()
+        || address.is_private()
+        || address.is_link_local()
+        || address.is_documentation()
+        || address.is_broadcast()
+        || address.is_multicast()
+    {
+        return false;
+    }
+    let [first, second, third, _] = address.octets();
+    !((first == 100 && (64..=127).contains(&second))
+        || (first == 192 && second == 0 && third == 0)
+        || (first == 192 && second == 88 && third == 99)
+        || (first == 198 && (second == 18 || second == 19))
+        || first >= 240)
+}
+
+fn is_global_ipv6(address: Ipv6Addr) -> bool {
+    let segments = address.segments();
+    segments[0] & 0xe000 == 0x2000
+        && !((segments[0] == 0x2001 && (segments[1] <= 0x01ff || segments[1] == 0x0db8))
+            || segments[0] & 0xfff0 == 0x3ff0)
+}