//! Allow-listed egress for a microVM job (oxygenflow#55). The VM has no

//! network device: `flow-init` serves a proxy on the guest's 127.0.0.1

//! and relays each connection over vsock to this end, a Unix socket that

//! Firecracker connects for the guest. Here each `CONNECT host:443` is

//! checked against the sandbox's allow-list, the host is resolved (the

//! guest never resolves), any address that is not public is refused, the

//! decision goes into the job log, and the bytes are relayed untouched (no

//! TLS is opened here).

#[cfg(unix)]

use std::io::{Read, Write};

use std::net::IpAddr;

#[cfg(unix)]

use std::net::{SocketAddr, TcpStream, ToSocketAddrs};

#[cfg(unix)]

use std::os::unix::net::{UnixListener, UnixStream};

use std::path::Path;

#[cfg(unix)]

use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

use std::sync::mpsc;

#[cfg(unix)]

use std::sync::Arc;

#[cfg(unix)]

use std::time::Duration;

/// The vsock port the guest's relay connects to on the host (CID 2).

pub const VSOCK_PORT: u32 = 3128;

/// Where the job's tools find the proxy inside the guest.

pub const GUEST_PROXY: &str = "http://127.0.0.1:3128";

/// The only port relayed: TLS.

const PORT: u16 = 443;

#[cfg(unix)]

/// Connections at once, per job.

const MAX_OPEN: usize = 64;

#[cfg(unix)]

const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);

#[cfg(unix)]

/// The longest request head read.

const MAX_HEAD: usize = 8 * 1024;

/// What a CONNECT asked for, or why it was refused.

#[derive(Debug, PartialEq, Eq)]

pub enum Ask {

    Connect { host: String, port: u16 },

    Refused(String),

}

/// Read a request head: `CONNECT host:port HTTP/1.x`, and only that.

pub fn parse_head(head: &str) -> Ask {

    let line = head.lines().next().unwrap_or("");

    let mut parts = line.split_whitespace();

    let (method, target, version) = (parts.next().unwrap_or(""), parts.next().unwrap_or(""), parts.next().unwrap_or(""));

    if method != "CONNECT" {

        return Ask::Refused(format!("only https is relayed (CONNECT), not {method:?}"));

    }

    if !version.starts_with("HTTP/1.") {

        return Ask::Refused("not an HTTP/1 request".into());

    }

    let Some((host, port)) = target.rsplit_once(':') else { return Ask::Refused(format!("CONNECT {target:?} names no port")) };

    let Ok(port) = port.parse::<u16>() else { return Ask::Refused(format!("CONNECT {target:?} has no valid port")) };

    let host = host.trim_start_matches('[').trim_end_matches(']').trim_end_matches('.').to_ascii_lowercase();

    if host.is_empty() {

        return Ask::Refused("CONNECT names no host".into());

    }

    Ask::Connect { host, port }

}

/// Whether the job may go to `host:port`: the port is 443, the host is a

/// name (an address is refused: the allow-list names hosts) the list allows.

pub fn decide(allow: &[String], host: &str, port: u16) -> Result<(), String> {

    if port != PORT {

        return Err(format!("only port {PORT} is relayed"));

    }

    if host.parse::<IpAddr>().is_ok() {

        return Err("an address, not a name: the allow-list names hosts".into());

    }

    if !flow_core::egress_allows(allow, host) {

        return Err("not in the sandbox's egress allow-list".into());

    }

    Ok(())

}

/// Whether an address is on the public internet: not loopback, private,

/// shared (CGNAT), link-local (cloud metadata among them), multicast,

/// unspecified or documentation, nor an IPv6 local range or a mapped form.

pub fn public(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()

                || v4.is_unspecified()

                || v4.is_documentation()

                || o[0] == 0

                || (o[0] == 100 && (o[1] & 0xc0) == 64)

                || (o[0] == 192 && o[1] == 0 && o[2] == 0)

                || o[0] >= 240)

        }

        IpAddr::V6(v6) => {

            if let Some(v4) = v6.to_ipv4_mapped() {

                return public(IpAddr::V4(v4));

            }

            let s = v6.segments();

            !(v6.is_loopback()

                || v6.is_unspecified()

                || v6.is_multicast()

                || (s[0] & 0xfe00) == 0xfc00

                || (s[0] & 0xffc0) == 0xfe80

                || (s[0] == 0x64 && s[1] == 0xff9b)

                || (s[0] == 0x2001 && s[1] == 0x0db8))

        }

    }

}

#[cfg(unix)]

/// The proxy for one job: listening until dropped.

pub struct Proxy {

    stop: Arc<AtomicBool>,

    path: std::path::PathBuf,

}

#[cfg(unix)]

impl Drop for Proxy {

    fn drop(&mut self) {

        self.stop.store(true, Ordering::SeqCst);

        // Wake the accept loop so it sees the stop.

        let _ = UnixStream::connect(&self.path);

        let _ = std::fs::remove_file(&self.path);

    }

}

#[cfg(unix)]

/// Listen on `socket` (Firecracker's `uds_path` and `_3128`) for the

/// guest's connections; each decision is sent to `log` as a line.

pub fn start(socket: &Path, allow: Vec<String>, log: mpsc::Sender<String>) -> std::io::Result<Proxy> {

    let _ = std::fs::remove_file(socket);

    let listener = UnixListener::bind(socket)?;

    let stop = Arc::new(AtomicBool::new(false));

    let open = Arc::new(AtomicUsize::new(0));

    let allow = Arc::new(allow);

    let s = stop.clone();

    std::thread::spawn(move || {

        for conn in listener.incoming() {

            if s.load(Ordering::SeqCst) {

                break;

            }

            let Ok(conn) = conn else { continue };

            if open.load(Ordering::SeqCst) >= MAX_OPEN {

                let _ = refuse(conn, "too many connections at once");

                let _ = log.send(format!("egress: refused a connection: {MAX_OPEN} already open"));

                continue;

            }

            open.fetch_add(1, Ordering::SeqCst);

            let (allow, log, open) = (allow.clone(), log.clone(), open.clone());

            std::thread::spawn(move || {

                serve(conn, &allow, &log);

                open.fetch_sub(1, Ordering::SeqCst);

            });

        }

    });

    Ok(Proxy { stop, path: socket.to_path_buf() })

}

#[cfg(unix)]

fn refuse(mut conn: UnixStream, why: &str) -> std::io::Result<()> {

    let body = format!("{why}\n");

    write!(conn, "HTTP/1.1 403 Forbidden\r\nContent-Type: text/plain\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len())

}

#[cfg(unix)]

/// Read up to the blank line that ends a request head.

fn read_head(conn: &mut UnixStream) -> std::io::Result<String> {

    let mut head = Vec::new();

    let mut byte = [0u8; 1];

    while !head.ends_with(b"\r\n\r\n") {

        if head.len() >= MAX_HEAD || conn.read(&mut byte)? == 0 {

            break;

        }

        head.push(byte[0]);

    }

    Ok(String::from_utf8_lossy(&head).into_owned())

}

#[cfg(unix)]

fn serve(mut conn: UnixStream, allow: &[String], log: &mpsc::Sender<String>) {

    let _ = conn.set_read_timeout(Some(Duration::from_secs(30)));

    let head = match read_head(&mut conn) {

        Ok(h) => h,

        Err(_) => return,

    };

    let (host, port) = match parse_head(&head) {

        Ask::Connect { host, port } => (host, port),

        Ask::Refused(why) => {

            let _ = log.send(format!("egress: refused: {why}"));

            let _ = refuse(conn, &why);

            return;

        }

    };

    if let Err(why) = decide(allow, &host, port) {

        let _ = log.send(format!("egress: refused {host}:{port}: {why}"));

        let _ = refuse(conn, &format!("{host}:{port}: {why}"));

        return;

    }

    let addrs: Vec<SocketAddr> = match (host.as_str(), port).to_socket_addrs() {

        Ok(a) => a.collect(),

        Err(e) => {

            let _ = log.send(format!("egress: refused {host}:{port}: it could not be resolved ({e})"));

            let _ = refuse(conn, &format!("{host} could not be resolved"));

            return;

        }

    };

    if let Some(bad) = addrs.iter().find(|a| !public(a.ip())) {

        let why = format!("it resolves to {}, which is not a public address", bad.ip());

        let _ = log.send(format!("egress: refused {host}:{port}: {why}"));

        let _ = refuse(conn, &format!("{host}: {why}"));

        return;

    }

    let out = addrs.iter().find_map(|a| TcpStream::connect_timeout(a, CONNECT_TIMEOUT).ok());

    let Some(out) = out else {

        let _ = log.send(format!("egress: {host}:{port} is allowed but did not answer"));

        let _ = conn.write_all(b"HTTP/1.1 502 Bad Gateway\r\nConnection: close\r\n\r\n");

        return;

    };

    let _ = log.send(format!("egress: allowed {host}:{port}"));

    if conn.write_all(b"HTTP/1.1 200 Connection established\r\n\r\n").is_err() {

        return;

    }

    let _ = conn.set_read_timeout(None);

    relay(conn, out);

}

#[cfg(unix)]

/// Copy both ways until either side closes.

fn relay(guest: UnixStream, out: TcpStream) {

    let (Ok(mut guest_r), Ok(mut out_w)) = (guest.try_clone(), out.try_clone()) else { return };

    let up = std::thread::spawn(move || {

        let _ = std::io::copy(&mut guest_r, &mut out_w);

        let _ = out_w.shutdown(std::net::Shutdown::Write);

    });

    let (mut out_r, mut guest_w) = (out, guest);

    let _ = std::io::copy(&mut out_r, &mut guest_w);

    let _ = guest_w.shutdown(std::net::Shutdown::Write);

    let _ = up.join();

}

/// The egress proxy serves a microVM's vsock, which only Linux has.

#[cfg(not(unix))]

pub struct Proxy;

#[cfg(not(unix))]

pub fn start(_socket: &Path, _allow: Vec<String>, _log: mpsc::Sender<String>) -> std::io::Result<Proxy> {

    Err(std::io::Error::new(std::io::ErrorKind::Unsupported, "egress serves microVMs, which run only on Linux"))

}

#[cfg(all(test, unix))]

mod tests {

    use super::*;

    fn allow() -> Vec<String> {

        vec!["*.googleapis.com".into(), "registry.opentofu.org".into()]

    }

    #[test]

    fn only_connect_to_an_allowed_name_on_443_is_let_through() {

        assert_eq!(

            parse_head("CONNECT storage.googleapis.com:443 HTTP/1.1\r\nHost: x\r\n\r\n"),

            Ask::Connect { host: "storage.googleapis.com".into(), port: 443 }

        );

        assert!(matches!(parse_head("GET http://example.com/ HTTP/1.1\r\n\r\n"), Ask::Refused(w) if w.contains("only https")));

        assert!(matches!(parse_head("CONNECT example.com HTTP/1.1\r\n\r\n"), Ask::Refused(w) if w.contains("no port")));

        assert_eq!(decide(&allow(), "storage.googleapis.com", 443), Ok(()));

        assert_eq!(decide(&allow(), "registry.opentofu.org", 443), Ok(()));

        assert!(decide(&allow(), "example.com", 443).unwrap_err().contains("allow-list"));

        assert!(decide(&allow(), "registry.opentofu.org", 22).unwrap_err().contains("only port 443"));

        assert!(decide(&allow(), "142.250.1.1", 443).unwrap_err().contains("names hosts"));

    }

    #[test]

    fn private_addresses_are_not_public() {

        for ip in [

            "127.0.0.1",

            "10.1.2.3",

            "192.168.1.1",

            "172.16.0.1",

            "169.254.169.254",

            "100.64.0.1",

            "0.0.0.0",

            "::1",

            "fd00::1",

            "fe80::1",

            "::ffff:10.0.0.1",

        ] {

            assert!(!public(ip.parse().unwrap()), "{ip}");

        }

        for ip in ["142.250.1.1", "2a00:1450:4009::200e"] {

            assert!(public(ip.parse().unwrap()), "{ip}");

        }

    }

    #[test]

    fn the_proxy_refuses_and_logs_a_host_not_allowed() {

        let dir = tempfile::tempdir().unwrap();

        let sock = dir.path().join("v.sock_3128");

        let (tx, rx) = mpsc::channel();

        let proxy = start(&sock, allow(), tx).unwrap();

        let ask = |head: &str| {

            let mut c = UnixStream::connect(&sock).unwrap();

            c.write_all(head.as_bytes()).unwrap();

            let mut answer = String::new();

            c.read_to_string(&mut answer).unwrap();

            answer

        };

        let a = ask("CONNECT example.com:443 HTTP/1.1\r\n\r\n");

        assert!(a.starts_with("HTTP/1.1 403") && a.contains("not in the sandbox's egress allow-list"), "{a}");

        assert_eq!(

            rx.recv_timeout(Duration::from_secs(5)).unwrap(),

            "egress: refused example.com:443: not in the sandbox's egress allow-list"

        );

        let a = ask("GET http://storage.googleapis.com/ HTTP/1.1\r\n\r\n");

        assert!(a.starts_with("HTTP/1.1 403"), "{a}");

        assert!(rx.recv_timeout(Duration::from_secs(5)).unwrap().contains("only https"));

        drop(proxy);

        assert!(!sock.exists(), "the socket goes with the job");

    }

}