wedpr_l_crypto_block_cipher_aes/
lib.rs

1// Copyright 2020 WeDPR Lab Project Authors. Licensed under Apache-2.0.
2
3//! AES block cipher functions.
4
5use aes::Aes256;
6use block_modes::{block_padding::Pkcs7, BlockMode, Cbc};
7
8use rand::RngCore;
9use std::borrow::BorrowMut;
10use wedpr_l_utils::{error::WedprError, traits::BlockCipher};
11
12/// Implements AES256 as a BlockCipher instance.
13#[derive(Default, Debug, Clone)]
14pub struct WedprBlockCipherAes256 {}
15
16// TODO: Add configurable implementation.
17type Aes256Cbc = Cbc<Aes256, Pkcs7>;
18const BLOCK_SIZE: usize = 16;
19const KEY_SIZE: usize = 32;
20
21impl BlockCipher for WedprBlockCipherAes256 {
22    fn encrypt<T: ?Sized + AsRef<[u8]>>(
23        &self,
24        message: &T,
25        key: &T,
26        iv: &T,
27    ) -> Result<Vec<u8>, WedprError> {
28        let cipher = match Aes256Cbc::new_var(key.as_ref(), iv.as_ref()) {
29            Ok(v) => v,
30            Err(_) => return Err(WedprError::FormatError),
31        };
32        // TODO: Find a better way to prepare the aligned buffer.
33        let input_length = message.as_ref().len();
34        let padding_length =
35            (input_length / BLOCK_SIZE + 1) * BLOCK_SIZE - input_length;
36        let mut buffer_vec = message.as_ref().to_vec();
37        buffer_vec.append([0u8].repeat(padding_length).borrow_mut());
38
39        return match cipher.encrypt(buffer_vec.as_mut_slice(), input_length) {
40            Ok(v) => Ok(v.to_vec()),
41            Err(_) => Err(WedprError::FormatError),
42        };
43    }
44
45    fn decrypt<T: ?Sized + AsRef<[u8]>>(
46        &self,
47        ciphertext: &T,
48        key: &T,
49        iv: &T,
50    ) -> Result<Vec<u8>, WedprError> {
51        let cipher = match Aes256Cbc::new_var(key.as_ref(), iv.as_ref()) {
52            Ok(v) => v,
53            Err(_) => return Err(WedprError::FormatError),
54        };
55
56        return match cipher.decrypt(ciphertext.as_ref().to_vec().as_mut_slice())
57        {
58            Ok(v) => Ok(v.to_vec()),
59            Err(_) => Err(WedprError::FormatError),
60        };
61    }
62
63    fn generate_key(&self) -> Vec<u8> {
64        let mut rng = rand::rngs::OsRng::default();
65        let mut key = [0u8; KEY_SIZE];
66        rng.fill_bytes(&mut key);
67        key.to_vec()
68    }
69
70    fn generate_iv(&self) -> Vec<u8> {
71        let mut rng = rand::rngs::OsRng::default();
72        let mut iv = [0u8; BLOCK_SIZE];
73        rng.fill_bytes(&mut iv);
74        iv.to_vec()
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    #[test]
83    fn test_aes256_long_message() {
84        let aes256 = WedprBlockCipherAes256::default();
85        let key = aes256.generate_key();
86        let iv = aes256.generate_iv();
87        let msg = b"helloworld1 helloworld2 helloworld3 helloworld4";
88
89        let ciphertext = aes256.encrypt(&msg.to_vec(), &key, &iv).unwrap();
90        let decrypted_msg = aes256.decrypt(&ciphertext, &key, &iv).unwrap();
91        assert_eq!(decrypted_msg, msg);
92    }
93
94    #[test]
95    fn test_aes256_short_message() {
96        let aes256 = WedprBlockCipherAes256::default();
97        let key = aes256.generate_key();
98        let iv = aes256.generate_iv();
99        let msg = b"hello";
100
101        let ciphertext = aes256.encrypt(&msg.to_vec(), &key, &iv).unwrap();
102        let decrypted_msg = aes256.decrypt(&ciphertext, &key, &iv).unwrap();
103        assert_eq!(decrypted_msg, msg);
104    }
105}
wedpr_l_crypto_block_cipher_aes/lib.rs

wedpr_l_crypto_block_cipher_aes/
lib.rs
