ironcore_documents/
aes.rs

1// This module is dedicated with things to do with aes encryption/decryption.
2use super::Error;
3use crate::{impl_secret_debug, take_lock};
4use aes_gcm::{AeadCore, Aes256Gcm, KeyInit, Nonce, aead::Aead, aead::Payload};
5use bytes::Bytes;
6use rand::{CryptoRng, RngCore};
7use std::{
8    ops::DerefMut,
9    sync::{Arc, Mutex},
10};
11
12type Result<T> = core::result::Result<T, super::Error>;
13pub(crate) const IV_LEN: usize = 12;
14
15/// These bytes are the IV + CIPHERTEXT.
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct IvAndCiphertext(pub Bytes);
18
19impl IvAndCiphertext {
20    pub fn len(&self) -> usize {
21        self.0.len()
22    }
23
24    pub fn is_empty(&self) -> bool {
25        self.0.is_empty()
26    }
27}
28
29impl Default for IvAndCiphertext {
30    fn default() -> IvAndCiphertext {
31        IvAndCiphertext(Bytes::new())
32    }
33}
34
35impl AsRef<[u8]> for IvAndCiphertext {
36    fn as_ref(&self) -> &[u8] {
37        self.0.as_ref()
38    }
39}
40
41impl From<Bytes> for IvAndCiphertext {
42    fn from(b: Bytes) -> Self {
43        IvAndCiphertext(b)
44    }
45}
46
47impl From<Vec<u8>> for IvAndCiphertext {
48    fn from(v: Vec<u8>) -> Self {
49        IvAndCiphertext(v.into())
50    }
51}
52
53impl From<IvAndCiphertext> for Bytes {
54    fn from(p: IvAndCiphertext) -> Self {
55        p.0
56    }
57}
58
59/// Holds bytes of an aes encrypted value
60#[derive(Debug, Clone, PartialEq, Eq)]
61pub struct EncryptedDocument(pub Vec<u8>);
62
63/// Holds bytes which are decrypted (The actual document bytes).
64#[derive(Debug, Clone, PartialEq, Eq)]
65pub struct PlaintextDocument(pub Vec<u8>);
66
67#[derive(Clone, Copy, PartialEq, Eq)]
68pub struct EncryptionKey(pub [u8; 32]);
69impl_secret_debug!(EncryptionKey);
70
71/// Decrypt the AES encrypted payload using the key. Note that the IV is on the front of the payload.
72pub fn decrypt_document_with_attached_iv(
73    key: &EncryptionKey,
74    aes_encrypted_payload: &IvAndCiphertext,
75) -> Result<PlaintextDocument> {
76    let (iv_slice, ciphertext) = aes_encrypted_payload.0.split_at(IV_LEN);
77    let iv = iv_slice
78        .try_into()
79        .expect("IV conversion will always have 12 bytes.");
80    aes_decrypt_core(key, iv, ciphertext, &[]).map(PlaintextDocument)
81}
82
83/// Encrypt a document and put the iv on the front of it.
84pub fn encrypt_document_and_attach_iv<R: RngCore + CryptoRng>(
85    rng: Arc<Mutex<R>>,
86    key: EncryptionKey,
87    document: PlaintextDocument,
88) -> Result<IvAndCiphertext> {
89    let (iv, mut enc_data) = aes_encrypt(key, &document.0, &[], rng)?;
90    let mut iv_vec = iv.to_vec();
91    iv_vec.append(&mut enc_data.0);
92    Ok(IvAndCiphertext(iv_vec.into()))
93}
94
95pub(crate) fn aes_encrypt<R: RngCore + CryptoRng>(
96    key: EncryptionKey,
97    plaintext: &[u8],
98    associated_data: &[u8],
99    rng: Arc<Mutex<R>>,
100) -> Result<([u8; 12], EncryptedDocument)> {
101    let iv = {
102        let mut guard = take_lock(&rng);
103        Aes256Gcm::generate_nonce(&mut *guard.deref_mut())
104    };
105    aes_encrypt_with_iv(key, plaintext, iv.into(), associated_data)
106}
107
108pub(crate) fn aes_encrypt_with_iv(
109    key: EncryptionKey,
110    plaintext: &[u8],
111    iv: [u8; IV_LEN],
112    associated_data: &[u8],
113) -> Result<([u8; 12], EncryptedDocument)> {
114    let cipher = Aes256Gcm::new(&key.0.into());
115    let encrypted_bytes = cipher
116        .encrypt(
117            &iv.into(),
118            Payload {
119                msg: plaintext,
120                aad: associated_data,
121            },
122        )
123        .map_err(|_| Error::EncryptError("Encryption failed.".to_string()))?;
124    Ok((iv, EncryptedDocument(encrypted_bytes)))
125}
126
127pub(crate) fn aes_decrypt_core(
128    key: &EncryptionKey,
129    iv: [u8; 12],
130    ciphertext: &[u8],
131    associated_data: &[u8],
132) -> Result<Vec<u8>> {
133    let cipher = Aes256Gcm::new(&key.0.into());
134
135    cipher
136        .decrypt(
137            Nonce::from_slice(&iv),
138            Payload {
139                msg: ciphertext,
140                aad: associated_data,
141            },
142        )
143        .map_err(|_| {
144            Error::DecryptError(
145                "Decryption failed. Ensure the data and key are correct.".to_string(),
146            )
147        })
148}
149
150#[cfg(test)]
151mod test {
152    use super::*;
153    use hex_literal::hex;
154    use rand::SeedableRng;
155    use rand_chacha::ChaCha20Rng;
156
157    #[test]
158    fn test_probabilistic_roundtrip() {
159        let key = EncryptionKey(hex!(
160            "fffefdfcfbfaf9f8f7f6f5f4f3f2f1f0f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff"
161        ));
162        let plaintext = hex!("112233445566778899aabbccddee");
163        let (iv, encrypt_result) = aes_encrypt(
164            key,
165            &plaintext,
166            &[],
167            Arc::new(Mutex::new(rand::thread_rng())),
168        )
169        .unwrap();
170        let decrypt_result = aes_decrypt_core(&key, iv, &encrypt_result.0, &[]).unwrap();
171        assert_eq!(decrypt_result, plaintext);
172    }
173
174    #[test]
175    fn encrypt_decrypt_attached_roundtrip() {
176        let rng = ChaCha20Rng::seed_from_u64(13u64);
177        let key = EncryptionKey(hex!(
178            "fffefdfcfbfaf9f8f7f6f5f4f3f2f1f0f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff"
179        ));
180        let document = vec![1u8];
181        let encrypted = encrypt_document_and_attach_iv(
182            Arc::new(Mutex::new(rng)),
183            key,
184            PlaintextDocument(document.clone()),
185        )
186        .unwrap();
187        let result = decrypt_document_with_attached_iv(&key, &encrypted).unwrap();
188        assert_eq!(result.0, document);
189    }
190}