//! An implementation of Tor's current relay cell cryptography.
//!
//! These are not very good algorithms; they were the best we could come up with
//! in ~2002.  They are somewhat inefficient, and vulnerable to tagging attacks.
//! They should get replaced within the next several years.  For information on
//! some older proposed alternatives so far, see proposals 261, 295, and 298.
//!
//! I am calling this design `tor1`; it does not have a generally recognized
//! name.

use crate::{Error, Result, client::circuit::CircuitBinding, crypto::binding::CIRC_BINDING_LEN};

use cipher::{KeyIvInit, StreamCipher};
use digest::Digest;
use tor_cell::{chancell::ChanCmd, relaycell::msg::SendmeTag};
use tor_error::internal;
use typenum::Unsigned;

use super::{
    ClientLayer, CryptInit, InboundClientLayer, InboundRelayLayer, OutboundClientLayer,
    OutboundRelayLayer, RelayCellBody, RelayLayer,
};

/// Length of SENDME tag generated by this encryption method.
const SENDME_TAG_LEN: usize = 20;

/// A CryptState represents one layer of shared cryptographic state between
/// a relay and a client for a single hop, in a single direction.
///
/// For example, if a client makes a 3-hop circuit, then it will have 6
/// `CryptState`s, one for each relay, for each direction of communication.
///
/// Note that although `CryptState` is used to implement [`OutboundClientLayer`],
/// [`InboundClientLayer`], [`OutboundRelayLayer`], and [`InboundRelayLayer`],
/// each instance will only be used for one of these roles.
///
/// It is parameterized on a stream cipher and a digest type: most circuits
/// will use AES-128-CTR and SHA1, but v3 onion services use AES-256-CTR and
/// SHA-3.
struct CryptState<SC: StreamCipher, D: Digest + Clone> {
    /// Stream cipher for en/decrypting cell bodies.
    ///
    /// This cipher is the one keyed with Kf or Kb in the spec.
    cipher: SC,
    /// Digest for authenticating cells to/from this hop.
    ///
    /// This digest is the one keyed with Df or Db in the spec.
    digest: D,
    /// Most recent digest value generated by this crypto.
    last_sendme_tag: SendmeTag,
}

/// A pair of CryptStates shared between a client and a relay, one for the
/// outbound (away from the client) direction, and one for the inbound
/// (towards the client) direction.
#[cfg_attr(feature = "bench", visibility::make(pub))]
pub(crate) struct CryptStatePair<SC: StreamCipher, D: Digest + Clone> {
    /// State for en/decrypting cells sent away from the client.
    fwd: CryptState<SC, D>,
    /// State for en/decrypting cells sent towards the client.
    back: CryptState<SC, D>,
    /// A circuit binding key.
    binding: CircuitBinding,
}

impl<SC: StreamCipher + KeyIvInit, D: Digest + Clone> CryptInit for CryptStatePair<SC, D> {
    fn seed_len() -> usize {
        SC::KeySize::to_usize() * 2 + D::OutputSize::to_usize() * 2 + CIRC_BINDING_LEN
    }
    fn initialize(mut seed: &[u8]) -> Result<Self> {
        // This corresponds to the use of the KDF algorithm as described in
        // tor-spec 5.2.2
        if seed.len() != Self::seed_len() {
            return Err(Error::from(internal!(
                "seed length {} was invalid",
                seed.len()
            )));
        }

        // Advances `seed` by `n` bytes, returning the advanced bytes
        let mut take_seed = |n: usize| -> &[u8] {
            let res = &seed[..n];
            seed = &seed[n..];
            res
        };

        let dlen = D::OutputSize::to_usize();
        let keylen = SC::KeySize::to_usize();

        let df = take_seed(dlen);
        let db = take_seed(dlen);
        let kf = take_seed(keylen);
        let kb = take_seed(keylen);
        let binding_key = take_seed(CIRC_BINDING_LEN);

        let fwd = CryptState {
            cipher: SC::new(
                kf.try_into().expect("Incorrect size, despite validation!"),
                &Default::default(),
            ),
            digest: D::new().chain_update(df),
            last_sendme_tag: [0_u8; SENDME_TAG_LEN].into(),
        };
        let back = CryptState {
            cipher: SC::new(
                kb.try_into().expect("Incorrect size, despite validation!"),
                &Default::default(),
            ),
            digest: D::new().chain_update(db),
            last_sendme_tag: [0_u8; SENDME_TAG_LEN].into(),
        };
        let binding = CircuitBinding::try_from(binding_key)?;

        Ok(CryptStatePair { fwd, back, binding })
    }
}

impl<SC, D> ClientLayer<ClientOutbound<SC, D>, ClientInbound<SC, D>> for CryptStatePair<SC, D>
where
    SC: StreamCipher,
    D: Digest + Clone,
{
    fn split_client_layer(self) -> (ClientOutbound<SC, D>, ClientInbound<SC, D>, CircuitBinding) {
        (self.fwd.into(), self.back.into(), self.binding)
    }
}

/// An inbound relay layer, encrypting relay cells for a client.
#[cfg_attr(feature = "bench", visibility::make(pub))]
#[derive(derive_more::From)]
pub(crate) struct RelayInbound<SC: StreamCipher, D: Digest + Clone>(CryptState<SC, D>);
impl<SC: StreamCipher, D: Digest + Clone> InboundRelayLayer for RelayInbound<SC, D> {
    fn originate(&mut self, cmd: ChanCmd, cell: &mut RelayCellBody) -> SendmeTag {
        cell.set_digest::<_>(&mut self.0.digest, &mut self.0.last_sendme_tag);
        self.encrypt_inbound(cmd, cell);
        self.0.last_sendme_tag
    }
    fn encrypt_inbound(&mut self, _cmd: ChanCmd, cell: &mut RelayCellBody) {
        // This is describe in tor-spec 5.5.3.1, "Relaying Backward at Onion Routers"
        self.0.cipher.apply_keystream(cell.as_mut());
    }
}

/// An outbound relay layer, decrypting relay cells from a client.
#[cfg_attr(feature = "bench", visibility::make(pub))]
#[derive(derive_more::From)]
pub(crate) struct RelayOutbound<SC: StreamCipher, D: Digest + Clone>(CryptState<SC, D>);
impl<SC: StreamCipher, D: Digest + Clone> OutboundRelayLayer for RelayOutbound<SC, D> {
    fn decrypt_outbound(&mut self, _cmd: ChanCmd, cell: &mut RelayCellBody) -> Option<SendmeTag> {
        // This is describe in tor-spec 5.5.2.2, "Relaying Forward at Onion Routers"
        self.0.cipher.apply_keystream(cell.as_mut());
        if cell.is_recognized::<_>(&mut self.0.digest, &mut self.0.last_sendme_tag) {
            Some(self.0.last_sendme_tag)
        } else {
            None
        }
    }
}
impl<SC: StreamCipher, D: Digest + Clone> RelayLayer<RelayOutbound<SC, D>, RelayInbound<SC, D>>
    for CryptStatePair<SC, D>
{
    fn split_relay_layer(self) -> (RelayOutbound<SC, D>, RelayInbound<SC, D>, CircuitBinding) {
        let CryptStatePair { fwd, back, binding } = self;
        (fwd.into(), back.into(), binding)
    }
}

/// An outbound client layer, encrypting relay cells for a relay.
#[cfg_attr(feature = "bench", visibility::make(pub))]
#[derive(derive_more::From)]
pub(crate) struct ClientOutbound<SC: StreamCipher, D: Digest + Clone>(CryptState<SC, D>);

impl<SC: StreamCipher, D: Digest + Clone> OutboundClientLayer for ClientOutbound<SC, D> {
    fn originate_for(&mut self, cmd: ChanCmd, cell: &mut RelayCellBody) -> SendmeTag {
        cell.set_digest::<_>(&mut self.0.digest, &mut self.0.last_sendme_tag);
        self.encrypt_outbound(cmd, cell);
        self.0.last_sendme_tag
    }
    fn encrypt_outbound(&mut self, _cmd: ChanCmd, cell: &mut RelayCellBody) {
        // This is a single iteration of the loop described in tor-spec
        // 5.5.2.1, "routing away from the origin."
        self.0.cipher.apply_keystream(&mut cell.0[..]);
    }
}

/// An outbound client layer, decryption relay cells from a relay.
#[cfg_attr(feature = "bench", visibility::make(pub))]
#[derive(derive_more::From)]
pub(crate) struct ClientInbound<SC: StreamCipher, D: Digest + Clone>(CryptState<SC, D>);
impl<SC: StreamCipher, D: Digest + Clone> InboundClientLayer for ClientInbound<SC, D> {
    fn decrypt_inbound(&mut self, _cmd: ChanCmd, cell: &mut RelayCellBody) -> Option<SendmeTag> {
        // This is a single iteration of the loop described in tor-spec
        // 5.5.3, "routing to the origin."
        self.0.cipher.apply_keystream(&mut cell.0[..]);
        if cell.is_recognized::<_>(&mut self.0.digest, &mut self.0.last_sendme_tag) {
            Some(self.0.last_sendme_tag)
        } else {
            None
        }
    }
}

/// Location in the relay cell for our "recognized" field.
pub(super) const RECOGNIZED_RANGE: std::ops::Range<usize> = 1..3;
/// Location in the relay cell for our "Digest" field.
pub(super) const DIGEST_RANGE: std::ops::Range<usize> = 5..9;
/// An all-zero digest value.
pub(super) const EMPTY_DIGEST: &[u8] = &[0, 0, 0, 0];

/// Functions on RelayCellBody that implement the digest/recognized
/// algorithm.
///
/// The current relay crypto protocol uses two wholly inadequate fields to
/// see whether a cell is intended for its current recipient: a two-byte
/// "recognized" field that needs to be all-zero; and a four-byte "digest"
/// field containing a running digest of all cells (for this recipient) to
/// this one, seeded with an initial value (either Df or Db in the spec).
///
/// These operations are described in tor-spec section 6.1 "Relay cells"
//
// TODO: It may be that we should un-parameterize the functions
// that use RCF: given our timeline for deployment of CGO encryption,
// it is likely that we will never actually want to  support `tor1` encryption
// with any other format than RelayCellFormat::V0.
impl RelayCellBody {
    /// Returns the byte slice of the `recognized` field.
    fn recognized(&self) -> &[u8] {
        &self.0[RECOGNIZED_RANGE]
    }
    /// Returns the mut byte slice of the `recognized` field.
    fn recognized_mut(&mut self) -> &mut [u8] {
        &mut self.0[RECOGNIZED_RANGE]
    }
    /// Returns the byte slice of the `digest` field.
    fn digest(&self) -> &[u8] {
        &self.0[DIGEST_RANGE]
    }
    /// Returns the mut byte slice of the `digest` field.
    fn digest_mut(&mut self) -> &mut [u8] {
        &mut self.0[DIGEST_RANGE]
    }
    /// Prepare a cell body by setting its digest and recognized field.
    #[cfg_attr(feature = "bench", visibility::make(pub))]
    fn set_digest<D: Digest + Clone>(&mut self, d: &mut D, sendme_tag: &mut SendmeTag) {
        self.recognized_mut().fill(0); // Set 'Recognized' to zero
        self.digest_mut().fill(0); // Set Digest to zero

        d.update(&self.0[..]);
        // TODO(nickm) can we avoid this clone?  Probably not.
        let computed_digest = d.clone().finalize();
        // TODO PERF: Make sure this compiles nicely.
        *sendme_tag = SendmeTag::try_from(&computed_digest[..SENDME_TAG_LEN])
            .expect("Somehow produced a SENDME tag of invalid length!");
        let used_digest_prefix = &computed_digest[0..DIGEST_RANGE.len()];
        self.digest_mut().copy_from_slice(used_digest_prefix);
    }
    /// Check whether this just-decrypted cell is now an authenticated plaintext.
    ///
    /// This method returns true if the `recognized` field is all zeros, and if the
    /// `digest` field is a digest of the correct material.
    /// If it returns true, it also sets `rcvg` to the appropriate authenticated
    /// SENDME tag to use if acknowledging this message.
    ///
    /// If this method returns false, then either further decryption is required,
    /// or the cell is corrupt.
    ///
    // TODO #1336: Further optimize and/or benchmark this.
    #[cfg_attr(feature = "bench", visibility::make(pub))]
    fn is_recognized<D: Digest + Clone>(&self, d: &mut D, rcvd: &mut SendmeTag) -> bool {
        use crate::util::ct;

        // Validate 'Recognized' field
        if !ct::is_zero(self.recognized()) {
            return false;
        }

        // Now also validate the 'Digest' field:

        let mut dtmp = d.clone();
        // Add bytes up to the 'Digest' field
        dtmp.update(&self.0[..DIGEST_RANGE.start]);
        // Add zeroes where the 'Digest' field is
        dtmp.update(EMPTY_DIGEST);
        // Add the rest of the bytes
        dtmp.update(&self.0[DIGEST_RANGE.end..]);
        // Clone the digest before finalize destroys it because we will use
        // it in the future
        let dtmp_clone = dtmp.clone();
        let result = dtmp.finalize();

        if ct::bytes_eq(self.digest(), &result[0..DIGEST_RANGE.len()]) {
            // Copy useful things out of this cell (we keep running digest)
            *d = dtmp_clone;
            *rcvd = SendmeTag::try_from(&result[..SENDME_TAG_LEN])
                .expect("Somehow generated a sendme tag of invalid length!");
            return true;
        }

        false
    }
}

/// Benchmark utilities for the `tor1` module.
#[cfg(feature = "bench")]
pub mod bench_utils {
    pub use super::ClientInbound;
    pub use super::ClientOutbound;
    pub use super::CryptStatePair;
    pub use super::RelayInbound;
    pub use super::RelayOutbound;

    /// The throughput for a relay cell in bytes with the Tor1 scheme.
    pub const TOR1_THROUGHPUT: u64 = 498;
}

#[cfg(test)]
mod test {
    // @@ begin test lint list maintained by maint/add_warning @@
    #![allow(clippy::bool_assert_comparison)]
    #![allow(clippy::clone_on_copy)]
    #![allow(clippy::dbg_macro)]
    #![allow(clippy::mixed_attributes_style)]
    #![allow(clippy::print_stderr)]
    #![allow(clippy::print_stdout)]
    #![allow(clippy::single_char_pattern)]
    #![allow(clippy::unwrap_used)]
    #![allow(clippy::unchecked_time_subtraction)]
    #![allow(clippy::useless_vec)]
    #![allow(clippy::needless_pass_by_value)]
    #![allow(clippy::string_slice)] // See arti#2571
    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->

    use crate::crypto::cell::{
        InboundClientCrypt, OutboundClientCrypt, Tor1RelayCrypto, test::add_layers,
    };

    use super::*;

    // From tor's test_relaycrypt.c

    #[test]
    fn testvec() {
        use digest::XofReader;
        use digest::{ExtendableOutput, Update};

        // (The ....s at the end here are the KH ca)
        const K1: &[u8; 92] =
            b"    'My public key is in this signed x509 object', said Tom assertively.      (N-PREG-VIRYL)";
        const K2: &[u8; 92] =
            b"'Let's chart the pedal phlanges in the tomb', said Tom cryptographically.  (PELCG-GBR-TENCU)";
        const K3: &[u8; 92] =
            b"     'Segmentation fault bugs don't _just happen_', said Tom seethingly.        (P-GUVAT-YL)";

        const SEED: &[u8;108] = b"'You mean to tell me that there's a version of Sha-3 with no limit on the output length?', said Tom shakily.";
        let cmd = ChanCmd::RELAY;

        // These test vectors were generated from Tor.
        let data: &[(usize, &str)] = &include!("../../../testdata/cell_crypt.rs");

        let mut cc_out = OutboundClientCrypt::new();
        let mut cc_in = InboundClientCrypt::new();
        let pair = Tor1RelayCrypto::initialize(&K1[..]).unwrap();
        add_layers(&mut cc_out, &mut cc_in, pair);
        let pair = Tor1RelayCrypto::initialize(&K2[..]).unwrap();
        add_layers(&mut cc_out, &mut cc_in, pair);
        let pair = Tor1RelayCrypto::initialize(&K3[..]).unwrap();
        add_layers(&mut cc_out, &mut cc_in, pair);

        let mut xof = tor_llcrypto::d::Shake256::default();
        xof.update(&SEED[..]);
        let mut stream = xof.finalize_xof();

        let mut j = 0;
        for cellno in 0..51 {
            let mut body = Box::new([0_u8; 509]);
            body[0] = 2; // command: data.
            body[4] = 1; // streamid: 1.
            body[9] = 1; // length: 498
            body[10] = 242;
            stream.read(&mut body[11..]);

            let mut cell = body.into();
            let _ = cc_out.encrypt(cmd, &mut cell, 2.into());

            if cellno == data[j].0 {
                let expected = hex::decode(data[j].1).unwrap();
                assert_eq!(cell.as_ref(), &expected[..]);
                j += 1;
            }
        }
    }
}