aes-elk 0.1.0

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

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

/// Error types for ELK operations.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Error {
    /// The IV provided to ELK mode cannot be all zeroes, as the LFSR would never advance.
    ZeroIv,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::ZeroIv => write!(f, "IV cannot be all zeroes in ELK mode"),
        }
    }
}

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

/// Trait to define the Galois field feedback polynomial mask for different block sizes.
pub trait ElkPolynomial {
    fn apply_mask(block: &mut [u8]);
}

impl ElkPolynomial for U16 {
    #[inline(always)]
    fn apply_mask(block: &mut [u8]) {
        // SELF-COMPUTED. NOT IN NIST SP 800-197!
        // Polynomial for n = 128:
        // X^128 + X^113 + X^97 + X^89 + X^71 + X^59 + X^37 + X^23 + 1

        block[1]  ^= 1 << 1; // X^113
        block[3]  ^= 1 << 1; // X^97
        block[4]  ^= 1 << 1; // X^89
        block[6]  ^= 1 << 7; // X^71
        block[8]  ^= 1 << 3; // X^59
        block[11] ^= 1 << 5; // X^37
        block[13] ^= 1 << 7; // X^23
        block[15] ^= 1 << 0; // 1 (X^0)
    }
}

impl ElkPolynomial for U32 {
    #[inline(always)]
    fn apply_mask(block: &mut [u8]) {
        // §4.4 (ELK: The Encrypted LFSR Keystream Mode of Operation) https://csrc.nist.gov/files/pubs/sp/800/197/iprd/docs/3_samvadini.pdf
        // Polynomial for n = 256:
        // X^256 + X^229 + X^193 + X^167 + X^139 + X^109 + X^71 + X^37 + 1

        block[3]  ^= 1 << 5; // X^229
        block[7]  ^= 1 << 1; // X^193
        block[11] ^= 1 << 7; // X^167
        block[14] ^= 1 << 3; // X^139
        block[18] ^= 1 << 5; // X^109
        block[23] ^= 1 << 7; // X^71
        block[27] ^= 1 << 5; // X^37
        block[31] ^= 1 << 0; // 1 (X^0)
    }
}

impl ElkPolynomial for U64 {
    #[inline(always)]
    fn apply_mask(block: &mut [u8]) {
        // §4.4 (ELK: The Encrypted LFSR Keystream Mode of Operation) https://csrc.nist.gov/files/pubs/sp/800/197/iprd/docs/3_samvadini.pdf
        // Polynomial for n = 512:
        // X^512 + X^461 + X^397 + X^331 + X^281 + X^211 + X^139 + X^71 + 1

        block[6]  ^= 1 << 5; // X^461
        block[14] ^= 1 << 5; // X^397
        block[22] ^= 1 << 3; // X^331
        block[28] ^= 1 << 1; // X^281
        block[37] ^= 1 << 3; // X^211
        block[46] ^= 1 << 3; // X^139
        block[55] ^= 1 << 7; // X^71
        block[63] ^= 1 << 0; // 1 (X^0)
    }
}

/// Advances the state array using a left-shifting LFSR.
#[inline]
fn step_lfsr<C>(state: &mut Block<C>)
where
    C: BlockSizeUser,
    C::BlockSize: ElkPolynomial,
{
    let overflow = state[0] >> 7;
    let len = state.len();
    
    // Shift left by 1 bit across the entire array
    for i in 0..(len - 1) {
        state[i] = (state[i] << 1) | (state[i + 1] >> 7);
    }
    state[len - 1] <<= 1;
    
    // Apply Galois Field primitive polynomial if an overflow occurred
    if overflow != 0 {
        C::BlockSize::apply_mask(state.as_mut_slice());
    }
}

/// Universal ELK (Encrypted LFSR Keystream) mode.
/// Supports any block cipher with 128, 256, 512, or 1024 bit block sizes.
#[derive(Clone)]
#[cfg_attr(feature = "zeroize", derive(ZeroizeOnDrop))]
pub struct Elk<C>
where
    C: BlockCipherEncrypt + BlockSizeUser,
    C::BlockSize: ElkPolynomial,
{
    #[cfg_attr(feature = "zeroize", zeroize(skip))]
    cipher: C,
    state: Block<C>,
}

impl<C> BlockSizeUser for Elk<C>
where
    C: BlockCipherEncrypt + BlockSizeUser,
    C::BlockSize: ElkPolynomial,
{
    type BlockSize = C::BlockSize;
}

impl<C> InnerUser for Elk<C>
where
    C: BlockCipherEncrypt + BlockSizeUser,
    C::BlockSize: ElkPolynomial,
{
    type Inner = C;
}

impl<C> IvSizeUser for Elk<C>
where
    C: BlockCipherEncrypt + BlockSizeUser,
    C::BlockSize: ElkPolynomial,
{
    type IvSize = C::BlockSize;
}

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

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

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

impl<C> Elk<C>
where
    C: BlockCipherEncrypt + BlockSizeUser,
    C::BlockSize: ElkPolynomial,
{
    /// Core ELK operation: apply the keystream to the `data` in-place.
    /// Because ELK is a stream cipher, encryption and decryption are identical.
    pub fn apply_keystream(&mut self, data: &mut [u8]) -> Result<(), Error> {
        // Prevent catastrophic infinite loop of zeroes
        if self.state.iter().all(|&b| b == 0) {
            return Err(Error::ZeroIv);
        }

        let block_size = self.state.len();
        let mut keystream = self.state.clone();

        for chunk in data.chunks_mut(block_size) {
            // 1. Generate the keystream by encrypting the current LFSR state
            keystream.clone_from(&self.state);
            self.cipher.encrypt_block(&mut keystream);

            // 2. XOR the keystream with the plaintext (or ciphertext) chunk
            for (byte, key_byte) in chunk.iter_mut().zip(keystream.iter()) {
                *byte ^= key_byte;
            }

            // 3. Step the LFSR to prepare for the next block
            step_lfsr::<C>(&mut self.state);
        }
        
        Ok(())
    }
}

#[cfg(kani)]
mod verification {
    use super::*;
    use cipher::{Block, BlockCipherEncrypt, BlockSizeUser};
    use cipher::typenum::U16;

    #[derive(Clone)]
    struct MockCipher<Size> {
        _marker: core::marker::PhantomData<Size>,
    }

    impl<Size> BlockSizeUser for MockCipher<Size>
    where
        Size: hybrid_array::ArraySize,
    {
        type BlockSize = Size;
    }

    impl<Size> BlockCipherEncrypt for MockCipher<Size>
    where
        Size: hybrid_array::ArraySize,
    {
        fn encrypt_block(&self, _block: &mut Block<Self>) {
            // Dummy encryption that does nothing
        }

        fn encrypt_with_backend(&self, _backend: impl cipher::BlockCipherEncClosure<BlockSize = Self::BlockSize>) {
            // Dummy implementation
        }
    }

    #[kani::proof]
    #[kani::unwind(22)]
    fn verify_elk_16_multi_block() {
        let state_bytes: [u8; 16] = kani::any();
        let state: Block<MockCipher<U16>> = hybrid_array::Array::from(state_bytes);
        let cipher = MockCipher { _marker: core::marker::PhantomData };
        let mut elk = Elk { cipher, state };

        // Test up to 3 full blocks (48 bytes) to force the LFSR to step multiple times
        let mut buffer = [0u8; 48];
        let data_len: usize = kani::any();
        kani::assume(data_len <= 48);
        let data = &mut buffer[..data_len];

        let _ = elk.apply_keystream(data);
    }
}