rns-crypto 0.2.2

use alloc::vec::Vec;
use core::fmt;

use aes::cipher::{block_padding::NoPadding, BlockModeDecrypt, BlockModeEncrypt, KeyIvInit};

const BLOCK_SIZE: usize = 16;

/// Error returned by the raw (unpadded) AES-CBC helpers.
#[derive(Debug, PartialEq)]
pub enum CbcError {
    /// The input length is not a multiple of the 16-byte block size.
    UnalignedLength,
}

impl fmt::Display for CbcError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            CbcError::UnalignedLength => {
                write!(f, "CBC input length is not a multiple of the block size")
            }
        }
    }
}

#[cfg(feature = "std")]
impl std::error::Error for CbcError {}

type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;

#[derive(Clone)]
pub struct Aes256 {
    key: [u8; 32],
}

impl Aes256 {
    /// Construct AES-256 with a compile-time-sized 256-bit key.
    ///
    /// The array reference deliberately makes an invalid key length
    /// unrepresentable at this API boundary.
    pub fn new(key: &[u8; 32]) -> Self {
        Aes256 { key: *key }
    }

    /// Encrypt block-aligned plaintext in CBC mode without padding.
    ///
    /// Returns `CbcError::UnalignedLength` if `plaintext.len()` is not a
    /// multiple of 16.
    pub fn encrypt_cbc(&self, plaintext: &[u8], iv: &[u8; 16]) -> Result<Vec<u8>, CbcError> {
        if !plaintext.len().is_multiple_of(BLOCK_SIZE) {
            return Err(CbcError::UnalignedLength);
        }
        let mut buf = plaintext.to_vec();
        self.encrypt_cbc_in_place(&mut buf, iv)?;
        Ok(buf)
    }

    /// [`encrypt_cbc`](Self::encrypt_cbc) over `buf` in place.
    pub fn encrypt_cbc_in_place(&self, buf: &mut [u8], iv: &[u8; 16]) -> Result<(), CbcError> {
        if !buf.len().is_multiple_of(BLOCK_SIZE) {
            return Err(CbcError::UnalignedLength);
        }
        #[cfg(target_arch = "x86_64")]
        if crate::aes256_ni::encrypt_cbc(&self.key, iv, buf) {
            return Ok(());
        }
        self.encrypt_cbc_portable(buf, iv)
    }

    /// [`encrypt_cbc_in_place`](Self::encrypt_cbc_in_place) through the
    /// generic cipher, on CPUs without a dedicated path.
    fn encrypt_cbc_portable(&self, buf: &mut [u8], iv: &[u8; 16]) -> Result<(), CbcError> {
        let message_len = buf.len();
        Aes256CbcEnc::new(&self.key.into(), iv.into())
            .encrypt_padded::<NoPadding>(buf, message_len)
            .map_err(|_| CbcError::UnalignedLength)?;
        Ok(())
    }

    /// A CBC decryptor that keeps its chaining state across calls, for
    /// ciphertext that arrives in pieces.
    pub(crate) fn cbc_decryptor(&self, iv: &[u8; 16]) -> cbc::Decryptor<aes::Aes256> {
        Aes256CbcDec::new(&self.key.into(), iv.into())
    }

    /// Decrypt block-aligned ciphertext in CBC mode without removing padding.
    ///
    /// Returns `CbcError::UnalignedLength` if `ciphertext.len()` is not a
    /// multiple of 16.
    pub fn decrypt_cbc(&self, ciphertext: &[u8], iv: &[u8; 16]) -> Result<Vec<u8>, CbcError> {
        if !ciphertext.len().is_multiple_of(BLOCK_SIZE) {
            return Err(CbcError::UnalignedLength);
        }
        let mut buf = ciphertext.to_vec();
        self.decrypt_cbc_in_place(&mut buf, iv)?;
        Ok(buf)
    }

    /// [`Self::decrypt_cbc`] in place: `buf` holds the ciphertext and is
    /// left holding the plaintext, padding included.
    pub fn decrypt_cbc_in_place(&self, buf: &mut [u8], iv: &[u8; 16]) -> Result<(), CbcError> {
        if !buf.len().is_multiple_of(BLOCK_SIZE) {
            return Err(CbcError::UnalignedLength);
        }
        Aes256CbcDec::new(&self.key.into(), iv.into())
            .decrypt_padded::<NoPadding>(buf)
            .map_err(|_| CbcError::UnalignedLength)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn dedicated_cbc_encryption_matches_the_generic_cipher() {
        let mut seed = 0x9e37_79b9u32;
        let mut next = move || {
            seed = seed.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
            (seed >> 24) as u8
        };
        for blocks in [1usize, 2, 3, 16, 255, 4096] {
            for _ in 0..8 {
                let key: [u8; 32] = core::array::from_fn(|_| next());
                let iv: [u8; 16] = core::array::from_fn(|_| next());
                let plain: Vec<u8> = (0..blocks * 16).map(|_| next()).collect();
                let aes = Aes256::new(&key);
                let mut fast = plain.clone();
                aes.encrypt_cbc_in_place(&mut fast, &iv).unwrap();
                let mut generic = plain.clone();
                aes.encrypt_cbc_portable(&mut generic, &iv).unwrap();
                assert_eq!(fast, generic);
                assert_eq!(aes.decrypt_cbc(&fast, &iv).unwrap(), plain);
            }
        }
    }

    #[test]
    fn test_aes256_encrypt_decrypt_block() {
        let key = [0u8; 32];
        let iv = [0u8; 16];
        let cipher = Aes256::new(&key);
        let plaintext = [0u8; 16];
        let encrypted = cipher.encrypt_cbc(&plaintext, &iv).unwrap();
        let decrypted = cipher.decrypt_cbc(&encrypted, &iv).unwrap();
        assert_eq!(decrypted, plaintext.to_vec());
    }

    #[test]
    fn test_aes256_cbc_roundtrip() {
        let key = [0x01u8; 32];
        let iv = [0x02u8; 16];
        let cipher = Aes256::new(&key);
        let plaintext = [0x03u8; 32];
        let encrypted = cipher.encrypt_cbc(&plaintext, &iv).unwrap();
        let decrypted = cipher.decrypt_cbc(&encrypted, &iv).unwrap();
        assert_eq!(decrypted, plaintext.to_vec());
    }

    #[test]
    fn test_aes256_known_vector() {
        // NIST AES-256 ECB test vector, verified via single-block CBC with zero IV
        let key: [u8; 32] = [
            0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
            0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
            0x09, 0x14, 0xdf, 0xf4,
        ];
        let plaintext: [u8; 16] = [
            0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96, 0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93,
            0x17, 0x2a,
        ];
        let expected: [u8; 16] = [
            0xf3, 0xee, 0xd1, 0xbd, 0xb5, 0xd2, 0xa0, 0x3c, 0x06, 0x4b, 0x5a, 0x7e, 0x3d, 0xb1,
            0x81, 0xf8,
        ];
        let cipher = Aes256::new(&key);
        let iv = [0u8; 16];
        let result = cipher.encrypt_cbc(&plaintext, &iv).unwrap();
        assert_eq!(&result[..16], &expected);
    }

    #[test]
    fn test_aes256_nist_sp800_38a_cbc_four_block_vector() {
        let key: [u8; 32] = decode_hex(concat!(
            "603deb1015ca71be2b73aef0857d7781",
            "1f352c073b6108d72d9810a30914dff4"
        ))
        .try_into()
        .unwrap();
        let iv: [u8; 16] = decode_hex("000102030405060708090a0b0c0d0e0f")
            .try_into()
            .unwrap();
        let plaintext = decode_hex(concat!(
            "6bc1bee22e409f96e93d7e117393172a",
            "ae2d8a571e03ac9c9eb76fac45af8e51",
            "30c81c46a35ce411e5fbc1191a0a52ef",
            "f69f2445df4f9b17ad2b417be66c3710"
        ));
        let expected = decode_hex(concat!(
            "f58c4c04d6e5f1ba779eabfb5f7bfbd6",
            "9cfc4e967edb808d679f777bc6702c7d",
            "39f23369a9d9bacfa530e26304231461",
            "b2eb05e2c39be9fcda6c19078c6a9d1b"
        ));
        let cipher = Aes256::new(&key);

        assert_eq!(cipher.encrypt_cbc(&plaintext, &iv).unwrap(), expected);
        assert_eq!(cipher.decrypt_cbc(&expected, &iv).unwrap(), plaintext);
    }

    #[test]
    fn test_aes256_empty_cbc_payload_is_stable() {
        let cipher = Aes256::new(&[0x11; 32]);
        assert!(cipher.encrypt_cbc(&[], &[0x22; 16]).unwrap().is_empty());
        assert!(cipher.decrypt_cbc(&[], &[0x22; 16]).unwrap().is_empty());
    }

    #[test]
    fn test_aes256_rejects_unaligned_ciphertext() {
        assert_eq!(
            Aes256::new(&[0; 32]).decrypt_cbc(&[0; 17], &[0; 16]),
            Err(CbcError::UnalignedLength)
        );
    }

    #[test]
    fn test_aes256_rejects_unaligned_plaintext() {
        assert_eq!(
            Aes256::new(&[0; 32]).encrypt_cbc(&[0; 15], &[0; 16]),
            Err(CbcError::UnalignedLength)
        );
    }

    fn decode_hex(hex: &str) -> Vec<u8> {
        (0..hex.len())
            .step_by(2)
            .map(|index| u8::from_str_radix(&hex[index..index + 2], 16).unwrap())
            .collect()
    }
}