//! The guest's half of allow-listed egress (oxygenflow#55): a proxy on

//! 127.0.0.1:3128 that relays each connection, unread, over vsock to the

//! runner (host CID 2), which decides what may pass. The VM has no network

//! device; this is its only way out.

use std::fs::File;

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

use std::net::{TcpListener, TcpStream};

use std::os::fd::{AsRawFd, FromRawFd};

/// Where the job's tools reach the proxy, and the vsock port behind it.

const LISTEN: &str = "127.0.0.1:3128";

const VSOCK_PORT: u32 = 3128;

/// Start relaying in the background (before the job starts, so its first

/// connection finds the proxy listening).

pub fn start() -> Result<(), String> {

    let listener = TcpListener::bind(LISTEN).map_err(|e| format!("egress proxy on {LISTEN}: {e}"))?;

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

        for conn in listener.incoming().flatten() {

            std::thread::spawn(move || relay(conn));

        }

    });

    Ok(())

}

fn vsock_to_host(port: u32) -> std::io::Result<File> {

    let fd = unsafe { libc::socket(libc::AF_VSOCK, libc::SOCK_STREAM | libc::SOCK_CLOEXEC, 0) };

    if fd < 0 {

        return Err(std::io::Error::last_os_error());

    }

    // Owned from here, so it is closed on every path.

    let file = unsafe { File::from_raw_fd(fd) };

    let mut addr: libc::sockaddr_vm = unsafe { std::mem::zeroed() };

    addr.svm_family = libc::AF_VSOCK as libc::sa_family_t;

    addr.svm_cid = libc::VMADDR_CID_HOST;

    addr.svm_port = port;

    let rc = unsafe {

        libc::connect(

            fd,

            &addr as *const libc::sockaddr_vm as *const libc::sockaddr,

            std::mem::size_of::<libc::sockaddr_vm>() as libc::socklen_t,

        )

    };

    if rc != 0 {

        return Err(std::io::Error::last_os_error());

    }

    Ok(file)

}

fn relay(tcp: TcpStream) {

    let Ok(host) = vsock_to_host(VSOCK_PORT) else {

        let mut tcp = tcp;

        let _ = tcp.write_all(b"HTTP/1.1 502 Bad Gateway\r\nConnection: close\r\n\r\nthe runner's egress proxy did not answer\n");

        return;

    };

    let (Ok(mut tcp_r), Ok(mut host_w)) = (tcp.try_clone(), host.try_clone()) else { return };

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

        copy(&mut tcp_r, &mut host_w);

        unsafe { libc::shutdown(host_w.as_raw_fd(), libc::SHUT_WR) };

    });

    let (mut host_r, mut tcp_w) = (host, tcp);

    copy(&mut host_r, &mut tcp_w);

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

    let _ = up.join();

}

fn copy(from: &mut impl Read, to: &mut impl Write) {

    let mut buf = [0u8; 16 * 1024];

    loop {

        match from.read(&mut buf) {

            Ok(0) | Err(_) => return,

            Ok(n) => {

                if to.write_all(&buf[..n]).is_err() {

                    return;

                }

            }

        }

    }

}