eme2-extended 0.1.0

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 i in 0..out.len() {
        out[i] = a[i] ^ b[i];
    }
}

#[inline]
fn xor_into<C: BlockSizeUser>(out: &mut Block<C>, b: &Block<C>) {
    for i in 0..out.len() {
        out[i] ^= b[i];
    }
}

/// Trait for block sizes that have defined EME2 Galois Field polynomials
pub trait EmePoly: hybrid_array::ArraySize {
    fn mult_by_two(val: &Array<u8, Self>) -> Array<u8, Self>;
}

impl EmePoly for U16 {
    #[inline]
    fn mult_by_two(val: &Array<u8, Self>) -> Array<u8, Self> {
        let mut res = val.clone();
        let carry_out = res[15] >> 7;
        let mut carry = 0;
        for i in 0..16 {
            let temp = res[i];
            res[i] = (temp << 1) | carry;
            carry = temp >> 7;
        }
        res[0] ^= 0u8.wrapping_sub(carry_out) & 0x87;
        res
    }
}

impl EmePoly for U32 {
    #[inline]
    fn mult_by_two(val: &Array<u8, Self>) -> Array<u8, Self> {
        let mut res = val.clone();
        let carry_out = res[31] >> 7;
        let mut carry = 0;
        for i in 0..32 {
            let temp = res[i];
            res[i] = (temp << 1) | carry;
            carry = temp >> 7;
        }
        let mask = 0u8.wrapping_sub(carry_out);
        res[0] ^= mask & 0x25;
        res[1] ^= mask & 0x04;
        res
    }
}

impl EmePoly for U64 {
    #[inline]
    fn mult_by_two(val: &Array<u8, Self>) -> Array<u8, Self> {
        let mut res = val.clone();
        let carry_out = res[63] >> 7;
        let mut carry = 0;
        for i in 0..64 {
            let temp = res[i];
            res[i] = (temp << 1) | carry;
            carry = temp >> 7;
        }
        let mask = 0u8.wrapping_sub(carry_out);
        res[0] ^= mask & 0x25;
        res[1] ^= mask & 0x01;
        res
    }
}

impl EmePoly for U128 {
    #[inline]
    fn mult_by_two(val: &Array<u8, Self>) -> Array<u8, Self> {
        let mut res = val.clone();
        let carry_out = res[127] >> 7;
        let mut carry = 0;
        for i in 0..128 {
            let temp = res[i];
            res[i] = (temp << 1) | carry;
            carry = temp >> 7;
        }
        let mask = 0u8.wrapping_sub(carry_out);
        res[0] ^= mask & 0x43;
        res[2] ^= mask & 0x08;
        res
    }
}

#[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
        let mut l_current = C::BlockSize::mult_by_two(&l_0);
        for i in 0..last_full {
            let mut block = Block::<C>::default();
            block.copy_from_slice(&data[i * bs..(i + 1) * bs]);
            xor_into::<C>(&mut block, &l_current);
            encrypt_block(&self.cipher, &mut block);
            data[i * bs..(i + 1) * bs].copy_from_slice(&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 mut ppp_i = Block::<C>::default();
            ppp_i.copy_from_slice(&data[i * bs..(i + 1) * bs]);
            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();
        ppp_0.copy_from_slice(&data[0..bs]);

        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();
        
        if last_full < m {
            let mut mm = mp_1.clone();
            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
        for i in 1..last_full {
            let mut ppp_i = Block::<C>::default();
            ppp_i.copy_from_slice(&data[i * bs..(i + 1) * bs]);
            let mut ccc_i = Block::<C>::default();
            
            let k = i % 128;
            if k == 0 {
                let mut mp_j = Block::<C>::default();
                xor_blocks::<C>(&mut mp_j, &ppp_i, &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>(&mut ccc_i, &mc_j, &m_1);
                current_m_k = current_m_j.clone();
            } else {
                current_m_k = C::BlockSize::mult_by_two(&current_m_k);
                xor_blocks::<C>(&mut ccc_i, &ppp_i, &current_m_k);
            }
            data[i * bs..(i + 1) * bs].copy_from_slice(&ccc_i);
        }

        let mut sc = Block::<C>::default();
        for i in 1..last_full {
            let mut ccc_i = Block::<C>::default();
            ccc_i.copy_from_slice(&data[i * bs..(i + 1) * bs]);
            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
        l_current = C::BlockSize::mult_by_two(&l_0);
        for i in 0..last_full {
            let mut cc_i = Block::<C>::default();
            cc_i.copy_from_slice(&data[i * bs..(i + 1) * bs]);
            encrypt_block(&self.cipher, &mut cc_i);
            xor_into::<C>(&mut cc_i, &l_current);
            data[i * bs..(i + 1) * bs].copy_from_slice(&cc_i);
            
            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();
            ppp_0.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();
        }

        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
        for i in 0..last_full {
            let mut block = Block::<C>::default();
            block.copy_from_slice(&data[i * bs..(i + 1) * bs]);
            xor_into::<C>(&mut block, &l_current);
            decrypt_block(&self.cipher, &mut block); 
            data[i * bs..(i + 1) * bs].copy_from_slice(&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 mut ccc_i = Block::<C>::default();
            ccc_i.copy_from_slice(&data[i * bs..(i + 1) * bs]);
            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();
        ccc_0.copy_from_slice(&data[0..bs]);

        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();
        
        if last_full < m {
            let mut mm = mc_1.clone();
            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
        for i in 1..last_full {
            let mut ccc_i = Block::<C>::default();
            ccc_i.copy_from_slice(&data[i * bs..(i + 1) * bs]);
            let mut ppp_i = Block::<C>::default();

            let k = i % 128;
            if k == 0 {
                let mut mc_j = Block::<C>::default();
                xor_blocks::<C>(&mut mc_j, &ccc_i, &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>(&mut ppp_i, &mp_j, &m_1);
                current_m_k = current_m_j.clone();
            } else {
                current_m_k = C::BlockSize::mult_by_two(&current_m_k);
                xor_blocks::<C>(&mut ppp_i, &ccc_i, &current_m_k);
            }
            data[i * bs..(i + 1) * bs].copy_from_slice(&ppp_i);
        }

        let mut sp = Block::<C>::default();
        for i in 1..last_full {
            let mut ppp_i = Block::<C>::default();
            ppp_i.copy_from_slice(&data[i * bs..(i + 1) * bs]);
            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
        l_current = C::BlockSize::mult_by_two(&l_0);
        for i in 0..last_full {
            let mut pp_i = Block::<C>::default();
            pp_i.copy_from_slice(&data[i * bs..(i + 1) * bs]);
            decrypt_block(&self.cipher, &mut pp_i);
            xor_into::<C>(&mut pp_i, &l_current);
            data[i * bs..(i + 1) * bs].copy_from_slice(&pp_i);
            
            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();
            ccc_0.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();
        }

        Ok(())
    }
}