mod cipher;
mod config;
mod handshake;
mod relay;

use std::net::{Shutdown, TcpListener, TcpStream, ToSocketAddrs};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use config::{Addr, Config, Mode};
use relay::CipherSide;

/// A counting gate that bounds how many connections we service at once. When the
/// limit is reached `try_acquire` returns `None`, and the caller hangs up the
/// newcomer at once rather than holding a connection open without servicing it.
struct Semaphore {
    count: Mutex<usize>,
}

impl Semaphore {
    fn new(limit: usize) -> Arc<Self> {
        Arc::new(Self {
            count: Mutex::new(limit),
        })
    }

    /// Take a slot if one is free, without ever blocking. Returns `None` when we
    /// are already at the limit.
    fn try_acquire(self: &Arc<Self>) -> Option<Permit> {
        let mut count = self.count.lock().unwrap_or_else(|e| e.into_inner());
        if *count == 0 {
            return None;
        }
        *count -= 1;
        Some(Permit {
            sem: Arc::clone(self),
        })
    }
}

struct Permit {
    sem: Arc<Semaphore>,
}

impl Drop for Permit {
    fn drop(&mut self) {
        let mut count = self.sem.count.lock().unwrap_or_else(|e| e.into_inner());
        *count += 1;
    }
}

/// Panic-safe diagnostics: write to stderr and swallow any error. The macro
/// `eprintln!` panics if stderr is gone (a closed pipe), which we never want.
fn diag(msg: &str) {
    use std::io::Write;
    let _ = writeln!(std::io::stderr(), "{msg}");
}

/// Panic-safe stdout write, used for `--help`.
fn out(msg: &str) {
    use std::io::Write;
    let _ = std::io::stdout().write_all(msg.as_bytes());
}

fn main() {
    let env_password = std::env::var(config::PASSWORD_ENV).ok();
    let config = match config::parse(std::env::args().skip(1), env_password) {
        Ok(c) => c,
        Err(msg) => {
            // --help routes through Err too; print it cleanly either way.
            if msg.starts_with("jokycrypt") {
                out(&msg);
                std::process::exit(0);
            }
            diag(&format!("error: {msg}\n\nTry 'jow --help' for usage."));
            std::process::exit(2);
        }
    };

    if let Err(e) = serve(config) {
        diag(&format!("fatal: {e}"));
        std::process::exit(1);
    }
}

fn serve(config: Config) -> std::io::Result<()> {
    let listener = TcpListener::bind((config.listen.host.as_str(), config.listen.port))?;
    let role = match config.mode {
        Mode::Client => "client (encrypting toward upstream)",
        Mode::Server => "server (decrypting toward target)",
    };
    diag(&format!(
        "jokycrypt {role}: listening on {}:{}",
        config.listen.host, config.listen.port
    ));
    match config.mode {
        Mode::Server => {
            let list = config
                .targets
                .iter()
                .map(Addr::to_string)
                .collect::<Vec<_>>()
                .join(", ");
            diag(&format!("  allowed targets: {list}"));
            if let Some(first) = config.targets.first() {
                diag(&format!(
                    "  default target (when the client asks for none): {first}"
                ));
            }
        }
        Mode::Client => {
            let server = &config.targets[0];
            diag(&format!("  upstream server: {server}"));
            match &config.requested_target {
                Some(t) => diag(&format!("  requesting target: {t}")),
                None => diag("  requesting target: the server's default"),
            }
        }
    }
    if let Some(n) = config.max_conns {
        diag(&format!("  max simultaneous connections: {n}"));
    }
    if let Some(t) = config.timeout {
        diag(&format!("  idle timeout: {}s", t.as_secs()));
    }
    if config.verbose {
        diag("  verbose: printing plaintext traffic");
    }

    let config = Arc::new(config);
    let semaphore = config.max_conns.map(Semaphore::new);

    for incoming in listener.incoming() {
        let stream = match incoming {
            Ok(s) => s,
            Err(e) => {
                // Transient accept failures (e.g. the file-descriptor table is
                // momentarily full) must not kill the server or spin the CPU.
                diag(&format!("accept failed: {e}"));
                std::thread::sleep(Duration::from_millis(50));
                continue;
            }
        };

        // Honour the connection cap before committing a thread to this peer. If
        // we are already at the limit, hang up immediately: leaving the socket
        // open without a handler would only make the client wait for a relay
        // that is never going to read from it.
        let permit = match &semaphore {
            None => None,
            Some(s) => match s.try_acquire() {
                Some(p) => Some(p),
                None => {
                    let _ = stream.shutdown(Shutdown::Both);
                    continue;
                }
            },
        };
        let config = Arc::clone(&config);
        let spawned = std::thread::Builder::new()
            .name("jokycrypt-conn".to_string())
            .spawn(move || {
                let peer = stream
                    .peer_addr()
                    .map(|a| a.to_string())
                    .unwrap_or_else(|_| "?".to_string());
                if let Err(e) = handle(&config, stream) {
                    diag(&format!("connection from {peer} ended: {e}"));
                }
                drop(permit);
            });
        if let Err(e) = spawned {
            // The dropped closure releases the permit and closes the socket.
            diag(&format!("could not spawn connection handler: {e}"));
        }
    }

    Ok(())
}

fn dial(target: &Addr, timeout: Option<Duration>) -> std::io::Result<TcpStream> {
    let addr = (target.host.as_str(), target.port)
        .to_socket_addrs()?
        .next()
        .ok_or_else(|| {
            std::io::Error::new(
                std::io::ErrorKind::NotFound,
                format!("could not resolve {target}"),
            )
        })?;
    match timeout {
        Some(t) => TcpStream::connect_timeout(&addr, t),
        None => TcpStream::connect(addr),
    }
}

/// Default ceiling on how long the handshake may take when no `--timeout` was
/// given. Deliberately generous: the handshake is a few round trips, and it has
/// to complete over genuinely bad links — overseas latency, hotel wifi, 3G with
/// packet loss and TCP retransmit backoff. Thirty seconds leaves ample headroom
/// there while still bounding how long a stalled or hostile peer can hold a slot.
const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30);

/// Configure the socket carrying the encrypted leg for the handshake — the hard
/// deadline and Nagle off (the handshake is a tiny ping-pong, the worst case for
/// Nagle plus delayed-ACK) — and return the deadline for the watchdog. Both roles
/// set it up identically, so that lives here rather than being repeated per role.
fn prepare_handshake_socket(sock: &TcpStream, config: &Config) -> std::io::Result<Duration> {
    let deadline = config.timeout.unwrap_or(HANDSHAKE_TIMEOUT);
    sock.set_read_timeout(Some(deadline))?;
    sock.set_write_timeout(Some(deadline))?;
    let _ = sock.set_nodelay(true);
    Ok(deadline)
}

/// Run a handshake closure under a hard wall-clock deadline.
///
/// Two layers of protection, because the handshake takes several round trips and
/// the server must never wedge on a misbehaving peer:
///
/// 1. Per-operation read and write timeouts (set by `prepare_handshake_socket`)
///    bound any single blocking call against a silent peer.
/// 2. A watchdog thread caps the *total* time. If the handshake has not finished
///    by the deadline it shuts the socket down, so any read or write in flight
///    returns immediately. This also defeats a peer that trickles one byte at a
///    time to keep resetting the per-operation timers.
fn run_handshake<R>(
    sock: &TcpStream,
    deadline: Duration,
    run: impl FnOnce(&mut &TcpStream) -> std::io::Result<R>,
) -> std::io::Result<R> {
    let watch = sock.try_clone()?;
    let (tx, rx) = std::sync::mpsc::channel::<()>();
    let watchdog = std::thread::Builder::new()
        .name("jokycrypt-hs-watchdog".to_string())
        .spawn(move || {
            if let Err(std::sync::mpsc::RecvTimeoutError::Timeout) = rx.recv_timeout(deadline) {
                let _ = watch.shutdown(Shutdown::Both);
            }
        });

    let mut io = sock;
    let result = run(&mut io);

    // Finished in time: drop the sender so the watchdog wakes at once and leaves
    // the (now active) socket alone.
    drop(tx);
    if let Ok(handle) = watchdog {
        let _ = handle.join();
    }
    result
}

fn handle(config: &Config, accepted: TcpStream) -> std::io::Result<()> {
    // The roles differ in which socket is the encrypted leg, when it is dialled,
    // and how the target is chosen: the server authenticates the peer, lets the
    // handshake settle on one of its allowed targets, then dials it; the client
    // dials its single upstream relay and authenticates that, naming the target
    // it wants to reach.
    match config.mode {
        Mode::Server => {
            let deadline = prepare_handshake_socket(&accepted, config)?;
            let (session, target) = run_handshake(&accepted, deadline, |io| {
                handshake::respond(io, &config.password, |host, port| {
                    std::str::from_utf8(host)
                        .ok()
                        .and_then(|h| config.select_target(h, port))
                        .cloned()
                })
            })?;
            let upstream = dial(&target, config.timeout)?;
            relay::run(
                accepted,
                upstream,
                session,
                CipherSide::A,
                config.verbose,
                config.timeout,
            );
        }
        Mode::Client => {
            let upstream = dial(&config.targets[0], config.timeout)?;
            let deadline = prepare_handshake_socket(&upstream, config)?;
            let (req_host, req_port) = match &config.requested_target {
                Some(addr) => (addr.host.as_bytes(), addr.port),
                None => (&b""[..], 0),
            };
            let session = run_handshake(&upstream, deadline, |io| {
                handshake::initiate(io, &config.password, req_host, req_port)
            })?;
            relay::run(
                accepted,
                upstream,
                session,
                CipherSide::B,
                config.verbose,
                config.timeout,
            );
        }
    }
    Ok(())
}
