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

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)
}