pdf-syntax 0.6.1

//! AES-128 / AES-256 CBC for the PDF standard security handler.

use alloc::vec::Vec;

use aes::cipher::{BlockDecrypt, BlockEncrypt, KeyInit, generic_array::GenericArray};
use aes::{Aes128, Aes256};

/// AES with CBC chaining, over the RustCrypto `aes` crate.
///
/// The block cipher is not ours: a hand-written one is crypto code to maintain,
/// and table-driven software AES leaks through cache timing. The `aes` crate
/// uses the hardware instructions where the CPU has them and a constant-time
/// bitsliced fallback where it does not, in pure Rust (MIT OR Apache-2.0, no C).
/// An earlier bitwise implementation here also cost about 40 ms to decrypt a
/// 570 KB embedded font, once per page that used it (#468).
// The two key sizes differ by one expanded key schedule (256 bytes of round keys); boxing the larger
// would add an allocation to every stream decryption for no gain, so the size difference is accepted.
#[allow(clippy::large_enum_variant)]
#[derive(Clone)]
enum Block {
    A128(Aes128),
    A256(Aes256),
}

#[derive(Clone)]
pub(crate) struct AESCipher<const KEY_SIZE: usize> {
    inner: Block,
}

pub(crate) type AES128Cipher = AESCipher<16>;
pub(crate) type AES256Cipher = AESCipher<32>;

impl<const KEY_SIZE: usize> AESCipher<KEY_SIZE> {
    pub(crate) fn new(key: &[u8]) -> Option<Self> {
        if key.len() != KEY_SIZE {
            return None;
        }
        let inner = match KEY_SIZE {
            16 => Block::A128(Aes128::new_from_slice(key).ok()?),
            32 => Block::A256(Aes256::new_from_slice(key).ok()?),
            _ => return None,
        };
        Some(Self { inner })
    }

    pub(crate) fn encrypt_block(&self, input: &[u8; 16]) -> [u8; 16] {
        let mut block = GenericArray::clone_from_slice(input);
        match &self.inner {
            Block::A128(c) => c.encrypt_block(&mut block),
            Block::A256(c) => c.encrypt_block(&mut block),
        }
        block.into()
    }

    #[cfg(test)]
    pub(crate) fn decrypt_block(&self, input: &[u8; 16]) -> [u8; 16] {
        let mut block = GenericArray::clone_from_slice(input);
        match &self.inner {
            Block::A128(c) => c.decrypt_block(&mut block),
            Block::A256(c) => c.decrypt_block(&mut block),
        }
        block.into()
    }

    pub(crate) fn encrypt_cbc(&self, data: &[u8], iv: &[u8; 16]) -> Vec<u8> {
        let mut result = Vec::with_capacity(data.len() + 16);
        let mut current_iv = *iv;

        let mut padded_data = data.to_vec();
        let pad_len = 16 - (data.len() % 16);
        padded_data.extend(core::iter::repeat_n(pad_len as u8, pad_len));

        for chunk in padded_data.chunks(16) {
            let mut block = [0_u8; 16];
            block.copy_from_slice(chunk);
            for i in 0..16 {
                block[i] ^= current_iv[i];
            }
            let encrypted = self.encrypt_block(&block);
            result.extend_from_slice(&encrypted);
            current_iv = encrypted;
        }

        result
    }

    pub(crate) fn decrypt_cbc(&self, data: &[u8], iv: &[u8; 16], unpad: bool) -> Vec<u8> {
        // CBC decryption parallelises: P[i] = D(C[i]) xor C[i-1], so every block
        // can go through the cipher at once and the chaining is applied after.
        let mut blocks: Vec<aes::Block> = data
            .chunks_exact(16)
            .map(GenericArray::clone_from_slice)
            .collect();
        match &self.inner {
            Block::A128(c) => c.decrypt_blocks(&mut blocks),
            Block::A256(c) => c.decrypt_blocks(&mut blocks),
        }

        let mut result = Vec::with_capacity(blocks.len() * 16);
        let mut prev: &[u8] = iv;
        for (plain, cipher) in blocks.iter().zip(data.chunks_exact(16)) {
            result.extend(plain.iter().zip(prev).map(|(p, c)| p ^ c));
            prev = cipher;
        }

        if unpad
            && let Some(&pad_len) = result.last()
            && pad_len > 0
            && pad_len <= 16
            && result.len() >= pad_len as usize
        {
            let start = result.len() - pad_len as usize;
            if result[start..].iter().all(|&b| b == pad_len) {
                result.truncate(start);
            }
        }

        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use aes::cipher::{
        BlockDecryptMut, BlockEncrypt, BlockEncryptMut, KeyInit, generic_array::GenericArray,
    };
    use aes::{Aes128, Aes256};
    use cbc::{
        Decryptor, Encryptor,
        cipher::{KeyIvInit, block_padding::Pkcs7},
    };

    type Aes128Cbc = Encryptor<Aes128>;
    type Aes128CbcDec = Decryptor<Aes128>;
    type Aes256Cbc = Encryptor<Aes256>;
    type Aes256CbcDec = Decryptor<Aes256>;

    fn aes_128_block_test(key: &[u8; 16], plaintext: &[u8; 16]) {
        let our_cipher = AES128Cipher::new(key).unwrap();

        let external_cipher = Aes128::new_from_slice(key).unwrap();
        let mut external_block = GenericArray::clone_from_slice(plaintext);
        external_cipher.encrypt_block(&mut external_block);

        let our_ciphertext = our_cipher.encrypt_block(plaintext);
        assert_eq!(
            our_ciphertext,
            external_block.as_slice(),
            "encryption should match"
        );

        let our_decrypted = our_cipher.decrypt_block(&our_ciphertext);
        assert_eq!(
            our_decrypted, *plaintext,
            "roundtrip should recover original plaintext"
        );
    }

    fn aes_256_block_test(key: &[u8; 32], plaintext: &[u8; 16]) {
        let our_cipher = AES256Cipher::new(key).unwrap();

        let external_cipher = Aes256::new_from_slice(key).unwrap();
        let mut external_block = GenericArray::clone_from_slice(plaintext);
        external_cipher.encrypt_block(&mut external_block);

        let our_ciphertext = our_cipher.encrypt_block(plaintext);
        assert_eq!(
            our_ciphertext,
            external_block.as_slice(),
            "encryption should match"
        );

        let our_decrypted = our_cipher.decrypt_block(&our_ciphertext);
        assert_eq!(
            our_decrypted, *plaintext,
            "roundtrip should recover original plaintext"
        );
    }

    fn aes_128_cbc_test(key: &[u8; 16], iv: &[u8; 16], plaintext: &[u8]) {
        let our_cipher = AES128Cipher::new(key).unwrap();

        let our_ciphertext = our_cipher.encrypt_cbc(plaintext, iv);

        let external_decryptor = Aes128CbcDec::new_from_slices(key, iv).unwrap();
        let mut external_buffer = our_ciphertext.clone();
        let external_decrypted = external_decryptor
            .decrypt_padded_mut::<Pkcs7>(&mut external_buffer)
            .unwrap();

        assert_eq!(
            external_decrypted, plaintext,
            "AES-128 CBC: our encryption should be decryptable by external crate"
        );

        let external_encryptor = Aes128Cbc::new_from_slices(key, iv).unwrap();
        let mut buffer = plaintext.to_vec();
        buffer.resize(plaintext.len() + 16, 0);
        let external_ciphertext = external_encryptor
            .encrypt_padded_mut::<Pkcs7>(&mut buffer, plaintext.len())
            .unwrap();

        let our_decrypted = our_cipher.decrypt_cbc(external_ciphertext, iv, true);
        assert_eq!(
            our_decrypted, plaintext,
            "AES-128 CBC: external encryption should be decryptable by our implementation"
        );

        let our_roundtrip = our_cipher.decrypt_cbc(&our_ciphertext, iv, true);
        assert_eq!(
            our_roundtrip, plaintext,
            "AES-128 CBC: our roundtrip should work"
        );
    }

    fn aes_256_cbc_test(key: &[u8; 32], iv: &[u8; 16], plaintext: &[u8]) {
        let our_cipher = AES256Cipher::new(key).unwrap();

        let our_ciphertext = our_cipher.encrypt_cbc(plaintext, iv);

        let external_decryptor = Aes256CbcDec::new_from_slices(key, iv).unwrap();
        let mut external_buffer = our_ciphertext.clone();
        let external_decrypted = external_decryptor
            .decrypt_padded_mut::<Pkcs7>(&mut external_buffer)
            .unwrap();

        assert_eq!(
            external_decrypted, plaintext,
            "AES-256 CBC: our encryption should be decryptable by external crate"
        );

        let external_encryptor = Aes256Cbc::new_from_slices(key, iv).unwrap();
        let mut buffer = plaintext.to_vec();
        buffer.resize(plaintext.len() + 16, 0);
        let external_ciphertext = external_encryptor
            .encrypt_padded_mut::<Pkcs7>(&mut buffer, plaintext.len())
            .unwrap();

        let our_decrypted = our_cipher.decrypt_cbc(external_ciphertext, iv, true);
        assert_eq!(
            our_decrypted, plaintext,
            "AES-256 CBC: external encryption should be decryptable by our implementation"
        );

        let our_roundtrip = our_cipher.decrypt_cbc(&our_ciphertext, iv, true);
        assert_eq!(
            our_roundtrip, plaintext,
            "AES-256 CBC: our roundtrip should work"
        );
    }

    #[test]
    fn test_aes128_block_operations() {
        let test_cases = [
            (
                [
                    0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09,
                    0xcf, 0x4f, 0x3c,
                ],
                [
                    0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0,
                    0x37, 0x07, 0x34,
                ],
            ),
            (
                [
                    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
                    0x0d, 0x0e, 0x0f,
                ],
                [
                    0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc,
                    0xdd, 0xee, 0xff,
                ],
            ),
            ([0xff; 16], [0x00; 16]),
            ([0x00; 16], [0xff; 16]),
        ];

        for (key, plaintext) in test_cases {
            aes_128_block_test(&key, &plaintext);
        }
    }

    #[test]
    fn test_aes256_block_operations() {
        let test_cases = [
            (
                [
                    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,
                ],
                [
                    0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96, 0xe9, 0x3d, 0x7e, 0x11, 0x73,
                    0x93, 0x17, 0x2a,
                ],
            ),
            (
                [
                    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
                    0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19,
                    0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
                ],
                [
                    0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc,
                    0xdd, 0xee, 0xff,
                ],
            ),
            ([0xff; 32], [0x00; 16]),
            ([0x00; 32], [0xff; 16]),
        ];

        for (key, plaintext) in test_cases {
            aes_256_block_test(&key, &plaintext);
        }
    }

    #[test]
    fn test_aes128_encryption_decryption() {
        let test_cases = [
            (
                [
                    0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09,
                    0xcf, 0x4f, 0x3c,
                ],
                [
                    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
                    0x0d, 0x0e, 0x0f,
                ],
                b"Hello, World!" as &[u8],
            ),
            (
                [
                    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
                    0x0d, 0x0e, 0x0f,
                ],
                [0xff; 16],
                b"The quick brown fox jumps over the lazy dog",
            ),
            (
                [0xff; 16],
                [0x00; 16],
                b"AES-128 encryption test vector for CBC mode with PKCS7 padding",
            ),
            ([0xaa; 16], [0x55; 16], &[0x42_u8; 100]),
        ];

        for (key, iv, plaintext) in test_cases {
            aes_128_cbc_test(&key, &iv, plaintext);
        }
    }

    #[test]
    fn test_aes256_encryption_decryption() {
        let test_cases = [
            (
                [
                    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,
                ],
                [
                    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
                    0x0d, 0x0e, 0x0f,
                ],
                b"Hello, World!" as &[u8],
            ),
            (
                [
                    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
                    0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19,
                    0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
                ],
                [0xff; 16],
                b"The quick brown fox jumps over the lazy dog",
            ),
            (
                [0xff; 32],
                [0x00; 16],
                b"AES-256 encryption test vector for CBC mode with PKCS7 padding",
            ),
            ([0xaa; 32], [0x55; 16], &[0x42_u8; 100]),
        ];

        for (key, iv, plaintext) in test_cases {
            aes_256_cbc_test(&key, &iv, plaintext);
        }
    }

    #[test]
    fn test_aes128_cbc_decryption() {
        let test_cases = [
            (
                [
                    0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09,
                    0xcf, 0x4f, 0x3c,
                ],
                [
                    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
                    0x0d, 0x0e, 0x0f,
                ],
                b"Hello, World!" as &[u8],
            ),
            (
                [
                    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
                    0x0d, 0x0e, 0x0f,
                ],
                [0xff; 16],
                b"The quick brown fox jumps over the lazy dog",
            ),
            (
                [0xff; 16],
                [0x00; 16],
                b"AES-128 test vector for CBC mode with PKCS7 padding",
            ),
            ([0xaa; 16], [0x55; 16], &[0x42_u8; 100]),
        ];

        for (key, iv, plaintext) in test_cases {
            aes_128_cbc_test(&key, &iv, plaintext);
        }
    }

    #[test]
    fn test_aes256_cbc_decryption() {
        let test_cases = [
            (
                [
                    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,
                ],
                [
                    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
                    0x0d, 0x0e, 0x0f,
                ],
                b"Hello, World!" as &[u8],
            ),
            (
                [
                    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
                    0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19,
                    0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
                ],
                [0xff; 16],
                b"The quick brown fox jumps over the lazy dog",
            ),
            (
                [0xff; 32],
                [0x00; 16],
                b"AES-256 test vector for CBC mode with PKCS7 padding",
            ),
            ([0xaa; 32], [0x55; 16], &[0x42_u8; 100]),
        ];

        for (key, iv, plaintext) in test_cases {
            aes_256_cbc_test(&key, &iv, plaintext);
        }
    }

    #[test]
    fn test_edge_cases() {
        let key128 = [
            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
            0x4f, 0x3c,
        ];
        let key256 = [
            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 iv = [0x00; 16];

        aes_128_cbc_test(&key128, &iv, b"");
        aes_256_cbc_test(&key256, &iv, b"");

        aes_128_cbc_test(&key128, &iv, b"a");
        aes_256_cbc_test(&key256, &iv, b"a");

        aes_128_cbc_test(&key128, &iv, &[0xff_u8; 15]);
        aes_256_cbc_test(&key256, &iv, &[0xff_u8; 15]);

        aes_128_cbc_test(&key128, &iv, &[0xff_u8; 16]);
        aes_256_cbc_test(&key256, &iv, &[0xff_u8; 16]);

        aes_128_cbc_test(&key128, &iv, &[0xff_u8; 17]);
        aes_256_cbc_test(&key256, &iv, &[0xff_u8; 17]);
    }

    /// A font-sized stream, so the cipher is exercised on the amount of
    /// data that made the earlier hand-written one show up in a profile, against the `aes`
    /// crate rather than against ourselves.
    #[test]
    fn a_font_sized_stream_decrypts_like_the_reference_crate() {
        let mut state = 0x9e37_79b9_u32;
        let plaintext: Vec<u8> = (0..300_000)
            .map(|_| {
                state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
                (state >> 24) as u8
            })
            .collect();
        aes_128_cbc_test(&[0x42; 16], &[0x07; 16], &plaintext);
        aes_256_cbc_test(&[0x42; 32], &[0x07; 16], &plaintext);
    }

    #[test]
    fn test_invalid_key_lengths() {
        assert!(AES128Cipher::new(&[0_u8; 15]).is_none());
        assert!(AES128Cipher::new(&[0_u8; 17]).is_none());
        assert!(AES256Cipher::new(&[0_u8; 31]).is_none());
        assert!(AES256Cipher::new(&[0_u8; 33]).is_none());
    }
}