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

use std::{
    collections::HashMap,
    io,
    net::IpAddr,
    sync::{Arc, Mutex},
};

use ipnet::IpNet;
use remora_outbound::{
    CanonicalUrl, OutboundError, OutboundPolicy, OutboundPolicyBuilder, Resolver, SystemResolver,
};

#[derive(Clone, Default)]
struct ScriptedResolver {
    state: Arc<Mutex<ResolverState>>,
}

#[derive(Default)]
struct ResolverState {
    answers: HashMap<String, Vec<Vec<IpAddr>>>,
    calls: HashMap<String, usize>,
}

impl ScriptedResolver {
    fn with(host: &str, answers: Vec<Vec<IpAddr>>) -> Self {
        let resolver = Self::default();
        resolver
            .state
            .lock()
            .expect("resolver state")
            .answers
            .insert(host.to_owned(), answers);
        resolver
    }

    fn calls(&self, host: &str) -> usize {
        self.state
            .lock()
            .expect("resolver state")
            .calls
            .get(host)
            .copied()
            .unwrap_or_default()
    }
}

impl Resolver for ScriptedResolver {
    async fn resolve(&self, host: &str, _port: u16) -> io::Result<Vec<IpAddr>> {
        let mut state = self.state.lock().expect("resolver state");
        let call = state.calls.entry(host.to_owned()).or_default();
        let index = *call;
        *call += 1;
        let sequence = state.answers.get(host).cloned().unwrap_or_default();
        Ok(sequence
            .get(index)
            .or_else(|| sequence.last())
            .cloned()
            .unwrap_or_default())
    }
}

fn ip(input: &str) -> IpAddr {
    input.parse().expect("valid test address")
}

async fn authorize_literal(policy: &OutboundPolicy, address: &str) -> Result<(), OutboundError> {
    let brackets = if address.contains(':') {
        ("[", "]")
    } else {
        ("", "")
    };
    let url = CanonicalUrl::parse(&format!("https://{}{address}{}/", brackets.0, brackets.1))?;
    policy
        .authorize(&url, &ScriptedResolver::default())
        .await
        .map(|_| ())
}

#[test]
fn canonical_url_rejects_ambiguous_or_sensitive_forms() {
    let canonical = CanonicalUrl::parse("https://EXAMPLE.com.:443/a?q=1").expect("canonical URL");
    assert_eq!(canonical.to_string(), "https://example.com/a?q=1");

    for invalid in [
        "ftp://example.com/object",
        "https://user:[email protected]/object",
        "https://example.com/object#fragment",
        "https://%31%32%37.0.0.1/object",
        "https://2130706433/object",
        "https://0x7f000001/object",
        "https://0177.0.0.1/object",
        "https://127.1/object",
        "https://[2001:0DB8::1]/object",
    ] {
        assert_eq!(CanonicalUrl::parse(invalid), Err(OutboundError::InvalidUrl));
    }
}

#[tokio::test]
async fn complete_answer_must_be_nonempty_and_fully_allowed() {
    let policy = OutboundPolicy::https_only();
    let mixed = ScriptedResolver::with(
        "mixed.test",
        vec![vec![ip("93.184.216.34"), ip("10.0.0.1")]],
    );
    let url = CanonicalUrl::parse("https://mixed.test/object").expect("URL");
    assert_eq!(
        policy.authorize(&url, &mixed).await,
        Err(OutboundError::AddressDenied)
    );

    let empty = ScriptedResolver::with("empty.test", vec![vec![]]);
    let url = CanonicalUrl::parse("https://empty.test/object").expect("URL");
    assert_eq!(
        policy.authorize(&url, &empty).await,
        Err(OutboundError::EmptyResolution)
    );
}

#[tokio::test]
async fn default_policy_denies_special_ipv4_and_ipv6_classes() {
    let policy = OutboundPolicy::https_only();
    for address in [
        "0.0.0.0",
        "10.0.0.1",
        "100.64.0.1",
        "127.0.0.1",
        "169.254.1.1",
        "172.16.0.1",
        "192.168.0.1",
        "192.0.2.1",
        "198.18.0.1",
        "224.0.0.1",
        "240.0.0.1",
        "::",
        "::1",
        "fc00::1",
        "fe80::1",
        "ff00::1",
        "2001:db8::1",
        "::ffff:127.0.0.1",
    ] {
        assert_eq!(
            authorize_literal(&policy, address).await,
            Err(OutboundError::AddressDenied),
            "{address}"
        );
    }
    for address in ["93.184.216.34", "2606:4700:4700::1111"] {
        authorize_literal(&policy, address)
            .await
            .expect("global address");
    }
}

#[tokio::test]
async fn allowlists_never_override_metadata_denial() {
    let policy = OutboundPolicyBuilder::new()
        .allow_non_global("10.0.0.0/8".parse::<IpNet>().expect("network"))
        .allow_non_global("169.254.0.0/16".parse::<IpNet>().expect("network"))
        .allow_non_global("fc00::/7".parse::<IpNet>().expect("network"))
        .build();
    authorize_literal(&policy, "10.1.2.3")
        .await
        .expect("allowlisted private address");
    authorize_literal(&policy, "fd12:3456:789a::1")
        .await
        .expect("allowlisted ULA address");
    for address in [
        "169.254.169.254",
        "169.254.170.2",
        "169.254.0.23",
        "::ffff:169.254.170.2",
        "100.100.100.200",
        "::ffff:100.100.100.200",
        "192.0.0.192",
        "fd00:ec2::254",
        "fd20:ce::254",
        "fd00:db8:deca:deed::fe",
    ] {
        assert_eq!(
            authorize_literal(&policy, address).await,
            Err(OutboundError::AddressDenied)
        );
    }
}

#[tokio::test]
async fn system_resolver_returns_complete_deduplicated_localhost_answers() {
    let answers = SystemResolver
        .resolve("localhost", 80)
        .await
        .expect("localhost resolution");
    assert!(!answers.is_empty());
    assert!(answers.iter().all(IpAddr::is_loopback));
    for (index, address) in answers.iter().enumerate() {
        assert!(!answers[..index].contains(address));
    }
}

#[tokio::test]
async fn scheme_port_policy_is_exact() {
    let resolver = ScriptedResolver::with("example.test", vec![vec![ip("93.184.216.34")]]);
    let policy = OutboundPolicy::https_only();
    for denied in ["http://example.test/", "https://example.test:444/"] {
        assert_eq!(
            policy
                .authorize(&CanonicalUrl::parse(denied).expect("URL"), &resolver)
                .await,
            Err(OutboundError::SchemePortDenied)
        );
    }
}

#[tokio::test]
async fn each_authorization_performs_one_fresh_resolution() {
    let resolver = ScriptedResolver::with(
        "rebind.test",
        vec![vec![ip("93.184.216.34")], vec![ip("10.0.0.1")]],
    );
    let policy = OutboundPolicy::https_only();
    let url = CanonicalUrl::parse("https://rebind.test/").expect("URL");
    policy
        .authorize(&url, &resolver)
        .await
        .expect("first authorization");
    assert_eq!(
        policy.authorize(&url, &resolver).await,
        Err(OutboundError::AddressDenied)
    );
    assert_eq!(resolver.calls("rebind.test"), 2);
}