1use aes::{
15 cipher::{generic_array::GenericArray, BlockDecrypt, BlockEncrypt, KeyInit},
16 Aes128,
17};
18use rand::Rng;
19use zenoh_result::{bail, ZResult};
20
21use super::PseudoRng;
22
23pub struct BlockCipher {
24 inner: Aes128,
25}
26
27impl BlockCipher {
28 pub const BLOCK_SIZE: usize = 16;
29
30 pub fn new(key: [u8; Self::BLOCK_SIZE]) -> BlockCipher {
31 BlockCipher {
32 inner: aes::Aes128::new(&key.into()),
33 }
34 }
35
36 pub fn encrypt(&self, mut bytes: Vec<u8>, padding: &mut PseudoRng) -> Vec<u8> {
37 let modulo = bytes.len() % Self::BLOCK_SIZE;
38 if modulo != 0 {
39 let missing = Self::BLOCK_SIZE - modulo;
40 bytes.resize_with(bytes.len() + missing, || padding.gen::<u8>());
41 }
42
43 let mut start: usize = 0;
44 while start < bytes.len() {
45 let end = start + Self::BLOCK_SIZE;
46 let block = GenericArray::from_mut_slice(&mut bytes[start..end]);
47 self.inner.encrypt_block(block);
48 start += Self::BLOCK_SIZE;
49 }
50
51 bytes
52 }
53
54 pub fn decrypt(&self, mut bytes: Vec<u8>) -> ZResult<Vec<u8>> {
55 if bytes.len() % Self::BLOCK_SIZE != 0 {
56 bail!("Invalid bytes length to decode: {}", bytes.len());
57 }
58
59 let mut start: usize = 0;
60 while start < bytes.len() {
61 let end = start + Self::BLOCK_SIZE;
62 let block = GenericArray::from_mut_slice(&mut bytes[start..end]);
63 self.inner.decrypt_block(block);
64 start += Self::BLOCK_SIZE;
65 }
66
67 Ok(bytes)
68 }
69}
70
71mod tests {
72 #[test]
73 fn cipher() {
74 use rand::{RngCore, SeedableRng};
75
76 use super::{BlockCipher, PseudoRng};
77
78 fn encrypt_decrypt(cipher: &BlockCipher, prng: &mut PseudoRng) {
79 println!("\n[1]");
80 let t1 = "A".as_bytes().to_vec();
81 println!("Clear: {t1:?}");
82 let encrypted = cipher.encrypt(t1.clone(), prng);
83 println!("Encrypted: {encrypted:?}");
84 let decrypted = cipher.decrypt(encrypted).unwrap();
85 println!("Decrypted: {decrypted:?}");
86 assert_eq!(&t1[..], &decrypted[..t1.len()]);
87
88 println!("\n[2]");
89 let t2 = "Short string".as_bytes().to_vec();
90 println!("Clear: {t2:?}");
91 let encrypted = cipher.encrypt(t2.clone(), prng);
92 println!("Encrypted: {encrypted:?}");
93 let decrypted = cipher.decrypt(encrypted).unwrap();
94 println!("Decrypted: {decrypted:?}");
95 assert_eq!(&t2[..], &decrypted[..t2.len()]);
96
97 println!("\n[3]");
98 let t3 = "This is a medium string with some text".as_bytes().to_vec();
99 println!("Clear: {t3:?}");
100 let encrypted = cipher.encrypt(t3.clone(), prng);
101 println!("Encrypted: {encrypted:?}");
102 let decrypted = cipher.decrypt(encrypted).unwrap();
103 println!("Decrypted: {decrypted:?}");
104 assert_eq!(&t3[..], &decrypted[..t3.len()]);
105
106 println!("\n[4]");
107 let t4 = "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.".as_bytes().to_vec();
108 println!("Clear: {t4:?}");
109 let encrypted = cipher.encrypt(t4.clone(), prng);
110 println!("Encrypted: {encrypted:?}");
111 let decrypted = cipher.decrypt(encrypted).unwrap();
112 println!("Decrypted: {decrypted:?}");
113 assert_eq!(&t4[..], &decrypted[..t4.len()]);
114 }
115
116 const RUN: usize = 16;
117
118 let mut prng = PseudoRng::from_entropy();
119 let mut key = [0_u8; BlockCipher::BLOCK_SIZE];
120 prng.fill_bytes(&mut key);
121 let cipher = BlockCipher::new(key);
122
123 for _ in 0..RUN {
124 encrypt_decrypt(&cipher, &mut prng);
125 }
126 }
127}