Skip to main content

ironcore_documents/
aes.rs

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