supermachine 0.7.70

// Per-snapshot egress policy enforcement for TSI outbound CONNECT.
// A single process-global string set via
// CLI (--egress-policy) or per-dispatch (pool-worker RESTORE
// command). `check(target)` returns Ok if allowed, Err otherwise;
// the muxer's TSI_CONNECT path consults this before opening TCP.
//
// Policies:
//   "allow_all"            no restrictions (default)
//   "deny_private"         block RFC 1918 / link-local / loopback /
//                          multicast / current network / TEST-NET /
//                          CGNAT (100.64.0.0/10) / 240.0.0.0/4
//                          For cloud deploys: keeps the metadata
//                          endpoint (169.254.169.254) and internal
//                          DBs unreachable from customer code.
//   "denylist:<cidrs>"     deny_private + extra comma-separated CIDRs
//   "allowlist:<cidrs>"    only allow IPs in these CIDRs

#![allow(dead_code)]

use std::net::IpAddr;
use std::sync::Mutex;

static POLICY: Mutex<String> = Mutex::new(String::new());

/// Set the process-global egress policy. STICKY by design: it is set
/// once when the worker boots (from `--egress-policy`, threaded from the
/// snapshot's `metadata.json["egress_policy"]`) and persists for that
/// worker's whole life, INCLUDING across cycle-restores — the restore
/// path deliberately does NOT reset it, because a pool worker always
/// serves the same snapshot, so its policy is constant. The per-dispatch
/// `RESTORE … egress_policy=<p>` path (vmm::pool) re-`set`s it only when
/// it explicitly carries one. Consequence to be aware of: if a single
/// reused worker were ever driven across DIFFERENT snapshots with
/// mismatched policies, the last `set` wins and a bare restore leaves the
/// prior policy in place — fine for the one-worker-per-snapshot model
/// used everywhere today.
pub fn set(value: &str) {
    *POLICY.lock().unwrap() = value.to_string();
}

pub fn get() -> String {
    let s = POLICY.lock().unwrap();
    if s.is_empty() {
        "allow_all".to_string()
    } else {
        s.clone()
    }
}

/// Check whether an outbound TCP target is allowed under the current
/// policy. `target` may be a SocketAddr (`"ip:port"` or `"[v6]:port"`) —
/// we only need the IP for policy checks.
pub fn check_addr(addr: std::net::SocketAddr) -> Result<(), String> {
    let ip = addr.ip();
    let policy = get();
    if policy == "allow_all" {
        return Ok(());
    }
    let private = is_private(ip);
    if policy == "deny_private" {
        return if private {
            Err(format!("policy: {ip} is in deny_private range"))
        } else {
            Ok(())
        };
    }
    if let Some(cidrs) = policy.strip_prefix("denylist:") {
        if private {
            return Err(format!("policy: {ip} is in deny_private range"));
        }
        for cidr in cidrs.split(',') {
            if cidr_contains(cidr.trim(), ip) {
                return Err(format!("policy: {ip} matched denylist {cidr}"));
            }
        }
        return Ok(());
    }
    if let Some(cidrs) = policy.strip_prefix("allowlist:") {
        for cidr in cidrs.split(',') {
            if cidr_contains(cidr.trim(), ip) {
                return Ok(());
            }
        }
        return Err(format!("policy: {ip} not in allowlist"));
    }
    Err(format!("policy: unknown {policy:?}"))
}

fn is_private(ip: IpAddr) -> bool {
    match ip {
        IpAddr::V4(v4) => {
            let o = v4.octets();
            v4.is_loopback() || v4.is_private() || v4.is_link_local()
                || v4.is_multicast() || v4.is_broadcast()
                || o[0] == 0
                || (o[0] == 100 && (o[1] & 0xc0) == 64)              // CGNAT
                || (o[0] == 192 && o[1] == 0 && o[2] == 0)
                || (o[0] == 192 && o[1] == 0 && o[2] == 2)
                || (o[0] == 198 && (o[1] == 18 || o[1] == 19))
                || (o[0] == 198 && o[1] == 51 && o[2] == 100)
                || (o[0] == 203 && o[1] == 0 && o[2] == 113)
                || o[0] >= 240
        }
        IpAddr::V6(v6) => {
            // Defense in depth: an IPv4-mapped IPv6 address
            // (`::ffff:a.b.c.d`) routes to the embedded IPv4 host, so it
            // must be classified by the (stricter) IPv4 rules — otherwise
            // `::ffff:169.254.169.254` (cloud metadata) or
            // `::ffff:10.0.0.1` would read as public and slip past
            // deny_private. Every current caller already unwraps mapped
            // addresses before calling, but folding it in here makes the
            // classifier correct on its own so a future call site can't
            // reopen the bypass.
            if let Some(v4) = v6.to_ipv4_mapped() {
                return is_private(IpAddr::V4(v4));
            }
            v6.is_loopback() || v6.is_unspecified() || v6.is_multicast() || {
                let s = v6.segments();
                (s[0] & 0xfe00) == 0xfc00 || (s[0] & 0xffc0) == 0xfe80
            }
        }
    }
}

fn cidr_contains(cidr: &str, ip: IpAddr) -> bool {
    let (net_str, prefix) = match cidr.split_once('/') {
        Some((n, p)) => (n, p.parse::<u32>().unwrap_or(u32::MAX)),
        None => (cidr, u32::MAX),
    };
    let net: IpAddr = match net_str.parse() {
        Ok(n) => n,
        Err(_) => return false,
    };
    match (net, ip) {
        (IpAddr::V4(n4), IpAddr::V4(i4)) => {
            let pfx = prefix.min(32);
            let mask: u32 = if pfx == 0 { 0 } else { !0u32 << (32 - pfx) };
            let n = u32::from_be_bytes(n4.octets());
            let i = u32::from_be_bytes(i4.octets());
            (n & mask) == (i & mask)
        }
        (IpAddr::V6(n6), IpAddr::V6(i6)) => {
            let pfx = prefix.min(128);
            let n = u128::from_be_bytes(n6.octets());
            let i = u128::from_be_bytes(i6.octets());
            let mask: u128 = if pfx == 0 { 0 } else { !0u128 << (128 - pfx) };
            (n & mask) == (i & mask)
        }
        _ => false,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::SocketAddr;
    use std::sync::Mutex;

    fn ip(s: &str) -> IpAddr {
        s.parse().unwrap()
    }
    fn sa(s: &str) -> SocketAddr {
        s.parse().unwrap()
    }

    // ── is_private: the security-critical range classifier ─────────

    #[test]
    fn is_private_public_v4_addresses_are_public() {
        for s in [
            "8.8.8.8",
            "1.1.1.1",
            "93.184.216.34",
            "100.128.0.1",
            "200.1.2.3",
        ] {
            assert!(!is_private(ip(s)), "{s} should be PUBLIC");
        }
    }

    #[test]
    fn is_private_rfc1918_ranges_are_private() {
        for s in [
            "10.0.0.1",
            "10.255.255.255", // 10/8
            "172.16.0.1",
            "172.31.255.255", // 172.16/12
            "192.168.0.1",
            "192.168.255.255", // 192.168/16
        ] {
            assert!(is_private(ip(s)), "{s} should be PRIVATE (RFC1918)");
        }
    }

    #[test]
    fn is_private_loopback_linklocal_and_metadata_endpoint() {
        assert!(is_private(ip("127.0.0.1")));
        assert!(is_private(ip("127.255.255.255")));
        assert!(is_private(ip("169.254.0.1")));
        // The cloud metadata endpoint — the whole point of deny_private.
        assert!(
            is_private(ip("169.254.169.254")),
            "metadata endpoint MUST be private"
        );
    }

    #[test]
    fn is_private_cgnat_boundary() {
        // 100.64.0.0/10 = 100.64.0.0 .. 100.127.255.255.
        assert!(is_private(ip("100.64.0.1")), "low edge of CGNAT");
        assert!(is_private(ip("100.127.255.255")), "high edge of CGNAT");
        // Just outside the /10 — these are public.
        assert!(
            !is_private(ip("100.63.255.255")),
            "just below CGNAT is public"
        );
        assert!(!is_private(ip("100.128.0.0")), "just above CGNAT is public");
    }

    #[test]
    fn is_private_special_and_reserved_v4() {
        for s in [
            "0.0.0.0",
            "0.1.2.3", // 0/8 "this network"
            "224.0.0.1",
            "239.255.255.255", // multicast
            "255.255.255.255", // broadcast
            "240.0.0.1",
            "255.0.0.0", // 240/4 reserved
            "192.0.0.1", // 192.0.0.0/24 IETF
            "192.0.2.5", // TEST-NET-1
            "198.18.0.1",
            "198.19.255.255", // benchmark 198.18/15
            "198.51.100.7",   // TEST-NET-2
            "203.0.113.9",    // TEST-NET-3
        ] {
            assert!(is_private(ip(s)), "{s} should be PRIVATE/reserved");
        }
    }

    #[test]
    fn is_private_v6_classification() {
        // Private / non-routable.
        for s in ["::1", "::", "fc00::1", "fd12:3456::1", "fe80::1", "ff02::1"] {
            assert!(is_private(ip(s)), "{s} should be PRIVATE (v6)");
        }
        // Global unicast — public.
        for s in ["2001:4860:4860::8888", "2606:4700:4700::1111"] {
            assert!(!is_private(ip(s)), "{s} should be PUBLIC (v6)");
        }
    }

    // ── cidr_contains: prefix math + parse robustness ──────────────

    #[test]
    fn cidr_exact_match_without_prefix() {
        assert!(cidr_contains("1.2.3.4", ip("1.2.3.4")));
        assert!(!cidr_contains("1.2.3.4", ip("1.2.3.5")));
    }

    #[test]
    fn cidr_v4_prefix_boundaries() {
        assert!(cidr_contains("10.0.0.0/8", ip("10.255.1.1")));
        assert!(!cidr_contains("10.0.0.0/8", ip("11.0.0.1")));
        assert!(cidr_contains("192.168.1.0/24", ip("192.168.1.255")));
        assert!(!cidr_contains("192.168.1.0/24", ip("192.168.2.0")));
        // /0 matches everything; /32 is exact.
        assert!(cidr_contains("0.0.0.0/0", ip("203.0.113.1")));
        assert!(cidr_contains("1.2.3.4/32", ip("1.2.3.4")));
        assert!(!cidr_contains("1.2.3.4/32", ip("1.2.3.5")));
    }

    #[test]
    fn cidr_prefix_out_of_range_is_clamped() {
        // /33 on a v4 net clamps to /32 (exact), doesn't panic/overflow.
        assert!(cidr_contains("1.2.3.4/33", ip("1.2.3.4")));
        assert!(!cidr_contains("1.2.3.4/99", ip("1.2.3.5")));
    }

    #[test]
    fn cidr_malformed_inputs_never_match_never_panic() {
        assert!(!cidr_contains("not-an-ip/24", ip("1.2.3.4")));
        assert!(!cidr_contains("", ip("1.2.3.4")));
        assert!(!cidr_contains("1.2.3.4/abc", ip("1.2.3.5"))); // bad prefix → u32::MAX → /32 exact
        assert!(!cidr_contains("999.999.999.999/8", ip("1.2.3.4")));
    }

    #[test]
    fn cidr_v4_v6_family_mismatch_never_matches() {
        assert!(!cidr_contains("10.0.0.0/8", ip("::1")));
        assert!(!cidr_contains("2001:db8::/32", ip("10.0.0.1")));
    }

    #[test]
    fn cidr_v6_prefix() {
        assert!(cidr_contains("2001:db8::/32", ip("2001:db8:dead:beef::1")));
        assert!(!cidr_contains("2001:db8::/32", ip("2001:db9::1")));
        assert!(cidr_contains("::/0", ip("2606:4700::1")));
    }

    #[test]
    fn cidr_v6_prefix_boundaries() {
        // /128 is an exact match; the shift (128-128=0) must not be UB.
        assert!(cidr_contains("2001:db8::1/128", ip("2001:db8::1")));
        assert!(!cidr_contains("2001:db8::1/128", ip("2001:db8::2")));
        // /127 covers ::0 and ::1 of the pair.
        assert!(cidr_contains("2001:db8::/127", ip("2001:db8::1")));
        assert!(!cidr_contains("2001:db8::/127", ip("2001:db8::2")));
        // /1 splits on the top bit only.
        assert!(cidr_contains("::/1", ip("7fff::1")));
        assert!(!cidr_contains("::/1", ip("8000::1")));
        // Over-range prefix clamps to /128 (exact), no overflow.
        assert!(cidr_contains("2001:db8::1/200", ip("2001:db8::1")));
    }

    // ── IPv4-mapped IPv6: the deny_private bypass guard ────────────────

    #[test]
    fn is_private_v4_mapped_v6_classified_by_v4_rules() {
        // `::ffff:a.b.c.d` routes to the embedded IPv4 host, so private /
        // metadata v4 targets stay private even when expressed as v6.
        assert!(
            is_private(ip("::ffff:169.254.169.254")),
            "v4-mapped metadata endpoint MUST be private"
        );
        assert!(
            is_private(ip("::ffff:10.0.0.1")),
            "v4-mapped RFC1918 is private"
        );
        assert!(
            is_private(ip("::ffff:127.0.0.1")),
            "v4-mapped loopback is private"
        );
        // A mapped PUBLIC v4 is still public.
        assert!(
            !is_private(ip("::ffff:8.8.8.8")),
            "v4-mapped public stays public"
        );
    }

    #[test]
    fn check_deny_private_blocks_v4_mapped_metadata() {
        with_policy("deny_private", || {
            assert!(
                check_addr(sa("[::ffff:169.254.169.254]:80")).is_err(),
                "deny_private MUST block the v4-mapped metadata endpoint"
            );
            assert!(check_addr(sa("[::ffff:10.0.0.1]:443")).is_err());
            // A mapped public address is still reachable.
            assert!(check_addr(sa("[::ffff:8.8.8.8]:443")).is_ok());
        });
    }

    #[test]
    fn check_allowlist_and_denylist_accept_v6_cidrs() {
        with_policy("allowlist:2001:db8::/32", || {
            assert!(check_addr(sa("[2001:db8:dead::1]:443")).is_ok());
            assert!(
                check_addr(sa("[2001:db9::1]:443")).is_err(),
                "outside allowlist"
            );
        });
        with_policy("denylist:2001:db8::/32", || {
            assert!(
                check_addr(sa("[2001:db8::99]:443")).is_err(),
                "in denied v6 CIDR"
            );
            assert!(
                check_addr(sa("[2606:4700::1]:443")).is_ok(),
                "public, not listed"
            );
        });
    }

    // ── check_addr: policy dispatch. Serialized — POLICY is global. ─
    static POLICY_LOCK: Mutex<()> = Mutex::new(());

    fn with_policy<T>(p: &str, f: impl FnOnce() -> T) -> T {
        let _g = POLICY_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        set(p);
        let r = f();
        set("allow_all"); // reset for the next test
        r
    }

    #[test]
    fn check_empty_policy_defaults_to_allow_all() {
        let _g = POLICY_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        set("");
        assert_eq!(get(), "allow_all");
        assert!(check_addr(sa("10.0.0.1:443")).is_ok());
        set("allow_all");
    }

    #[test]
    fn check_allow_all_permits_everything() {
        with_policy("allow_all", || {
            assert!(check_addr(sa("8.8.8.8:443")).is_ok());
            assert!(check_addr(sa("10.0.0.1:80")).is_ok());
            assert!(check_addr(sa("[::1]:80")).is_ok());
        });
    }

    #[test]
    fn check_deny_private_blocks_private_allows_public() {
        with_policy("deny_private", || {
            assert!(check_addr(sa("8.8.8.8:443")).is_ok(), "public allowed");
            assert!(check_addr(sa("10.0.0.1:443")).is_err());
            assert!(check_addr(sa("127.0.0.1:443")).is_err());
            // The headline security guarantee: metadata endpoint blocked.
            assert!(
                check_addr(sa("169.254.169.254:80")).is_err(),
                "deny_private MUST block the metadata endpoint"
            );
            assert!(check_addr(sa("[::1]:80")).is_err());
        });
    }

    #[test]
    fn check_denylist_blocks_private_and_extra_cidrs() {
        with_policy("denylist:1.2.3.0/24, 5.6.7.8/32", || {
            assert!(check_addr(sa("9.9.9.9:443")).is_ok(), "public, not listed");
            assert!(check_addr(sa("1.2.3.55:443")).is_err(), "in denied CIDR");
            assert!(check_addr(sa("5.6.7.8:443")).is_err(), "exact denied host");
            // deny_private base still applies under a denylist.
            assert!(
                check_addr(sa("192.168.1.1:443")).is_err(),
                "private still blocked"
            );
        });
    }

    #[test]
    fn check_allowlist_only_permits_listed() {
        with_policy("allowlist:8.8.8.0/24, 1.1.1.1/32", || {
            assert!(check_addr(sa("8.8.8.8:443")).is_ok());
            assert!(check_addr(sa("1.1.1.1:443")).is_ok());
            assert!(check_addr(sa("9.9.9.9:443")).is_err(), "not listed");
            // allowlist does NOT auto-permit private ranges.
            assert!(
                check_addr(sa("10.0.0.1:443")).is_err(),
                "private not auto-allowed"
            );
        });
    }

    #[test]
    fn check_empty_allowlist_blocks_everything() {
        with_policy("allowlist:", || {
            assert!(check_addr(sa("8.8.8.8:443")).is_err());
            assert!(check_addr(sa("1.1.1.1:443")).is_err());
        });
    }

    #[test]
    fn check_unknown_policy_denies_with_clear_error() {
        with_policy("garbage-policy", || {
            let e = check_addr(sa("8.8.8.8:443")).unwrap_err();
            assert!(e.contains("unknown"), "got: {e}");
        });
    }

    #[test]
    fn check_ignores_port() {
        with_policy("deny_private", || {
            assert!(check_addr(sa("10.0.0.1:1")).is_err());
            assert!(check_addr(sa("10.0.0.1:65535")).is_err());
            assert!(check_addr(sa("8.8.8.8:1")).is_ok());
        });
    }
}