eme2-extended 0.1.2

use cipher::{
    Block, BlockCipherDecrypt, BlockCipherEncrypt, BlockSizeUser, InnerIvInit, Iv,
    IvSizeUser, common::InnerUser, typenum::Unsigned, Array, typenum::{U16, U32, U64, U128},
};
use core::fmt;

#[cfg(feature = "zeroize")]
use zeroize::{Zeroize, ZeroizeOnDrop};

/// Error types for EME2 operations.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Error {
    /// The provided data is shorter than the minimum block size (16 bytes).
    DataTooShort,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::DataTooShort => write!(f, "Data must be at least one block size for EME2"),
        }
    }
}

impl core::error::Error for Error {}

#[inline]
fn xor_blocks<C: BlockSizeUser>(out: &mut Block<C>, a: &Block<C>, b: &Block<C>) {
    for (o, (x, y)) in out.iter_mut().zip(a.iter().zip(b.iter())) {
        *o = *x ^ *y;
    }
}

#[inline]
fn xor_into<C: BlockSizeUser>(out: &mut Block<C>, b: &Block<C>) {
    for (o, x) in out.iter_mut().zip(b.iter()) {
        *o ^= *x;
    }
}

/// A trait for block sizes that support EME2 Galois Field polynomial arithmetic.
pub trait EmePoly: hybrid_array::ArraySize {
    /// Multiplies the given value by the polynomial `x` within GF(2^n).
    fn mult_by_two(val: &Array<u8, Self>) -> Array<u8, Self>;
}

// 64-bit Limb GF(2^n) Arithmetic for massive performance gains
macro_rules! impl_eme_poly {
    ($size:ty, $limbs:expr, $mask_bytes:expr) => {
        impl EmePoly for $size {
            #[inline]
            fn mult_by_two(val: &Array<u8, Self>) -> Array<u8, Self> {
                let mut res = Array::<u8, Self>::default();
                let mut v = [0u64; $limbs];

                // Load little-endian 64-bit limbs
                for i in 0..$limbs {
                    let mut buf = [0u8; 8];
                    buf.copy_from_slice(&val[i * 8..(i + 1) * 8]);
                    v[i] = u64::from_le_bytes(buf);
                }

                // Shift left by 1, carrying overflow to the next limb
                let carry_out = (v[$limbs - 1] >> 63) as u8;
                for i in (1..$limbs).rev() {
                    v[i] = (v[i] << 1) | (v[i - 1] >> 63);
                }
                v[0] <<= 1;

                // Store back to little-endian bytes
                for i in 0..$limbs {
                    res[i * 8..(i + 1) * 8].copy_from_slice(&v[i].to_le_bytes());
                }

                // Apply polynomial mask if there was a carry
                if carry_out != 0 {
                    let mask: &[u8] = &$mask_bytes;
                    for (r, &m) in res.iter_mut().zip(mask.iter()) {
                        *r ^= m;
                    }
                }
                res
            }
        }
    };
}

impl_eme_poly!(U16, 2, [0x87, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]);
impl_eme_poly!(U32, 4, { let mut m = [0u8; 32]; m[0] = 0x25; m[1] = 0x04; m });
impl_eme_poly!(U64, 8, { let mut m = [0u8; 64]; m[0] = 0x25; m[1] = 0x01; m });
impl_eme_poly!(U128, 16, { let mut m = [0u8; 128]; m[0] = 0x01; m[2] = 0x08; m });

#[inline]
fn encrypt_block<C>(cipher: &C, block: &mut Block<C>)
where
    C: BlockCipherEncrypt + BlockSizeUser,
{
    cipher.encrypt_block(block);
}

#[inline]
fn decrypt_block<C>(cipher: &C, block: &mut Block<C>)
where
    C: BlockCipherDecrypt + BlockSizeUser,
{
    cipher.decrypt_block(block);
}

/// EME2 cipher mode.
#[derive(Clone)]
#[cfg_attr(feature = "zeroize", derive(ZeroizeOnDrop))]
pub struct Eme2<C>
where
    C: BlockCipherEncrypt + BlockCipherDecrypt + BlockSizeUser,
    C::BlockSize: EmePoly,
{
    #[cfg_attr(feature = "zeroize", zeroize(skip))]
    cipher: C,
    tweak: Block<C>,
}

impl<C> BlockSizeUser for Eme2<C>
where
    C: BlockCipherEncrypt + BlockCipherDecrypt + BlockSizeUser,
    C::BlockSize: EmePoly,
{
    type BlockSize = C::BlockSize;
}

impl<C> InnerUser for Eme2<C>
where
    C: BlockCipherEncrypt + BlockCipherDecrypt + BlockSizeUser,
    C::BlockSize: EmePoly,
{
    type Inner = C;
}

impl<C> IvSizeUser for Eme2<C>
where
    C: BlockCipherEncrypt + BlockCipherDecrypt + BlockSizeUser,
    C::BlockSize: EmePoly,
{
    type IvSize = C::BlockSize;
}

impl<C> InnerIvInit for Eme2<C>
where
    C: BlockCipherEncrypt + BlockCipherDecrypt + BlockSizeUser,
    C::BlockSize: EmePoly,
{
    #[inline]
    fn inner_iv_init(cipher: C, iv: &Iv<Self>) -> Self {
        Self {
            cipher,
            tweak: iv.clone(),
        }
    }
}

impl<C> cipher::AlgorithmName for Eme2<C>
where
    C: BlockCipherEncrypt + BlockCipherDecrypt + BlockSizeUser + cipher::AlgorithmName,
    C::BlockSize: EmePoly,
{
    fn write_alg_name(f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("Eme2<")?;
        <C as cipher::AlgorithmName>::write_alg_name(f)?;
        f.write_str(">")
    }
}

impl<C> fmt::Debug for Eme2<C>
where
    C: BlockCipherEncrypt + BlockCipherDecrypt + BlockSizeUser + cipher::AlgorithmName,
    C::BlockSize: EmePoly,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("Eme2<")?;
        <C as cipher::AlgorithmName>::write_alg_name(f)?;
        f.write_str("> { ... }")
    }
}

impl<C> Eme2<C>
where
    C: BlockCipherEncrypt + BlockCipherDecrypt + BlockSizeUser,
    C::BlockSize: EmePoly,
{
    /// Encrypts the `data` in-place using EME2 mode.
    /// `data` must be at least 16 bytes.
    pub fn encrypt(&self, data: &mut [u8]) -> Result<(), Error> {
        let bs = C::BlockSize::USIZE;
        let len = data.len();
        if len < bs {
            return Err(Error::DataTooShort);
        }

        let tweak = &self.tweak;
        let m = len.div_ceil(bs);
        let last_full = if len.is_multiple_of(bs) { m } else { m - 1 };

        let mut l_0 = Block::<C>::default();
        encrypt_block(&self.cipher, &mut l_0);

        // PASS 1: data[i] = ppp_i (Mutated strictly in-place)
        let mut l_current = C::BlockSize::mult_by_two(&l_0);
        for i in 0..last_full {
            let block: &mut Block<C> = (&mut data[i * bs..(i + 1) * bs]).try_into().expect("slice bounds are mathematically guaranteed to match block size");
            xor_into::<C>(block, &l_current);
            encrypt_block(&self.cipher, block);
            l_current = C::BlockSize::mult_by_two(&l_current);
        }

        let mut ppp_m = Block::<C>::default();
        if last_full < m {
            let rem = len % bs;
            ppp_m[..rem].copy_from_slice(&data[last_full * bs..len]);
            ppp_m[rem] = 0x80;
        }

        let mut sp = Block::<C>::default();
        for i in 1..last_full {
            let ppp_i: &Block<C> = (&data[i * bs..(i + 1) * bs]).try_into().expect("slice bounds are mathematically guaranteed to match block size");
            xor_into::<C>(&mut sp, ppp_i);
        }
        if last_full < m {
            xor_into::<C>(&mut sp, &ppp_m);
        }

        let ppp_0: &Block<C> = (&data[0..bs]).try_into().expect("slice bounds are mathematically guaranteed to match block size");
        let mut mp_1 = Block::<C>::default();
        xor_blocks::<C>(&mut mp_1, ppp_0, &sp);
        xor_into::<C>(&mut mp_1, tweak);

        let mut mc_1;
        let mut ccc_m = Block::<C>::default();
        let mut c_m = Block::<C>::default();
        
        let mut mm = Block::<C>::default(); 
        
        if last_full < m {
            mm.copy_from_slice(&mp_1); // Eliminates the unused_assignment warning
            encrypt_block(&self.cipher, &mut mm);
            mc_1 = mm.clone();
            encrypt_block(&self.cipher, &mut mc_1);

            let rem = len % bs;
            for i in 0..rem {
                c_m[i] = data[last_full * bs + i] ^ mm[i];
            }
            ccc_m[..rem].copy_from_slice(&c_m[..rem]);
            ccc_m[rem] = 0x80;
        } else {
            mc_1 = mp_1.clone();
            encrypt_block(&self.cipher, &mut mc_1);
        }

        let mut m_1 = Block::<C>::default();
        xor_blocks::<C>(&mut m_1, &mp_1, &mc_1);

        let mut current_m_j = m_1.clone();
        let mut current_m_k = m_1.clone();

        // PASS 2: data[i] = ccc_i (Mutated strictly in-place)
        for i in 1..last_full {
            let block: &mut Block<C> = (&mut data[i * bs..(i + 1) * bs]).try_into().expect("slice bounds are mathematically guaranteed to match block size");
            
            let k = i & 127;
            if k == 0 {
                let mut mp_j = Block::<C>::default();
                xor_blocks::<C>(&mut mp_j, block, &m_1);
                let mut mc_j = mp_j.clone();
                encrypt_block(&self.cipher, &mut mc_j);
                xor_blocks::<C>(&mut current_m_j, &mp_j, &mc_j);
                xor_blocks::<C>(block, &mc_j, &m_1);
                current_m_k = current_m_j.clone();

                #[cfg(feature = "zeroize")]
                {
                    mp_j.zeroize();
                    mc_j.zeroize();
                }
            } else {
                current_m_k = C::BlockSize::mult_by_two(&current_m_k);
                xor_into::<C>(block, &current_m_k);
            }
        }

        let mut sc = Block::<C>::default();
        for i in 1..last_full {
            let ccc_i: &Block<C> = (&data[i * bs..(i + 1) * bs]).try_into().expect("slice bounds are mathematically guaranteed to match block size");
            xor_into::<C>(&mut sc, ccc_i);
        }
        if last_full < m {
            xor_into::<C>(&mut sc, &ccc_m);
        }

        let mut ccc_0 = Block::<C>::default();
        xor_blocks::<C>(&mut ccc_0, &mc_1, &sc);
        xor_into::<C>(&mut ccc_0, tweak);
        data[0..bs].copy_from_slice(&ccc_0);

        // PASS 3: data[i] = ciphertext (Mutated strictly in-place)
        l_current = C::BlockSize::mult_by_two(&l_0);
        for i in 0..last_full {
            let block: &mut Block<C> = (&mut data[i * bs..(i + 1) * bs]).try_into().expect("slice bounds are mathematically guaranteed to match block size");
            encrypt_block(&self.cipher, block);
            xor_into::<C>(block, &l_current);
            l_current = C::BlockSize::mult_by_two(&l_current);
        }

        if last_full < m {
            let rem = len % bs;
            data[last_full * bs..len].copy_from_slice(&c_m[..rem]);
        }

        #[cfg(feature = "zeroize")]
        {
            l_0.zeroize();
            l_current.zeroize();
            ppp_m.zeroize();
            sp.zeroize();
            mp_1.zeroize();
            mc_1.zeroize();
            ccc_m.zeroize();
            c_m.zeroize();
            m_1.zeroize();
            current_m_j.zeroize();
            current_m_k.zeroize();
            sc.zeroize();
            ccc_0.zeroize();
            mm.zeroize();
        }

        Ok(())
    }

    /// Decrypts the `data` in-place using EME2 mode.
    /// `data` must be at least 16 bytes.
    pub fn decrypt(&self, data: &mut [u8]) -> Result<(), Error> {
        let bs = C::BlockSize::USIZE;
        let len = data.len();
        if len < bs {
            return Err(Error::DataTooShort);
        }

        let tweak = &self.tweak;
        let m = len.div_ceil(bs);
        let last_full = if len.is_multiple_of(bs) { m } else { m - 1 };

        let mut l_0 = Block::<C>::default();
        encrypt_block(&self.cipher, &mut l_0);

        let mut l_current = C::BlockSize::mult_by_two(&l_0);
        
        // PASS 1: data[i] = ccc_i (Mutated strictly in-place)
        for i in 0..last_full {
            let block: &mut Block<C> = (&mut data[i * bs..(i + 1) * bs]).try_into().expect("slice bounds are mathematically guaranteed to match block size");
            xor_into::<C>(block, &l_current);
            decrypt_block(&self.cipher, block); 
            l_current = C::BlockSize::mult_by_two(&l_current);
        }

        let mut ccc_m = Block::<C>::default();
        if last_full < m {
            let rem = len % bs;
            ccc_m[..rem].copy_from_slice(&data[last_full * bs..len]);
            ccc_m[rem] = 0x80;
        }

        let mut sc = Block::<C>::default();
        for i in 1..last_full {
            let ccc_i: &Block<C> = (&data[i * bs..(i + 1) * bs]).try_into().expect("slice bounds are mathematically guaranteed to match block size");
            xor_into::<C>(&mut sc, ccc_i);
        }
        if last_full < m {
            xor_into::<C>(&mut sc, &ccc_m);
        }

        let ccc_0: &Block<C> = (&data[0..bs]).try_into().expect("slice bounds are mathematically guaranteed to match block size");
        let mut mc_1 = Block::<C>::default();
        xor_blocks::<C>(&mut mc_1, ccc_0, &sc);
        xor_into::<C>(&mut mc_1, tweak);

        let mut mp_1;
        let mut ppp_m = Block::<C>::default();
        let mut p_m = Block::<C>::default();
        
        let mut mm = Block::<C>::default();

        if last_full < m {
            mm.copy_from_slice(&mc_1); // Eliminates the unused_assignment warning
            decrypt_block(&self.cipher, &mut mm);
            mp_1 = mm.clone();
            decrypt_block(&self.cipher, &mut mp_1);

            let rem = len % bs;
            for i in 0..rem {
                p_m[i] = data[last_full * bs + i] ^ mm[i];
            }
            ppp_m[..rem].copy_from_slice(&p_m[..rem]);
            ppp_m[rem] = 0x80;
        } else {
            mp_1 = mc_1.clone();
            decrypt_block(&self.cipher, &mut mp_1);
        }

        let mut m_1 = Block::<C>::default();
        xor_blocks::<C>(&mut m_1, &mp_1, &mc_1);

        let mut current_m_j = m_1.clone();
        let mut current_m_k = m_1.clone();

        // PASS 2: data[i] = ppp_i (Mutated strictly in-place)
        for i in 1..last_full {
            let block: &mut Block<C> = (&mut data[i * bs..(i + 1) * bs]).try_into().expect("slice bounds are mathematically guaranteed to match block size");

            let k = i & 127;
            if k == 0 {
                let mut mc_j = Block::<C>::default();
                xor_blocks::<C>(&mut mc_j, block, &m_1);
                let mut mp_j = mc_j.clone();
                decrypt_block(&self.cipher, &mut mp_j);
                xor_blocks::<C>(&mut current_m_j, &mp_j, &mc_j);
                xor_blocks::<C>(block, &mp_j, &m_1);
                current_m_k = current_m_j.clone();

                #[cfg(feature = "zeroize")]
                {
                    mc_j.zeroize();
                    mp_j.zeroize();
                }
            } else {
                current_m_k = C::BlockSize::mult_by_two(&current_m_k);
                xor_into::<C>(block, &current_m_k);
            }
        }

        let mut sp = Block::<C>::default();
        for i in 1..last_full {
            let ppp_i: &Block<C> = (&data[i * bs..(i + 1) * bs]).try_into().expect("slice bounds are mathematically guaranteed to match block size");
            xor_into::<C>(&mut sp, ppp_i);
        }
        if last_full < m {
            xor_into::<C>(&mut sp, &ppp_m);
        }

        let mut ppp_0 = Block::<C>::default();
        xor_blocks::<C>(&mut ppp_0, &mp_1, &sp);
        xor_into::<C>(&mut ppp_0, tweak);
        data[0..bs].copy_from_slice(&ppp_0);

        // PASS 3: data[i] = plaintext (Mutated strictly in-place)
        l_current = C::BlockSize::mult_by_two(&l_0);
        for i in 0..last_full {
            let block: &mut Block<C> = (&mut data[i * bs..(i + 1) * bs]).try_into().expect("slice bounds are mathematically guaranteed to match block size");
            decrypt_block(&self.cipher, block);
            xor_into::<C>(block, &l_current);
            l_current = C::BlockSize::mult_by_two(&l_current);
        }

        if last_full < m {
            let rem = len % bs;
            data[last_full * bs..len].copy_from_slice(&p_m[..rem]);
        }

        #[cfg(feature = "zeroize")]
        {
            l_0.zeroize();
            l_current.zeroize();
            ccc_m.zeroize();
            sc.zeroize();
            mc_1.zeroize();
            mp_1.zeroize();
            ppp_m.zeroize();
            p_m.zeroize();
            m_1.zeroize();
            current_m_j.zeroize();
            current_m_k.zeroize();
            sp.zeroize();
            ppp_0.zeroize();
            mm.zeroize();
        }

        Ok(())
    }
}