zenoh_crypto/
cipher.rs

1//
2// Copyright (c) 2023 ZettaScale Technology
3//
4// This program and the accompanying materials are made available under the
5// terms of the Eclipse Public License 2.0 which is available at
6// http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0
7// which is available at https://www.apache.org/licenses/LICENSE-2.0.
8//
9// SPDX-License-Identifier: EPL-2.0 OR Apache-2.0
10//
11// Contributors:
12//   ZettaScale Zenoh Team, <zenoh@zettascale.tech>
13//
14use 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}