//! The egress proxy: the only way `http_fetch`'s container reaches the network. It speaks the part
//! of SOCKS5 (RFC 1928) that `curl --proxy socks5h://` uses, allows only the call's hosts, only
//! port 443, only public addresses, and after the handshake copies bytes both ways without looking
//! at them. Everything read during the handshake comes from the untrusted side: read exactly what
//! the protocol says, never index or allocate by a number that has not been checked. Spec section 5.

use std::io::{Read, Write};
use std::net::{SocketAddr, TcpStream, ToSocketAddrs};
use std::os::unix::net::{UnixListener, UnixStream};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};

pub const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
pub const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
pub const MAX_CONNECTIONS: usize = 8;
// SOCKS5 reply codes the proxy sends:
pub const NOT_ALLOWED: u8 = 2;
pub const HOST_UNREACHABLE: u8 = 4;
pub const CONNECTION_REFUSED: u8 = 5;
pub const COMMAND_NOT_SUPPORTED: u8 = 7;
pub const ADDRESS_TYPE_NOT_SUPPORTED: u8 = 8;

/// The call's hosts: one or more patterns separated by ','. An empty piece is an error.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Allow {
    patterns: Vec<String>,
}

impl Allow {
    /// One or more host patterns separated by ','. Every piece must pass
    /// `proto::hosts::valid_host_pattern`; an empty piece (",", "a.com,", "") is an error.
    pub fn parse(list: &str) -> Result<Allow, String> {
        let mut patterns = Vec::new();
        for piece in list.split(',') {
            if !proto::hosts::valid_host_pattern(piece) {
                return Err(format!("'{piece}' is not a host or a *.host"));
            }
            patterns.push(piece.to_string());
        }
        Ok(Allow { patterns })
    }

    /// `valid_host(host)` and some pattern `host_matches` it.
    pub fn permits(&self, host: &str) -> bool {
        proto::hosts::valid_host(host)
            && self
                .patterns
                .iter()
                .any(|p| proto::hosts::host_matches(p, host))
    }
}

/// Name resolution and connecting, so tests need no network.
pub trait Dial: Send + Sync {
    fn resolve(&self, host: &str, port: u16) -> std::io::Result<Vec<SocketAddr>>;
    fn connect(&self, addr: SocketAddr, timeout: Duration) -> std::io::Result<TcpStream>;
}
/// The system's: `(host, port).to_socket_addrs()` and `TcpStream::connect_timeout`.
pub struct SystemDial;

impl Dial for SystemDial {
    fn resolve(&self, host: &str, port: u16) -> std::io::Result<Vec<SocketAddr>> {
        let addrs = format!("{}:{}", host, port).to_socket_addrs()?;
        Ok(addrs.collect())
    }
    fn connect(&self, addr: SocketAddr, timeout: Duration) -> std::io::Result<TcpStream> {
        TcpStream::connect_timeout(&addr, timeout)
    }
}

pub struct Proxy {
    allow: Allow,
    dial: Arc<dyn Dial>,
    handshake: Duration,
}
impl Proxy {
    pub fn new(allow: Allow, dial: Arc<dyn Dial>) -> Proxy {
        Proxy {
            allow,
            dial,
            handshake: HANDSHAKE_TIMEOUT,
        }
    }
    pub fn with_handshake_timeout(self, handshake: Duration) -> Proxy {
        Proxy { handshake, ..self }
    }
    /// Accept for ever; at most MAX_CONNECTIONS handled at once, more are closed at once.
    pub fn serve(self: Arc<Self>, listener: UnixListener) -> std::io::Result<()> {
        let count = Arc::new(AtomicUsize::new(0));
        for stream in listener.incoming() {
            let stream = stream?;
            if count.load(Ordering::SeqCst) >= MAX_CONNECTIONS {
                drop(stream);
                continue;
            }
            count.fetch_add(1, Ordering::SeqCst);
            let proxy = Arc::clone(&self);
            let closer = Arc::clone(&count);
            let decrement = Arc::clone(&count);
            let worker = match stream.try_clone() {
                Ok(w) => w,
                Err(_) => {
                    decrement.fetch_sub(1, Ordering::SeqCst);
                    drop(stream);
                    continue;
                }
            };
            match start(move || {
                proxy.handle(worker);
                closer.fetch_sub(1, Ordering::SeqCst);
            }) {
                Ok(_) => {}
                Err(_) => {
                    decrement.fetch_sub(1, Ordering::SeqCst);
                    drop(stream);
                }
            }
        }
        Ok(())
    }
    /// One connection, from greeting to the end of the copy.
    pub fn handle(&self, client: UnixStream) {
        let mut client = client;
        let deadline = Instant::now() + self.handshake;

        // 1. version, count. Not SOCKS5 -> stop, no reply.
        let greeting = match read_n(&mut client, 2, deadline) {
            Some(b) if b[0] == 5 => b,
            _ => return,
        };

        // 2. methods. No auth method 0 (unauthenticated) -> refuse, else accept.
        let methods = match read_n(&mut client, greeting[1] as usize, deadline) {
            Some(b) => b,
            None => return,
        };
        if methods.contains(&0) {
            let _ = client.write_all(&[5, 0]);
        } else {
            let _ = client.write_all(&[5, 0xff]);
            return;
        }

        // 3. version, command, reserved, kind.
        let req = match read_n(&mut client, 4, deadline) {
            Some(b) => b,
            None => return,
        };
        if req[0] != 5 || req[2] != 0 {
            return;
        }

        // 4. only connect.
        if req[1] != 1 {
            return refuse(&mut client, COMMAND_NOT_SUPPORTED);
        }
        // 5. only domain names; do not read the address.
        if req[3] != 3 {
            return refuse(&mut client, ADDRESS_TYPE_NOT_SUPPORTED);
        }
        // 6. name length; zero is not allowed.
        let length = match read_n(&mut client, 1, deadline) {
            Some(b) => b[0],
            None => return,
        } as usize;
        if length == 0 {
            return refuse(&mut client, NOT_ALLOWED);
        }
        // 7. name, then port.
        let name = match read_n(&mut client, length, deadline) {
            Some(b) => b,
            None => return,
        };
        let port = match read_n(&mut client, 2, deadline) {
            Some(b) => b,
            None => return,
        };
        let port = u16::from_be_bytes([port[0], port[1]]);
        // 8. UTF-8, port 443, host allowed.
        let name = match std::str::from_utf8(&name) {
            Ok(s) => s,
            Err(_) => return refuse(&mut client, NOT_ALLOWED),
        };
        if port != 443 || !self.allow.permits(name) {
            return refuse(&mut client, NOT_ALLOWED);
        }
        // 9. resolve; first public address, others skipped never tried.
        let addrs = match self.dial.resolve(name, port) {
            Ok(a) => a,
            Err(_) => return refuse(&mut client, HOST_UNREACHABLE),
        };
        let addr = match addrs
            .iter()
            .find(|a| crate::addr::is_public(a.ip()))
            .copied()
        {
            Some(a) => a,
            None => return refuse(&mut client, HOST_UNREACHABLE),
        };
        // 10. connect.
        let server = match self.dial.connect(addr, CONNECT_TIMEOUT) {
            Ok(s) => s,
            Err(_) => return refuse(&mut client, CONNECTION_REFUSED),
        };
        // 11. success reply.
        if client.write_all(&[5, 0, 0, 1, 0, 0, 0, 0, 0, 0]).is_err() {
            return;
        }

        copy_both_ways(client, server);
    }
}

/// Start a thread that never panics on input: `std::thread::Builder`, which returns an error
/// instead of unwinding, unlike `spawn`.
fn start<F>(f: F) -> std::io::Result<JoinHandle<()>>
where
    F: FnOnce() + Send + 'static,
{
    std::thread::Builder::new()
        .name("egress-handle".into())
        .spawn(f)
}

/// Refuse: write the SOCKS5 reply and return, which closes the connection.
fn refuse(client: &mut UnixStream, code: u8) {
    let _ = client.write_all(&[5, code, 0, 1, 0, 0, 0, 0, 0, 0]);
}

/// Read exactly `n` bytes before `deadline`, never more. The whole handshake shares one deadline:
/// before every read, set the read timeout to the time left; if none is left, or a read times out,
/// fails, or reads 0 bytes (the peer closed), stop: return `None`.
fn read_n(stream: &mut UnixStream, n: usize, deadline: Instant) -> Option<Vec<u8>> {
    let mut out = Vec::with_capacity(n.min(255));
    while out.len() < n {
        let remaining = deadline.checked_duration_since(Instant::now())?;
        if stream.set_read_timeout(Some(remaining)).is_err() {
            return None;
        }
        let mut byte = [0u8; 1];
        match stream.read(&mut byte) {
            Ok(0) | Err(_) => return None,
            Ok(_) => out.push(byte[0]),
        }
    }
    Some(out)
}

/// After the handshake: copy bytes both ways, passing a half-close on. The client's read timeout
/// is cleared so the copy can run to its natural end.
fn copy_both_ways(mut client: UnixStream, mut server: TcpStream) {
    let _ = client.set_read_timeout(None);
    let mut server_write = match server.try_clone() {
        Ok(s) => s,
        Err(_) => return,
    };
    let mut client_read = match client.try_clone() {
        Ok(s) => s,
        Err(_) => return,
    };

    // Second thread: client -> server, then half-close the server's write side.
    let upstream = match std::thread::Builder::new()
        .name("egress-upstream".into())
        .spawn(move || {
            let _ = std::io::copy(&mut client_read, &mut server_write);
            let _ = server_write.shutdown(std::net::Shutdown::Write);
        }) {
        Ok(handle) => handle,
        Err(_) => return,
    };

    // Own thread: server -> client, then half-close the client's write side, then join.
    let _ = std::io::copy(&mut server, &mut client);
    let _ = client.shutdown(std::net::Shutdown::Write);
    let _ = upstream.join();
}