//! 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).

use std::io::{Read, Write};
use std::net::{IpAddr, SocketAddr, TcpStream, ToSocketAddrs};
use std::os::unix::net::{UnixListener, UnixStream};
use std::path::Path;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{mpsc, Arc};
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;
/// Connections at once, per job.
const MAX_OPEN: usize = 64;
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
/// 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))
        }
    }
}

/// The proxy for one job: listening until dropped.
pub struct Proxy {
    stop: Arc<AtomicBool>,
    path: std::path::PathBuf,
}

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

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

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

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

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

/// 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();
}

#[cfg(test)]
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");
    }
}