spacedls 0.4.0

use crate::core::{EncParams, EncProvider, EncSpec, GenericDecryptError, GenericEncryptError};
use cbc::cipher::KeyIvInit;
use cipher::block_padding::Padding;
use cipher::{
    BlockCipherDecrypt, BlockCipherEncrypt, BlockModeDecrypt, BlockModeEncrypt, IvSizeUser,
    KeySizeUser,
};
use core::marker::PhantomData;

/// Type-level specification for AES-CBC, deriving key and IV sizes from the cipher.
#[derive(Debug, Clone, Copy)]
pub struct AesCbcSpec<C>(PhantomData<C>);

impl<C> EncSpec for AesCbcSpec<C>
where
    C: BlockCipherDecrypt + BlockCipherEncrypt,
    cbc::Encryptor<C>: KeyIvInit,
    cbc::Decryptor<C>: KeyIvInit<
            KeySize = <cbc::Encryptor<C> as KeySizeUser>::KeySize,
            IvSize = <cbc::Encryptor<C> as IvSizeUser>::IvSize,
        >,
{
    type KeySize = <cbc::Encryptor<C> as KeySizeUser>::KeySize;
    type IvSize = <cbc::Encryptor<C> as IvSizeUser>::IvSize;
}

/// Software AES-CBC encryption provider with configurable block padding.
#[derive(Debug, Clone, Copy)]
pub struct AesCbc<C, Pad>(PhantomData<(C, Pad)>);

impl<C, Pad> AesCbc<C, Pad> {
    const BLOCK_SIZE: usize = 16;
}

impl<C, Pad> Default for AesCbc<C, Pad> {
    fn default() -> Self { Self(PhantomData) }
}

impl<C, Pad> EncProvider for AesCbc<C, Pad>
where
    C: BlockCipherDecrypt + BlockCipherEncrypt,
    cbc::Encryptor<C>: KeyIvInit,
    cbc::Decryptor<C>: KeyIvInit<
            KeySize = <cbc::Encryptor<C> as KeySizeUser>::KeySize,
            IvSize = <cbc::Encryptor<C> as IvSizeUser>::IvSize,
        >,
    Pad: Padding,
{
    type Spec = AesCbcSpec<C>;
    type EncryptError = GenericEncryptError;
    type DecryptError = GenericDecryptError;

    fn encrypt(
        &self,
        p: &EncParams<Self::Spec>,
        plain: &[u8],
        cipher: &mut [u8],
    ) -> Result<usize, Self::EncryptError> {
        let padded_len = ((plain.len() / Self::BLOCK_SIZE) + 1) * Self::BLOCK_SIZE;
        cipher.len().checked_sub(padded_len).ok_or(GenericEncryptError::OutputTooSmall)?;

        cipher[..plain.len()].copy_from_slice(plain);

        let enc = cbc::Encryptor::<C>::new(&p.key, &p.iv);

        Ok(enc
            .encrypt_padded::<Pad>(&mut cipher[..padded_len], plain.len())
            .map_err(|_| GenericEncryptError::OutputTooSmall)?
            .len())
    }

    fn decrypt(
        &self,
        p: &EncParams<Self::Spec>,
        cipher: &mut [u8],
    ) -> Result<usize, Self::DecryptError> {
        if cipher.len() < Self::BLOCK_SIZE {
            return Err(GenericDecryptError::InputTooShort);
        }
        if !cipher.len().is_multiple_of(Self::BLOCK_SIZE) {
            return Err(GenericDecryptError::InputNotBlockAligned);
        }

        let dec = cbc::Decryptor::<C>::new(&p.key, &p.iv);

        Ok(dec
            .decrypt_padded::<Pad>(cipher)
            .map_err(|_| GenericDecryptError::InvalidCiphertext)?
            .len())
    }

    fn pad_len(size: usize) -> u16 {
        (((size / Self::BLOCK_SIZE) + 1) * Self::BLOCK_SIZE - size) as u16
    }
}

#[cfg(test)]
mod tests {
    use super::{AesCbc, AesCbcSpec};
    use crate::core::{EncParams, EncProvider, GenericDecryptError, GenericEncryptError};
    use aes::Aes128;
    use assert_matches::assert_matches;
    use cipher::block_padding::Pkcs7;
    use std::vec;

    type Provider = AesCbc<Aes128, Pkcs7>;
    type Params = EncParams<AesCbcSpec<Aes128>>;

    fn provider() -> Provider { Provider::default() }

    fn params() -> Params {
        Params {
            key: [
                0x2c, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
                0x4f, 0x3c,
            ]
            .into(),
            iv: [
                0x03, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
                0x0e, 0x0f,
            ]
            .into(),
        }
    }

    #[test]
    fn round_trip_multiple_lengths() {
        let provider = provider();
        let params = params();
        let messages: [&[u8]; 6] = [
            b"",
            b"A",
            b"123456789012345",
            b"1234567890123456",
            b"12345678901234567",
            b"Space Data Link Security",
        ];

        for plain in messages {
            let mut cipher = vec![0_u8; ((plain.len() / 16) + 1) * 16];
            let written = provider.encrypt(&params, plain, &mut cipher).unwrap();

            assert!(written.is_multiple_of(16));
            assert!(written > 0);

            let mut decrypt_buf = cipher[..written].to_vec();
            let recovered = provider.decrypt(&params, &mut decrypt_buf).unwrap();
            assert_eq!(&decrypt_buf[..recovered], plain);
        }
    }

    #[test]
    fn rejects_too_small_output_buffer() {
        let provider = provider();
        let params = params();
        let plain = b"1234567890123456";
        let mut cipher = vec![0_u8; 16];

        let result = provider.encrypt(&params, plain, &mut cipher);
        assert_matches!(result, Err(GenericEncryptError::OutputTooSmall));
    }

    #[test]
    fn rejects_non_aligned_or_too_short_buffers() {
        let provider = provider();
        let params = params();
        let mut short = vec![0_u8; 15];
        let mut non_aligned = vec![0_u8; 17];

        assert_matches!(
            provider.decrypt(&params, &mut short),
            Err(GenericDecryptError::InputTooShort)
        );
        assert_matches!(
            provider.decrypt(&params, &mut non_aligned),
            Err(GenericDecryptError::InputNotBlockAligned)
        );
    }

    #[test]
    fn rejects_invalid_ciphertext_padding() {
        let provider = provider();
        let params = params();
        let plain = b"invalid-padding-check";
        let mut cipher = vec![0_u8; ((plain.len() / 16) + 1) * 16];
        let written = provider.encrypt(&params, plain, &mut cipher).unwrap();

        let mut tampered = cipher[..written].to_vec();
        let last = tampered.len() - 1;
        tampered[last] ^= 0xff;

        assert_matches!(
            provider.decrypt(&params, &mut tampered),
            Err(GenericDecryptError::InvalidCiphertext)
        );
    }
}