//! Bidirectional pumping between the two ends of a connection.
//!
//! Each connection has a "cipher side" (the encrypted network leg) and a "plain
//! side" (the local application or service). One thread pumps each direction.
//! The cipher is a stream, so each direction owns exactly one keystream and is
//! handled by exactly one thread — that keeps the keystream position glued to
//! the byte sequence.
//!
//! Memory per direction is a single fixed buffer, never a growing queue: we read
//! up to `BUF_SIZE`, then `write_all` it before reading more. `write_all` blocks
//! when the far side is slow, which stops us reading, which lets the source's TCP
//! window close and throttles the sender. So no peer can make the relay buffer an
//! unbounded stream — a flood is pushed back at the wire, not soaked up in RAM.

use std::io::{Read, Write};
use std::net::{Shutdown, TcpStream};
use std::sync::Mutex;
use std::time::Duration;

use crate::cipher::StreamCipher;
use crate::handshake::Session;

static PRINT_LOCK: Mutex<()> = Mutex::new(());

const BUF_SIZE: usize = 32 * 1024;

/// Which physical socket carries ciphertext.
#[derive(Clone, Copy)]
pub enum CipherSide {
    /// Server role: the accepted connection (`a`) is encrypted, the dialled
    /// target (`b`) is plaintext.
    A,
    /// Client role: the accepted connection (`a`) is plaintext, the dialled
    /// upstream (`b`) is encrypted.
    B,
}

/// Run the relay until either side closes. `a` is always the locally accepted
/// socket; `b` is the socket this process dialled.
pub fn run(
    a: TcpStream,
    b: TcpStream,
    session: Session,
    cipher_side: CipherSide,
    verbose: bool,
    timeout: Option<Duration>,
) {
    let Session { c2s, s2c } = session;

    for s in [&a, &b] {
        let _ = s.set_read_timeout(timeout);
        let _ = s.set_write_timeout(timeout);
        let _ = s.set_nodelay(true);
    }

    // a -> b carries client-to-server traffic, b -> a the replies. Whether the
    // bytes we read are already ciphertext depends on which side is encrypted.
    let (a_to_b_incoming_cipher, b_to_a_incoming_cipher) = match cipher_side {
        CipherSide::A => (true, false),
        CipherSide::B => (false, true),
    };

    let a_read = a;
    let b_write = match b.try_clone() {
        Ok(s) => s,
        Err(_) => return,
    };
    let b_read = b;
    let a_write = match a_read.try_clone() {
        Ok(s) => s,
        Err(_) => return,
    };

    let forward = std::thread::Builder::new()
        .name("jokycrypt-forward".to_string())
        .spawn(move || {
            pump(a_read, b_write, c2s, a_to_b_incoming_cipher, "c2s", verbose);
        });

    let forward = match forward {
        Ok(handle) => handle,
        Err(_) => {
            // Out of threads: tear down cleanly rather than run a half-duplex
            // relay that would silently stall one direction. The dropped closure
            // closes the other two socket handles.
            let _ = a_write.shutdown(Shutdown::Both);
            let _ = b_read.shutdown(Shutdown::Both);
            return;
        }
    };

    pump(b_read, a_write, s2c, b_to_a_incoming_cipher, "s2c", verbose);
    let _ = forward.join();
}

fn pump(
    mut src: TcpStream,
    mut dst: TcpStream,
    mut cipher: StreamCipher,
    incoming_is_cipher: bool,
    label: &str,
    verbose: bool,
) {
    let mut buf = vec![0u8; BUF_SIZE];
    loop {
        let n = match src.read(&mut buf) {
            Ok(0) => break,
            Ok(n) => n,
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
            Err(_) => break,
        };
        let chunk = &mut buf[..n];

        // We always log plaintext. On the encrypted leg that means logging after
        // decryption; on the plaintext leg, before encryption.
        if !incoming_is_cipher && verbose {
            log_traffic(label, chunk);
        }
        cipher.apply(chunk);
        if incoming_is_cipher && verbose {
            log_traffic(label, chunk);
        }

        if dst.write_all(chunk).is_err() {
            break;
        }
    }
    let _ = src.shutdown(Shutdown::Both);
    let _ = dst.shutdown(Shutdown::Both);
}

fn log_traffic(label: &str, data: &[u8]) {
    let rendered = render(data);
    let _guard = PRINT_LOCK.lock().unwrap_or_else(|e| e.into_inner());
    // Write straight to the handle and ignore failures: a closed stderr (for
    // example a piped reader that exited) must not panic a relay thread.
    let mut err = std::io::stderr();
    let _ = writeln!(err, "[{label}] {} bytes: {rendered}", data.len());
}

/// Render a chunk for the verbose log. Printable ASCII passes through verbatim,
/// control characters (tabs, newlines, and the like) collapse to a single space
/// so the line stays on one line, and any other byte is dropped entirely.
fn render(data: &[u8]) -> String {
    let mut rendered = String::with_capacity(data.len());
    for &byte in data {
        match byte {
            0x20..=0x7e => rendered.push(byte as char),
            _ if byte.is_ascii_control() => rendered.push(' '),
            _ => {}
        }
    }
    rendered
}

#[cfg(test)]
mod tests {
    use super::render;

    #[test]
    fn printable_passes_through() {
        assert_eq!(render(b"GET / HTTP/1.1"), "GET / HTTP/1.1");
    }

    #[test]
    fn control_chars_become_spaces() {
        assert_eq!(render(b"a\tb\r\nc"), "a b  c");
    }

    #[test]
    fn non_printable_bytes_are_dropped() {
        assert_eq!(render(&[0x80, b'h', 0xff, b'i', 0x7f]), "hi ");
    }
}
