//! A challenge handshake that proves, at the start of a connection, which key is at the
//! other end: for transports that don't authenticate their peers themselves (an onion
//! service learns nothing about who connects to it, `transport::tor`).
//!
//! Both sides send a HELLO: their public key and a fresh random nonce. Then each sends
//! a PROOF: its signature over a label, its role (dialer or listener), the connection's
//! context, and both HELLOs (its own first). Each checks the other's signature against
//! the key in the other's HELLO, and from then on knows that key is at the other end.
//!
//! * A recorded PROOF is useless elsewhere: it signs the other side's nonce, fresh for
//!   each connection.
//! * The roles keep a side's own PROOF from being reflected back to it.
//! * The context binds the handshake to the connection as the dialer meant it: the
//!   listener's address, as the dialer dialed it and as the listener knows itself. A
//!   party in the middle that passes one side's handshake on to someone else fails
//!   there: the context it was signed for is a different one.
//!
//! What it doesn't do: protect what follows. That is the connection's job, so the
//! handshake is only enough over one that is encrypted and tamper-proof end to end, and
//! authenticates the listener to the dialer (a Tor stream to an onion service). Over a
//! bare network (UDP, say), it would need key agreement and protected messages as well.
//!
//! The signing key is the node key (`crate::identity`); the label keeps these
//! signatures apart from the node key's others (WHO and ME).

use std::time::Duration;

use anyhow::{Result, anyhow, bail};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};

use crate::identity::{self, NodeKey};

const LABEL: &[u8] = b"continuum/challenge/1";
/// The first byte of a HELLO: this handshake's version.
const VERSION: u8 = 1;
const HELLO_BYTES: usize = 1 + 32 + 32;
const PROOF_BYTES: usize = 64;

/// Which end of the connection a side is.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum Role {
    Dialer,
    Listener,
}

impl Role {
    fn byte(self) -> u8 {
        match self {
            Self::Dialer => 0,
            Self::Listener => 1,
        }
    }

    fn other(self) -> Self {
        match self {
            Self::Dialer => Self::Listener,
            Self::Listener => Self::Dialer,
        }
    }
}

/// One side's HELLO.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct Hello {
    pub(crate) key: [u8; 32],
    pub(crate) nonce: [u8; 32],
}

impl Hello {
    /// Ours, with a fresh nonce.
    pub(crate) fn new(key: &NodeKey) -> Result<Self> {
        let mut nonce = [0u8; 32];
        getrandom::fill(&mut nonce).map_err(|error| anyhow!("No randomness: {error}"))?;
        Ok(Self {
            key: key.public(),
            nonce,
        })
    }

    pub(crate) fn to_bytes(self) -> [u8; HELLO_BYTES] {
        let mut bytes = [0u8; HELLO_BYTES];
        bytes[0] = VERSION;
        bytes[1..33].copy_from_slice(&self.key);
        bytes[33..].copy_from_slice(&self.nonce);
        bytes
    }

    pub(crate) fn from_bytes(bytes: &[u8; HELLO_BYTES]) -> Result<Self> {
        if bytes[0] != VERSION {
            bail!(
                "challenge handshake version {} (we speak {VERSION})",
                bytes[0]
            );
        }
        let mut hello = Self {
            key: [0; 32],
            nonce: [0; 32],
        };
        hello.key.copy_from_slice(&bytes[1..33]);
        hello.nonce.copy_from_slice(&bytes[33..]);
        Ok(hello)
    }
}

/// What the side in `role` signs: the label, its role, the context, then its own HELLO
/// and the other side's.
fn signed(role: Role, context: &str, own: &Hello, other: &Hello) -> Vec<u8> {
    let mut bytes = LABEL.to_vec();
    bytes.push(role.byte());
    bytes.extend_from_slice(&(context.len() as u32).to_be_bytes());
    bytes.extend_from_slice(context.as_bytes());
    for hello in [own, other] {
        bytes.extend_from_slice(&hello.key);
        bytes.extend_from_slice(&hello.nonce);
    }
    bytes
}

/// Our PROOF, as the side in `role`.
pub(crate) fn prove(
    key: &NodeKey,
    role: Role,
    context: &str,
    own: &Hello,
    other: &Hello,
) -> [u8; PROOF_BYTES] {
    key.sign(&signed(role, context, own, other))
}

/// Whether `proof` is the other side's (the one not in `role`) for this handshake.
pub(crate) fn verify(
    role: Role,
    context: &str,
    own: &Hello,
    other: &Hello,
    proof: &[u8; PROOF_BYTES],
) -> bool {
    identity::verify(
        &other.key,
        &signed(role.other(), context, other, own),
        proof,
    )
}

/// Runs the handshake on `stream` as the side in `role`, within `timeout`: returns the
/// public key at the other end. `context`: the listener's address, as dialed (module
/// docs).
pub(crate) async fn run<S>(
    stream: &mut S,
    key: &NodeKey,
    role: Role,
    context: &str,
    timeout: Duration,
) -> Result<[u8; 32]>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let exchange = async {
        let own = Hello::new(key)?;
        stream.write_all(&own.to_bytes()).await?;
        stream.flush().await?;
        let mut hello = [0u8; HELLO_BYTES];
        stream.read_exact(&mut hello).await?;
        let other = Hello::from_bytes(&hello)?;
        if other.key == own.key {
            bail!("the other end has our own key");
        }
        stream
            .write_all(&prove(key, role, context, &own, &other))
            .await?;
        stream.flush().await?;
        let mut proof = [0u8; PROOF_BYTES];
        stream.read_exact(&mut proof).await?;
        if !verify(role, context, &own, &other, &proof) {
            bail!("the other end's proof doesn't hold");
        }
        Ok(other.key)
    };
    tokio::time::timeout(timeout, exchange)
        .await
        .map_err(|_| anyhow!("no handshake within {timeout:?}"))?
}

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

    fn key(seed: u8) -> NodeKey {
        NodeKey::from_seed([seed; 32])
    }

    #[test]
    fn a_proof_holds_for_its_handshake_only() {
        let (a, b) = (key(1), key(2));
        let (ha, hb) = (Hello::new(&a).unwrap(), Hello::new(&b).unwrap());
        let context = "example.onion:9663";
        let proof = prove(&a, Role::Dialer, context, &ha, &hb);
        // b, the listener, checks a's proof.
        assert!(verify(Role::Listener, context, &hb, &ha, &proof));
        // Another context (passed on to another listener).
        assert!(!verify(
            Role::Listener,
            "other.onion:9663",
            &hb,
            &ha,
            &proof
        ));
        // Another connection: b's nonce is a new one.
        let hb2 = Hello::new(&b).unwrap();
        assert!(!verify(Role::Listener, context, &hb2, &ha, &proof));
        // Reflected: a can't be fooled by its own proof as if b had sent it.
        assert!(!verify(Role::Dialer, context, &ha, &hb, &proof));
        // Someone else's key in the HELLO.
        let hc = Hello {
            key: key(3).public(),
            nonce: ha.nonce,
        };
        assert!(!verify(Role::Listener, context, &hb, &hc, &proof));
    }

    #[test]
    fn hellos_round_trip_and_versions_are_checked() {
        let hello = Hello::new(&key(1)).unwrap();
        assert_eq!(Hello::from_bytes(&hello.to_bytes()).unwrap(), hello);
        let mut bytes = hello.to_bytes();
        bytes[0] = VERSION + 1;
        assert!(Hello::from_bytes(&bytes).is_err());
    }

    #[tokio::test]
    async fn both_ends_learn_each_others_key() {
        let (mut dialer, mut listener) = tokio::io::duplex(1024);
        let (a, b) = (key(1), key(2));
        let context = "example.onion:9663";
        let (at_a, at_b) = tokio::join!(
            run(
                &mut dialer,
                &a,
                Role::Dialer,
                context,
                Duration::from_secs(5)
            ),
            run(
                &mut listener,
                &b,
                Role::Listener,
                context,
                Duration::from_secs(5)
            ),
        );
        assert_eq!(at_a.unwrap(), b.public());
        assert_eq!(at_b.unwrap(), a.public());
    }

    #[tokio::test]
    async fn a_different_context_fails_on_both_ends() {
        let (mut dialer, mut listener) = tokio::io::duplex(1024);
        let (a, b) = (key(1), key(2));
        let (at_a, at_b) = tokio::join!(
            run(
                &mut dialer,
                &a,
                Role::Dialer,
                "meant.onion:1",
                Duration::from_secs(5)
            ),
            run(
                &mut listener,
                &b,
                Role::Listener,
                "actual.onion:1",
                Duration::from_secs(5)
            ),
        );
        assert!(at_a.is_err() && at_b.is_err());
    }
}