//! A tiny mutual challenge-response handshake.
//!
//! The server speaks first and issues a fresh random challenge. The client never
//! sends anything derived from the password until it has that challenge, and its
//! reply is bound to it — so the proof cannot be precomputed, and a reply
//! captured from one connection is useless on the next (every challenge is new).
//! The client folds in its own nonce too, which lets it challenge the server
//! right back, so both ends authenticate and both contribute to the session key.
//!
//! The client also names, in that same authenticated reply, which of the
//! server's allowed targets it wants to reach. Binding the target into the proof
//! means an eavesdropper cannot redirect a connection from one allowed target to
//! another by flipping bytes on the wire. The server answers with a status byte
//! that says whether the requested target is permitted, alongside its own proof,
//! so the client can tell an honest refusal apart from an impostor.
//!
//! The responses are computed with the same keyed table as the cipher, so this
//! pulls in no real cryptography. An eavesdropper can still mount an offline
//! dictionary attack against a captured exchange; that is acceptable for a
//! teaching tool whose point is to be fast and transparent rather than safe.
//!
//! The wire flow (challenges, nonces, and the requested target travel in the
//! clear; the password never does):
//!
//! ```text
//!   server -> client:  MAGIC || server_challenge
//!   client -> server:  client_nonce || host_len || host || port || response(client)
//!   server -> client:  status || response(server)
//! ```

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

use crate::cipher::{KeyedTable, StreamCipher};

const MAGIC: &[u8; 4] = b"JKR2";
const NONCE_LEN: usize = 8;
const TAG_LEN: usize = 16;

/// Handshake verdict the server sends in its final message.
const STATUS_OK: u8 = 0;
const STATUS_TARGET_REFUSED: u8 = 1;

/// The two keystreams for a connection, one per logical direction.
///
/// `c2s` carries data travelling from the client end toward the server end;
/// `s2c` carries the replies. Both endpoints derive the identical pair, so
/// whatever one side encrypts with `c2s` the other decrypts with `c2s`.
pub struct Session {
    pub c2s: StreamCipher,
    pub s2c: StreamCipher,
}

fn random_nonce() -> io::Result<[u8; NONCE_LEN]> {
    let mut buf = [0u8; NONCE_LEN];
    let mut f = std::fs::File::open("/dev/urandom")?;
    f.read_exact(&mut buf)?;
    Ok(buf)
}

/// The proof one side sends: a keyed tag over the password, both challenges, and
/// the requested target. Binding to the server's challenge is what makes it a
/// response rather than a reusable token; binding the target stops it being moved
/// to a different allowed destination; the `label` keeps the client's and
/// server's proofs distinct so neither can be reflected back at the other. The
/// host length is mixed in too, so the host and port fields can never be
/// reinterpreted across their boundary.
fn response(
    password: &[u8],
    server_challenge: &[u8],
    client_nonce: &[u8],
    req_host: &[u8],
    req_port: u16,
    label: &[u8],
) -> [u8; TAG_LEN] {
    let host_len = [req_host.len() as u8];
    let port = req_port.to_be_bytes();
    let material = [
        password,
        server_challenge,
        client_nonce,
        &host_len,
        req_host,
        &port,
        label,
    ]
    .concat();
    let mut out = [0u8; TAG_LEN];
    KeyedTable::new(&material).fill(&mut out);
    out
}

fn derive_stream(password: &[u8], nonce: &[u8], label: u8) -> StreamCipher {
    let key = [password, nonce, &[label]].concat();
    let mut seed = [0u8; 32];
    KeyedTable::new(&key).fill(&mut seed);
    StreamCipher::from_seed(seed)
}

fn session_from_nonces(password: &[u8], client_nonce: &[u8], server_challenge: &[u8]) -> Session {
    let mut nonce = [0u8; NONCE_LEN * 2];
    nonce[..NONCE_LEN].copy_from_slice(client_nonce);
    nonce[NONCE_LEN..].copy_from_slice(server_challenge);
    Session {
        c2s: derive_stream(password, &nonce, 0x01),
        s2c: derive_stream(password, &nonce, 0x02),
    }
}

fn auth_failed(detail: &str) -> io::Error {
    io::Error::new(io::ErrorKind::PermissionDenied, detail.to_string())
}

/// The server proved itself but will not reach the target we asked for. This is
/// not an authentication failure — the password was right — so it gets its own
/// error kind and a message naming the target.
fn target_refused(host: &[u8], port: u16) -> io::Error {
    let host = String::from_utf8_lossy(host);
    io::Error::new(
        io::ErrorKind::ConnectionRefused,
        format!("server does not allow target {host}:{port}"),
    )
}

const CLIENT_LABEL: &[u8] = b"client";
const SERVER_LABEL: &[u8] = b"server";

/// Client side: read the server's challenge, answer it (bound to a fresh client
/// nonce and to the target we want), then check the server answers our nonce in
/// return and permits that target.
///
/// `req_host`/`req_port` name the target to reach. An empty host with port 0
/// means "no preference" — the server falls back to its first configured target.
pub fn initiate<S: Read + Write>(
    stream: &mut S,
    password: &[u8],
    req_host: &[u8],
    req_port: u16,
) -> io::Result<Session> {
    if req_host.len() > u8::MAX as usize {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "requested target host is too long for the handshake",
        ));
    }

    let mut greeting = [0u8; 4 + NONCE_LEN];
    stream.read_exact(&mut greeting)?;
    if &greeting[..4] != MAGIC {
        return Err(auth_failed("bad protocol magic from server"));
    }
    let server_challenge = &greeting[4..];

    let client_nonce = random_nonce()?;
    let resp_c = response(
        password,
        server_challenge,
        &client_nonce,
        req_host,
        req_port,
        CLIENT_LABEL,
    );

    let mut reply = Vec::with_capacity(NONCE_LEN + 1 + req_host.len() + 2 + TAG_LEN);
    reply.extend_from_slice(&client_nonce);
    reply.push(req_host.len() as u8);
    reply.extend_from_slice(req_host);
    reply.extend_from_slice(&req_port.to_be_bytes());
    reply.extend_from_slice(&resp_c);
    stream.write_all(&reply)?;
    stream.flush()?;

    let mut status = [0u8; 1];
    stream.read_exact(&mut status)?;
    let mut resp_s = [0u8; TAG_LEN];
    stream.read_exact(&mut resp_s)?;

    let expected = response(
        password,
        server_challenge,
        &client_nonce,
        req_host,
        req_port,
        SERVER_LABEL,
    );
    if resp_s != expected {
        return Err(auth_failed("server failed the challenge"));
    }

    // Only now that the server has proven itself do we trust its verdict on the
    // target we asked for.
    match status[0] {
        STATUS_OK => Ok(session_from_nonces(
            password,
            &client_nonce,
            server_challenge,
        )),
        STATUS_TARGET_REFUSED => Err(target_refused(req_host, req_port)),
        other => Err(auth_failed(&format!(
            "server sent an unknown handshake status: {other}"
        ))),
    }
}

/// Server side: issue a fresh challenge, verify the client's bound response,
/// resolve the target it asked for, then prove our own knowledge of the password
/// and report whether that target is allowed.
///
/// `select` is handed the requested host (as raw bytes off the wire) and port,
/// and returns the chosen target when it is permitted or `None` to refuse. The
/// session and the selected target are returned together on success.
pub fn respond<S: Read + Write, T>(
    stream: &mut S,
    password: &[u8],
    select: impl FnOnce(&[u8], u16) -> Option<T>,
) -> io::Result<(Session, T)> {
    let server_challenge = random_nonce()?;

    let mut greeting = Vec::with_capacity(4 + NONCE_LEN);
    greeting.extend_from_slice(MAGIC);
    greeting.extend_from_slice(&server_challenge);
    stream.write_all(&greeting)?;
    stream.flush()?;

    let mut client_nonce = [0u8; NONCE_LEN];
    stream.read_exact(&mut client_nonce)?;

    let mut host_len = [0u8; 1];
    stream.read_exact(&mut host_len)?;
    let mut req_host = vec![0u8; host_len[0] as usize];
    stream.read_exact(&mut req_host)?;

    let mut port_bytes = [0u8; 2];
    stream.read_exact(&mut port_bytes)?;
    let req_port = u16::from_be_bytes(port_bytes);

    let mut resp_c = [0u8; TAG_LEN];
    stream.read_exact(&mut resp_c)?;

    let expected = response(
        password,
        &server_challenge,
        &client_nonce,
        &req_host,
        req_port,
        CLIENT_LABEL,
    );
    if resp_c != expected {
        return Err(auth_failed("client failed the challenge"));
    }

    // The client has authenticated. Decide whether the target it named is one we
    // will reach; the verdict gives us both the status byte to send and the
    // outcome to return, so it is expressed once here.
    let resp_s = response(
        password,
        &server_challenge,
        &client_nonce,
        &req_host,
        req_port,
        SERVER_LABEL,
    );
    let (status, outcome) = match select(&req_host, req_port) {
        Some(target) => (
            STATUS_OK,
            Ok((
                session_from_nonces(password, &client_nonce, &server_challenge),
                target,
            )),
        ),
        None => (
            STATUS_TARGET_REFUSED,
            Err(target_refused(&req_host, req_port)),
        ),
    };

    // Answer either way, proving ourselves so an honest refusal is
    // distinguishable from an impostor. One write keeps it to a single segment.
    let mut tail = Vec::with_capacity(1 + TAG_LEN);
    tail.push(status);
    tail.extend_from_slice(&resp_s);
    stream.write_all(&tail)?;
    stream.flush()?;

    outcome
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;

    /// Run both halves of the handshake against in-memory pipes and check the
    /// derived streams actually cancel out.
    struct Pipe {
        to_peer: std::sync::mpsc::Sender<u8>,
        from_peer: std::sync::mpsc::Receiver<u8>,
    }

    impl Read for Pipe {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            for (n, slot) in buf.iter_mut().enumerate() {
                match self.from_peer.recv() {
                    Ok(b) => *slot = b,
                    Err(_) if n > 0 => return Ok(n),
                    Err(_) => return Ok(0),
                }
            }
            Ok(buf.len())
        }
    }

    impl Write for Pipe {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            for &b in buf {
                let _ = self.to_peer.send(b);
            }
            Ok(buf.len())
        }
        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    fn pipe_pair() -> (Pipe, Pipe) {
        let (ctx, srx) = std::sync::mpsc::channel();
        let (stx, crx) = std::sync::mpsc::channel();
        (
            Pipe {
                to_peer: ctx,
                from_peer: crx,
            },
            Pipe {
                to_peer: stx,
                from_peer: srx,
            },
        )
    }

    #[test]
    fn matching_passwords_agree_on_streams() {
        let (mut client, mut server) = pipe_pair();

        let srv = std::thread::spawn(move || respond(&mut server, b"swordfish", |_h, _p| Some(())));
        let mut cs = initiate(&mut client, b"swordfish", b"", 0).unwrap();
        let (mut ss, ()) = srv.join().unwrap().unwrap();

        let msg = b"ping across the tunnel".to_vec();
        let mut wire = msg.clone();
        cs.c2s.apply(&mut wire);
        ss.c2s.apply(&mut wire);
        assert_eq!(wire, msg, "c2s streams must match");
    }

    #[test]
    fn wrong_password_is_rejected() {
        let (mut client, mut server) = pipe_pair();

        let srv =
            std::thread::spawn(move || respond(&mut server, b"the-real-secret", |_h, _p| Some(())));
        let client_result = initiate(&mut client, b"a-guess", b"", 0);
        let server_result = srv.join().unwrap();

        assert!(
            server_result.is_err(),
            "server should reject the wrong password"
        );
        // The client is left hanging once the server bails, so its read of the
        // (never-sent) server proof must also fail rather than succeed.
        assert!(client_result.is_err(), "client must not believe it authed");
    }

    #[test]
    fn client_rejects_bad_server_magic() {
        // The client reads (and validates) the greeting before sending anything,
        // so a read-only stream is enough to exercise the rejection path.
        let mut bogus = Cursor::new(b"XXXXnot-a-challenge".to_vec());
        assert!(initiate(&mut bogus, b"pw", b"", 0).is_err());
    }

    #[test]
    fn server_sees_the_requested_target() {
        let (mut client, mut server) = pipe_pair();

        let srv = std::thread::spawn(move || {
            respond(&mut server, b"pw", |host, port| Some((host.to_vec(), port)))
        });
        let _cs = initiate(&mut client, b"pw", b"192.168.1.50", 8888).unwrap();
        let (_ss, (host, port)) = srv.join().unwrap().unwrap();

        assert_eq!(host, b"192.168.1.50");
        assert_eq!(port, 8888);
    }

    #[test]
    fn refused_target_fails_both_ends() {
        let (mut client, mut server) = pipe_pair();

        // The password matches, but the server refuses the requested target.
        let srv = std::thread::spawn(move || respond(&mut server, b"pw", |_h, _p| None::<()>));
        let client_result = initiate(&mut client, b"pw", b"127.0.0.1", 22);
        let server_result = srv.join().unwrap();

        assert!(server_result.is_err(), "server refuses the target");
        assert!(
            client_result.is_err(),
            "client learns the target was refused"
        );
    }
}
