clatter 0.1.0-alpha

use zeroize::{Zeroize, ZeroizeOnDrop};

use crate::bytearray::ByteArray;
use crate::error::{CipherError, CipherResult};
use crate::traits::{Cipher, CryptoComponent};

/// Pair of [`CipherState`] instances for sending and receiving transport messages
pub struct CipherStates<C: Cipher> {
    pub initiator_to_responder: CipherState<C>,
    pub responder_to_initiator: CipherState<C>,
}

/// Cipherstate for encrypting and decrypting messages
#[derive(ZeroizeOnDrop, Zeroize)]
pub struct CipherState<C: Cipher> {
    k: C::Key,
    n: u64,
    overflowed: bool,
}

impl<C: Cipher> CryptoComponent for CipherState<C> {
    fn name() -> &'static str {
        C::name()
    }
}

impl<C: Cipher> CipherState<C> {
    /// Initialize cipherstate with given key and nonce
    ///
    /// # Panics
    /// Panics if key data has incorrect length
    pub fn new(k: &[u8], n: u64) -> Self {
        Self {
            k: C::Key::from_slice(k),
            n,
            overflowed: false,
        }
    }

    fn nonce_inc_check(&mut self) {
        // "If incrementing n results in 2^(64)-1, then any further EncryptWithAd()
        // or DecryptWithAd() calls will signal an error to the caller"
        if self.n.checked_add(1).is_none() {
            self.overflowed = true;
        }
    }

    /// AEAD encryption
    pub fn encrypt_with_ad(
        &mut self,
        ad: &[u8],
        plaintext: &[u8],
        out: &mut [u8],
    ) -> CipherResult<()> {
        if self.overflowed {
            return Err(CipherError::NonceOverflow);
        }

        C::encrypt(&self.k, self.n, ad, plaintext, out);
        self.nonce_inc_check();

        Ok(())
    }

    /// AEAD encryption in place
    pub fn encrypt_with_ad_in_place(
        &mut self,
        ad: &[u8],
        in_out: &mut [u8],
        plaintext_len: usize,
    ) -> CipherResult<usize> {
        if self.overflowed {
            return Err(CipherError::NonceOverflow);
        }

        let size = C::encrypt_in_place(&self.k, self.n, ad, in_out, plaintext_len);
        self.nonce_inc_check();

        Ok(size)
    }

    /// AEAD decryption
    pub fn decrypt_with_ad(
        &mut self,
        ad: &[u8],
        ciphertext: &[u8],
        out: &mut [u8],
    ) -> CipherResult<()> {
        if self.overflowed {
            return Err(CipherError::NonceOverflow);
        }

        C::decrypt(&self.k, self.n, ad, ciphertext, out)?;
        self.nonce_inc_check();

        Ok(())
    }

    /// AEAD decryption in place
    pub fn decrypt_with_ad_in_place(
        &mut self,
        ad: &[u8],
        in_out: &mut [u8],
        ciphertext_len: usize,
    ) -> CipherResult<usize> {
        if self.overflowed {
            return Err(CipherError::NonceOverflow);
        }

        let size = C::decrypt_in_place(&self.k, self.n, ad, in_out, ciphertext_len)?;
        self.nonce_inc_check();
        Ok(size)
    }

    /// Get current nonce value
    pub fn get_nonce(&self) -> u64 {
        self.n
    }

    /// Set nonce value
    ///
    /// # Warning
    /// **Do not reuse nonces**
    pub fn set_nonce(&mut self, nonce: u64) {
        self.n = nonce;
    }

    /// Take ownership of key and nonce of this state
    pub fn take(self) -> (C::Key, u64) {
        (self.k.clone(), self.n)
    }

    /// Rekey
    ///
    /// Rekeys as per Noise spec parts 4.2 and 11.3
    pub fn rekey(&mut self) {
        self.k = C::rekey(&self.k)
    }
}