//! The egress proxy against a fake resolver and a local echo server: every reply code, the host,
//! port and address checks, the byte copy both ways with half-close, the handshake deadline, the
//! connection limit, and `toolkit egress-proxy` as a program. No test needs a network. Do not edit.

mod support;

use std::collections::HashMap;
use std::io::{Read, Write};
use std::net::{Shutdown, SocketAddr, TcpListener, TcpStream};
use std::os::unix::net::{UnixListener, UnixStream};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use support::TempDir;
use toolkit::egress::{
    ADDRESS_TYPE_NOT_SUPPORTED, Allow, COMMAND_NOT_SUPPORTED, CONNECTION_REFUSED, Dial,
    HOST_UNREACHABLE, MAX_CONNECTIONS, NOT_ALLOWED, Proxy,
};

/// Names resolve from a table; every connection goes to one local echo server, except to
/// addresses listed as refusing. Records every address it was asked to connect to.
struct FakeDial {
    names: HashMap<String, Vec<SocketAddr>>,
    refusing: Vec<SocketAddr>,
    echo: SocketAddr,
    connected: Mutex<Vec<SocketAddr>>,
}

impl Dial for FakeDial {
    fn resolve(&self, host: &str, _port: u16) -> std::io::Result<Vec<SocketAddr>> {
        self.names
            .get(host)
            .cloned()
            .ok_or_else(|| std::io::Error::other("no such name"))
    }
    fn connect(&self, addr: SocketAddr, _timeout: Duration) -> std::io::Result<TcpStream> {
        self.connected.lock().unwrap().push(addr);
        if self.refusing.contains(&addr) {
            return Err(std::io::Error::from(std::io::ErrorKind::ConnectionRefused));
        }
        TcpStream::connect(self.echo)
    }
}

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

/// An echo server: copies back what it reads, and closes its side after reading the end.
fn echo_server() -> SocketAddr {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap();
    std::thread::spawn(move || {
        for stream in listener.incoming() {
            let Ok(mut stream) = stream else { continue };
            std::thread::spawn(move || {
                let mut back = stream.try_clone().unwrap();
                let _ = std::io::copy(&mut stream, &mut back);
                let _ = back.shutdown(Shutdown::Write);
            });
        }
    });
    addr
}

fn dial() -> Arc<FakeDial> {
    let mut names = HashMap::new();
    names.insert("example.com".to_string(), vec![sa("93.184.216.34:443")]);
    names.insert(
        "www.example.org".to_string(),
        vec![sa("[2606:2800::1]:443")],
    );
    names.insert(
        "mixed.example.com".to_string(),
        vec![
            sa("10.0.0.1:443"),
            sa("100.100.100.100:443"),
            sa("1.1.1.1:443"),
        ],
    );
    names.insert(
        "inside.example.com".to_string(),
        vec![
            sa("100.101.102.103:443"),
            sa("127.0.0.1:443"),
            sa("[::1]:443"),
        ],
    );
    names.insert("refusing.example.com".to_string(), vec![sa("8.8.8.8:443")]);
    Arc::new(FakeDial {
        names,
        refusing: vec![sa("8.8.8.8:443")],
        echo: echo_server(),
        connected: Mutex::new(Vec::new()),
    })
}

fn allow() -> Allow {
    Allow::parse(
        "example.com,*.example.org,mixed.example.com,inside.example.com,refusing.example.com,nowhere.example.com",
    )
    .unwrap()
}

/// A proxy on one end of a socket pair, handling that one connection on its own thread.
fn start(dial: Arc<FakeDial>) -> (UnixStream, std::thread::JoinHandle<()>) {
    let (client, server) = UnixStream::pair().unwrap();
    client
        .set_read_timeout(Some(Duration::from_secs(5)))
        .unwrap();
    let proxy = Proxy::new(allow(), dial).with_handshake_timeout(Duration::from_millis(300));
    let handle = std::thread::spawn(move || proxy.handle(server));
    (client, handle)
}

fn request(host: &str, port: u16) -> Vec<u8> {
    let mut bytes = vec![5, 1, 0, 3, u8::try_from(host.len()).unwrap()];
    bytes.extend_from_slice(host.as_bytes());
    bytes.extend_from_slice(&port.to_be_bytes());
    bytes
}

fn read_exactly(client: &mut UnixStream, n: usize) -> Vec<u8> {
    let mut buf = vec![0u8; n];
    client.read_exact(&mut buf).unwrap();
    buf
}

/// Read until the proxy closes; the bytes read. A close that leaves some of our bytes unread by
/// the proxy arrives as "connection reset" rather than as the end: both mean closed. A timeout
/// does not: the proxy did not close.
fn read_to_close(client: &mut UnixStream) -> Vec<u8> {
    let mut rest = Vec::new();
    let mut buf = [0u8; 4096];
    loop {
        match client.read(&mut buf) {
            Ok(0) => return rest,
            Ok(n) => rest.extend_from_slice(&buf[..n]),
            Err(e) if e.kind() == std::io::ErrorKind::ConnectionReset => return rest,
            Err(e) => panic!("the proxy did not close: {e}"),
        }
    }
}

/// Greeting, then `req`; the reply's code, and that the proxy then closed (for failures).
fn reply_code(req: &[u8]) -> u8 {
    let (mut client, handle) = start(dial());
    client.write_all(&[5, 1, 0]).unwrap();
    assert_eq!(read_exactly(&mut client, 2), [5, 0]);
    client.write_all(req).unwrap();
    let reply = read_exactly(&mut client, 10);
    assert_eq!(reply[0], 5);
    assert_eq!(&reply[2..], &[0, 1, 0, 0, 0, 0, 0, 0]);
    if reply[1] != 0 {
        assert_eq!(
            read_to_close(&mut client),
            b"",
            "a refusal is followed by the close"
        );
        handle.join().unwrap();
    }
    reply[1]
}

#[test]
fn an_allowed_host_is_connected_and_bytes_flow_both_ways_until_both_ends_close() {
    let dial = dial();
    let (mut client, handle) = start(Arc::clone(&dial));
    client.write_all(&[5, 2, 2, 0]).unwrap();
    assert_eq!(
        read_exactly(&mut client, 2),
        [5, 0],
        "method 0 among others"
    );
    client.write_all(&request("example.com", 443)).unwrap();
    assert_eq!(
        read_exactly(&mut client, 10),
        [5, 0, 0, 1, 0, 0, 0, 0, 0, 0]
    );
    client.write_all(b"hello through the tunnel").unwrap();
    assert_eq!(read_exactly(&mut client, 24), b"hello through the tunnel");
    client.shutdown(Shutdown::Write).unwrap();
    assert_eq!(
        read_to_close(&mut client),
        b"",
        "the half-close reached the server and back"
    );
    handle.join().unwrap();
    assert_eq!(
        *dial.connected.lock().unwrap(),
        vec![sa("93.184.216.34:443")]
    );
}

#[test]
fn a_wildcard_pattern_allows_a_name_under_it() {
    assert_eq!(reply_code(&request("www.example.org", 443)), 0);
}

#[test]
fn a_greeting_without_method_zero_is_answered_ff_and_closed() {
    let (mut client, handle) = start(dial());
    client.write_all(&[5, 2, 1, 2]).unwrap();
    assert_eq!(read_exactly(&mut client, 2), [5, 0xff]);
    assert_eq!(read_to_close(&mut client), b"");
    handle.join().unwrap();
}

#[test]
fn another_version_is_closed_without_a_word() {
    let (mut client, handle) = start(dial());
    client.write_all(&[4, 1, 0]).unwrap();
    assert_eq!(read_to_close(&mut client), b"");
    handle.join().unwrap();
}

#[test]
fn only_connect_by_name_to_an_allowed_host_on_443() {
    let mut bind = request("example.com", 443);
    bind[1] = 2;
    assert_eq!(reply_code(&bind), COMMAND_NOT_SUPPORTED);
    let mut udp = request("example.com", 443);
    udp[1] = 3;
    assert_eq!(reply_code(&udp), COMMAND_NOT_SUPPORTED);
    assert_eq!(
        reply_code(&[5, 1, 0, 1, 93, 184, 216, 34, 1, 187]),
        ADDRESS_TYPE_NOT_SUPPORTED
    );
    let mut v6 = vec![5, 1, 0, 4];
    v6.extend_from_slice(&[0; 16]);
    v6.extend_from_slice(&443u16.to_be_bytes());
    assert_eq!(reply_code(&v6), ADDRESS_TYPE_NOT_SUPPORTED);
    for (host, port) in [
        ("example.com", 80),
        ("example.com", 8443),
        ("evil.test", 443),
        ("example.org", 443), // *.example.org does not cover example.org
        ("Example.com", 443), // not a valid host name
        ("example.com.", 443),
        ("127.0.0.1", 443), // an IP literal sent as a name
        ("wwwexample.com", 443),
    ] {
        assert_eq!(
            reply_code(&request(host, port)),
            NOT_ALLOWED,
            "{host}:{port}"
        );
    }
}

#[test]
fn an_empty_or_non_utf8_name_is_not_allowed() {
    assert_eq!(reply_code(&[5, 1, 0, 3, 0, 1, 187]), NOT_ALLOWED);
    assert_eq!(
        reply_code(&[5, 1, 0, 3, 2, 0xff, 0xfe, 1, 187]),
        NOT_ALLOWED
    );
}

#[test]
fn only_public_addresses_are_used() {
    let dial = dial();
    let (mut client, handle) = start(Arc::clone(&dial));
    client.write_all(&[5, 1, 0]).unwrap();
    read_exactly(&mut client, 2);
    client
        .write_all(&request("mixed.example.com", 443))
        .unwrap();
    assert_eq!(read_exactly(&mut client, 10)[1], 0);
    drop(client);
    handle.join().unwrap();
    assert_eq!(
        *dial.connected.lock().unwrap(),
        vec![sa("1.1.1.1:443")],
        "the private and tailnet addresses were skipped, not tried"
    );
}

#[test]
fn a_name_with_no_public_address_or_no_address_is_unreachable() {
    assert_eq!(
        reply_code(&request("inside.example.com", 443)),
        HOST_UNREACHABLE
    );
    assert_eq!(
        reply_code(&request("nowhere.example.com", 443)),
        HOST_UNREACHABLE
    );
}

#[test]
fn a_refused_connection_is_reply_5() {
    assert_eq!(
        reply_code(&request("refusing.example.com", 443)),
        CONNECTION_REFUSED
    );
}

#[test]
fn a_stalled_handshake_is_closed_at_the_deadline() {
    let (mut client, handle) = start(dial());
    client.write_all(&[5, 1, 0]).unwrap();
    assert_eq!(read_exactly(&mut client, 2), [5, 0]);
    client.write_all(&[5, 1]).unwrap(); // half a request, then nothing
    let started = Instant::now();
    assert_eq!(read_to_close(&mut client), b"");
    let took = started.elapsed();
    handle.join().unwrap();
    assert!(
        took < Duration::from_secs(2),
        "one deadline for the handshake: {took:?}"
    );
}

#[test]
fn a_handshake_that_trickles_still_ends_at_the_deadline() {
    let (mut client, handle) = start(dial());
    let started = Instant::now();
    // One byte every 100 ms: each read is quick, but the handshake as a whole is not.
    for byte in [5u8, 1, 0, 5, 1, 0, 3, 11] {
        if client.write_all(&[byte]).is_err() {
            break;
        }
        std::thread::sleep(Duration::from_millis(100));
    }
    let _ = read_to_close(&mut client);
    handle.join().unwrap();
    assert!(started.elapsed() < Duration::from_secs(2));
}

#[test]
fn allow_lists_parse_strictly() {
    assert!(Allow::parse("example.com").is_ok());
    assert!(Allow::parse("example.com,*.example.org").is_ok());
    for bad in [
        "",
        ",",
        "example.com,",
        "a.com,,b.com",
        "Example.com",
        "*",
        "10.0.0.1",
    ] {
        assert!(Allow::parse(bad).is_err(), "{bad:?}");
    }
    let allow = Allow::parse("*.example.org").unwrap();
    assert!(allow.permits("a.example.org"));
    assert!(!allow.permits("example.org"));
    assert!(!allow.permits("A.example.org"));
}

#[test]
fn more_than_the_limit_of_connections_are_closed_at_once() {
    let dir = TempDir::new("egress-limit");
    let path = dir.path().join("egress.sock");
    let listener = UnixListener::bind(&path).unwrap();
    let proxy =
        Arc::new(Proxy::new(allow(), dial()).with_handshake_timeout(Duration::from_secs(5)));
    std::thread::spawn(move || {
        let _ = proxy.serve(listener);
    });
    // MAX_CONNECTIONS clients that greet and then wait, holding their places.
    let mut held = Vec::new();
    for _ in 0..MAX_CONNECTIONS {
        let mut c = UnixStream::connect(&path).unwrap();
        c.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
        c.write_all(&[5, 1, 0]).unwrap();
        assert_eq!(read_exactly(&mut c, 2), [5, 0]);
        held.push(c);
    }
    let mut extra = UnixStream::connect(&path).unwrap();
    extra
        .set_read_timeout(Some(Duration::from_secs(2)))
        .unwrap();
    let _ = extra.write_all(&[5, 1, 0]);
    let mut buf = [0u8; 2];
    assert!(
        matches!(extra.read(&mut buf), Ok(0) | Err(_)),
        "the connection over the limit gets no answer"
    );
    // When one place is freed, a new connection is served again.
    drop(held.pop());
    std::thread::sleep(Duration::from_millis(200));
    let mut again = UnixStream::connect(&path).unwrap();
    again
        .set_read_timeout(Some(Duration::from_secs(5)))
        .unwrap();
    again.write_all(&[5, 1, 0]).unwrap();
    assert_eq!(read_exactly(&mut again, 2), [5, 0]);
}

#[test]
fn the_program_listens_where_it_is_told_and_refuses_what_is_not_allowed() {
    let dir = TempDir::new("egress-prog");
    let path = dir.path().join("egress.sock");
    let mut child = std::process::Command::new(env!("CARGO_BIN_EXE_toolkit"))
        .args([
            "egress-proxy",
            "--socket",
            path.to_str().unwrap(),
            "--allow",
            "example.com",
        ])
        .spawn()
        .unwrap();
    let until = Instant::now() + Duration::from_secs(5);
    let mut client = loop {
        if let Ok(c) = UnixStream::connect(&path) {
            break c;
        }
        assert!(Instant::now() < until, "the proxy never listened");
        std::thread::sleep(Duration::from_millis(20));
    };
    client
        .set_read_timeout(Some(Duration::from_secs(5)))
        .unwrap();
    client.write_all(&[5, 1, 0]).unwrap();
    assert_eq!(read_exactly(&mut client, 2), [5, 0]);
    client.write_all(&request("evil.test", 443)).unwrap();
    assert_eq!(read_exactly(&mut client, 10)[1], NOT_ALLOWED);
    child.kill().unwrap();
    child.wait().unwrap();
}

#[test]
fn the_program_refuses_bad_arguments() {
    let dir = TempDir::new("egress-args");
    let path = dir.path().join("egress.sock");
    let sock = path.to_str().unwrap();
    std::fs::write(dir.path().join("taken"), "").unwrap();
    let taken = dir.path().join("taken");
    for args in [
        vec!["egress-proxy", "--socket", sock, "--allow", "Example.com"],
        vec!["egress-proxy", "--socket", sock, "--allow", ""],
        vec![
            "egress-proxy",
            "--socket",
            taken.to_str().unwrap(),
            "--allow",
            "example.com",
        ],
        vec!["egress-proxy", "--socket", sock],
        vec!["egress-proxy", "--allow", "example.com", "--socket", sock],
    ] {
        let out = std::process::Command::new(env!("CARGO_BIN_EXE_toolkit"))
            .args(&args)
            .output()
            .unwrap();
        assert_eq!(out.status.code(), Some(2), "{args:?}");
    }
}