eme2 0.1.3

use cipher::{
    Block, BlockCipherDecrypt, BlockCipherEncrypt, BlockSizeUser, InnerIvInit, Iv,
    IvSizeUser, consts::U16, common::InnerUser,
};
use core::fmt;

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

const BLOCK_SIZE: usize = 16;

/// 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 16 bytes for EME2"),
        }
    }
}

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

#[inline]
fn xor_blocks(out: &mut [u8; BLOCK_SIZE], a: &[u8; BLOCK_SIZE], b: &[u8; BLOCK_SIZE]) {
    for i in 0..BLOCK_SIZE {
        out[i] = a[i] ^ b[i];
    }
}

#[inline]
fn xor_into(out: &mut [u8; BLOCK_SIZE], b: &[u8; BLOCK_SIZE]) {
    for i in 0..BLOCK_SIZE {
        out[i] ^= b[i];
    }
}

#[inline]
fn mult_by_two(val: &[u8; BLOCK_SIZE]) -> [u8; BLOCK_SIZE] {
    let mut res = [0u8; BLOCK_SIZE];

    // Load as two little-endian 64-bit limbs
    let mut v0 = u64::from_le_bytes(val[0..8].try_into().unwrap());
    let mut v1 = u64::from_le_bytes(val[8..16].try_into().unwrap());
    
    // Check if the most significant bit is set before shifting
    let carry_out = (v1 >> 63) as u8;
    
    // Shift left by 1, carrying the overflow from v0 into v1
    v1 = (v1 << 1) | (v0 >> 63);
    v0 <<= 1;
    
    // Store back into bytes
    res[0..8].copy_from_slice(&v0.to_le_bytes());
    res[8..16].copy_from_slice(&v1.to_le_bytes());
    
    // Apply polynomial mask (x^128 + x^7 + x^2 + x + 1 = 0x87) if carry_out was 1
    if carry_out != 0 {
        res[0] ^= 0x87;
    }
    
    res
}

#[inline]
fn encrypt_block<C>(cipher: &C, block: &mut [u8; 16])
where
    C: BlockCipherEncrypt + BlockSizeUser<BlockSize = U16>,
{
    let mut b: Block<C> = (*block).into();
    cipher.encrypt_block(&mut b);
    *block = b.into();
}

#[inline]
fn decrypt_block<C>(cipher: &C, block: &mut [u8; 16])
where
    C: BlockCipherDecrypt + BlockSizeUser<BlockSize = U16>,
{
    let mut b: Block<C> = (*block).into();
    cipher.decrypt_block(&mut b);
    *block = b.into();
}

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

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

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

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

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

impl<C> cipher::AlgorithmName for Eme2<C>
where
    C: BlockCipherEncrypt + BlockCipherDecrypt + BlockSizeUser<BlockSize = U16> + cipher::AlgorithmName,
{
    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<BlockSize = U16> + cipher::AlgorithmName,
{
    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<BlockSize = U16>,
{
    /// 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 len = data.len();
        if len < BLOCK_SIZE {
            return Err(Error::DataTooShort);
        }

        let tweak_slice: &[u8; 16] = self.tweak.as_slice().try_into().expect("slice bounds are mathematically guaranteed to match block size");
        let m = len.div_ceil(BLOCK_SIZE);
        let last_full = if len.is_multiple_of(BLOCK_SIZE) { m } else { m - 1 };

        let mut l_0 = [0u8; BLOCK_SIZE];
        encrypt_block(&self.cipher, &mut l_0);

        // PASS 1: data[i] = ppp_i
        let mut l_current = mult_by_two(&l_0);
        for i in 0..last_full {
            let mut block = [0u8; BLOCK_SIZE];
            block.copy_from_slice(&data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE]);
            xor_into(&mut block, &l_current);
            encrypt_block(&self.cipher, &mut block);
            data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE].copy_from_slice(&block);
            
            l_current = mult_by_two(&l_current);
        }

        let mut ppp_m = [0u8; BLOCK_SIZE];
        if last_full < m {
            let rem = len % BLOCK_SIZE;
            ppp_m[..rem].copy_from_slice(&data[last_full * BLOCK_SIZE..len]);
            ppp_m[rem] = 0x80;
        }

        let mut sp = [0u8; BLOCK_SIZE];
        for i in 1..last_full {
            let mut ppp_i = [0u8; BLOCK_SIZE];
            ppp_i.copy_from_slice(&data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE]);
            xor_into(&mut sp, &ppp_i);
        }
        if last_full < m {
            xor_into(&mut sp, &ppp_m);
        }

        let mut ppp_0 = [0u8; BLOCK_SIZE];
        ppp_0.copy_from_slice(&data[0..BLOCK_SIZE]);

        let mut mp_1 = [0u8; BLOCK_SIZE];
        xor_blocks(&mut mp_1, &ppp_0, &sp);
        xor_into(&mut mp_1, tweak_slice);

        let mut mc_1;
        let mut ccc_m = [0u8; BLOCK_SIZE];
        let mut c_m = [0u8; BLOCK_SIZE];
        let mut mm = [0u8; BLOCK_SIZE];
        
        if last_full < m {
            mm.copy_from_slice(&mp_1);
            encrypt_block(&self.cipher, &mut mm);
            mc_1 = mm;
            encrypt_block(&self.cipher, &mut mc_1);

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

        let mut m_1 = [0u8; BLOCK_SIZE];
        xor_blocks(&mut m_1, &mp_1, &mc_1);

        let mut current_m_j = m_1;
        let mut current_m_k = m_1;

        // PASS 2: data[i] = ccc_i
        for i in 1..last_full {
            let mut ppp_i = [0u8; BLOCK_SIZE];
            ppp_i.copy_from_slice(&data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE]);
            let mut ccc_i = [0u8; BLOCK_SIZE];
            
            // [PERFORMANCE FIX]: Bitwise AND for clock cycle efficiency
            let k = i & 127;
            if k == 0 {
                let mut mp_j = [0u8; BLOCK_SIZE];
                xor_blocks(&mut mp_j, &ppp_i, &m_1);
                let mut mc_j = mp_j;
                encrypt_block(&self.cipher, &mut mc_j);
                xor_blocks(&mut current_m_j, &mp_j, &mc_j);
                xor_blocks(&mut ccc_i, &mc_j, &m_1);
                current_m_k = current_m_j;

                // [SECURITY FIX]: Clean up the internal key refresh state
                #[cfg(feature = "zeroize")]
                {
                    mp_j.zeroize();
                    mc_j.zeroize();
                }
            } else {
                current_m_k = mult_by_two(&current_m_k);
                xor_blocks(&mut ccc_i, &ppp_i, &current_m_k);
            }
            data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE].copy_from_slice(&ccc_i);
        }

        let mut sc = [0u8; BLOCK_SIZE];
        for i in 1..last_full {
            let mut ccc_i = [0u8; BLOCK_SIZE];
            ccc_i.copy_from_slice(&data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE]);
            xor_into(&mut sc, &ccc_i);
        }
        if last_full < m {
            xor_into(&mut sc, &ccc_m);
        }

        let mut ccc_0 = [0u8; BLOCK_SIZE];
        xor_blocks(&mut ccc_0, &mc_1, &sc);
        xor_into(&mut ccc_0, tweak_slice);
        data[0..BLOCK_SIZE].copy_from_slice(&ccc_0);

        // PASS 3: data[i] = ciphertext
        l_current = mult_by_two(&l_0);
        for i in 0..last_full {
            let mut cc_i = [0u8; BLOCK_SIZE];
            cc_i.copy_from_slice(&data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE]);
            encrypt_block(&self.cipher, &mut cc_i);
            xor_into(&mut cc_i, &l_current);
            data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE].copy_from_slice(&cc_i);
            
            l_current = mult_by_two(&l_current);
        }

        if last_full < m {
            let rem = len % BLOCK_SIZE;
            data[last_full * BLOCK_SIZE..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();
            mm.zeroize(); // [SECURITY FIX]
        }

        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 len = data.len();
        if len < BLOCK_SIZE {
            return Err(Error::DataTooShort);
        }

        let tweak_slice: &[u8; 16] = self.tweak.as_slice().try_into().expect("slice bounds are mathematically guaranteed to match block size");
        let m = len.div_ceil(BLOCK_SIZE);
        let last_full = if len.is_multiple_of(BLOCK_SIZE) { m } else { m - 1 };

        let mut l_0 = [0u8; BLOCK_SIZE];
        encrypt_block(&self.cipher, &mut l_0);

        let mut l_current = mult_by_two(&l_0);
        // PASS 1: data[i] = ccc_i
        for i in 0..last_full {
            let mut block = [0u8; BLOCK_SIZE];
            block.copy_from_slice(&data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE]);
            xor_into(&mut block, &l_current);
            decrypt_block(&self.cipher, &mut block); 
            data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE].copy_from_slice(&block);
            
            l_current = mult_by_two(&l_current);
        }

        let mut ccc_m = [0u8; BLOCK_SIZE];
        if last_full < m {
            let rem = len % BLOCK_SIZE;
            ccc_m[..rem].copy_from_slice(&data[last_full * BLOCK_SIZE..len]);
            ccc_m[rem] = 0x80;
        }

        let mut sc = [0u8; BLOCK_SIZE];
        for i in 1..last_full {
            let mut ccc_i = [0u8; BLOCK_SIZE];
            ccc_i.copy_from_slice(&data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE]);
            xor_into(&mut sc, &ccc_i);
        }
        if last_full < m {
            xor_into(&mut sc, &ccc_m);
        }

        let mut ccc_0 = [0u8; BLOCK_SIZE];
        ccc_0.copy_from_slice(&data[0..BLOCK_SIZE]);

        let mut mc_1 = [0u8; BLOCK_SIZE];
        xor_blocks(&mut mc_1, &ccc_0, &sc);
        xor_into(&mut mc_1, tweak_slice);

        let mut mp_1;
        let mut ppp_m = [0u8; BLOCK_SIZE];
        let mut p_m = [0u8; BLOCK_SIZE];
        let mut mm = [0u8; BLOCK_SIZE];

        if last_full < m {
            mm.copy_from_slice(&mc_1);
            decrypt_block(&self.cipher, &mut mm);
            mp_1 = mm;
            decrypt_block(&self.cipher, &mut mp_1);

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

        let mut m_1 = [0u8; BLOCK_SIZE];
        xor_blocks(&mut m_1, &mp_1, &mc_1);

        let mut current_m_j = m_1;
        let mut current_m_k = m_1;

        // PASS 2: data[i] = ppp_i
        for i in 1..last_full {
            let mut ccc_i = [0u8; BLOCK_SIZE];
            ccc_i.copy_from_slice(&data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE]);
            let mut ppp_i = [0u8; BLOCK_SIZE];

            let k = i & 127;
            if k == 0 {
                let mut mc_j = [0u8; BLOCK_SIZE];
                xor_blocks(&mut mc_j, &ccc_i, &m_1);
                let mut mp_j = mc_j;
                decrypt_block(&self.cipher, &mut mp_j);
                xor_blocks(&mut current_m_j, &mp_j, &mc_j);
                xor_blocks(&mut ppp_i, &mp_j, &m_1);
                current_m_k = current_m_j;

                #[cfg(feature = "zeroize")]
                {
                    mc_j.zeroize();
                    mp_j.zeroize();
                }
            } else {
                current_m_k = mult_by_two(&current_m_k);
                xor_blocks(&mut ppp_i, &ccc_i, &current_m_k);
            }
            data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE].copy_from_slice(&ppp_i);
        }

        let mut sp = [0u8; BLOCK_SIZE];
        for i in 1..last_full {
            let mut ppp_i = [0u8; BLOCK_SIZE];
            ppp_i.copy_from_slice(&data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE]);
            xor_into(&mut sp, &ppp_i);
        }
        if last_full < m {
            xor_into(&mut sp, &ppp_m);
        }

        let mut ppp_0 = [0u8; BLOCK_SIZE];
        xor_blocks(&mut ppp_0, &mp_1, &sp);
        xor_into(&mut ppp_0, tweak_slice);
        data[0..BLOCK_SIZE].copy_from_slice(&ppp_0);

        // PASS 3: data[i] = plaintext
        l_current = mult_by_two(&l_0);
        for i in 0..last_full {
            let mut pp_i = [0u8; BLOCK_SIZE];
            pp_i.copy_from_slice(&data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE]);
            decrypt_block(&self.cipher, &mut pp_i);
            xor_into(&mut pp_i, &l_current);
            data[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE].copy_from_slice(&pp_i);
            
            l_current = mult_by_two(&l_current);
        }

        if last_full < m {
            let rem = len % BLOCK_SIZE;
            data[last_full * BLOCK_SIZE..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();
            mm.zeroize();
        }

        Ok(())
    }
}