merge main

core/storage/encryption.rs

@@ -1,14 +1,13 @@
 #![allow(unused_variables, dead_code)]
 use crate::{LimboError, Result};
 use aegis::aegis256::Aegis256;
-use aes_gcm::{
-    aead::{Aead, AeadCore, KeyInit, OsRng},
-    Aes256Gcm, Key, Nonce,
-};
+use aes_gcm::aead::{AeadCore, OsRng};
 use std::ops::Deref;
 use turso_macros::match_ignore_ascii_case;
-
-pub const ENCRYPTED_PAGE_SIZE: usize = 4096;
+// AEGIS-256 supports both 16 and 32 byte tags, we use the 16 byte variant, it is faster
+// and provides sufficient security for our use case.
+const AEGIS_TAG_SIZE: usize = 16;
+const AES256GCM_TAG_SIZE: usize = 16;
 
 #[repr(transparent)]
 #[derive(Clone)]
@@ -74,10 +73,25 @@ impl Drop for EncryptionKey {
     }
 }
 
+pub trait AeadCipher {
+    fn encrypt(&self, plaintext: &[u8], ad: &[u8]) -> Result<(Vec<u8>, Vec<u8>)>;
+    fn decrypt(&self, ciphertext: &[u8], nonce: &[u8], ad: &[u8]) -> Result<Vec<u8>>;
+
+    fn encrypt_detached(&self, plaintext: &[u8], ad: &[u8]) -> Result<(Vec<u8>, Vec<u8>, Vec<u8>)>;
+
+    fn decrypt_detached(
+        &self,
+        ciphertext: &[u8],
+        nonce: &[u8],
+        tag: &[u8],
+        ad: &[u8],
+    ) -> Result<Vec<u8>>;
+}
+
 // wrapper struct for AEGIS-256 cipher, because the crate we use is a bit low-level and we add
 // some nice abstractions here
 // note, the AEGIS has many variants and support for hardware acceleration. Here we just use the
-// vanilla version, which is still order of maginitudes faster than AES-GCM in software. Hardware
+// vanilla version, which is still order of magnitudes faster than AES-GCM in software. Hardware
 // based compilation is left for future work.
 #[derive(Clone)]
 pub struct Aegis256Cipher {
@@ -85,39 +99,154 @@ pub struct Aegis256Cipher {
 }
 
 impl Aegis256Cipher {
-    // AEGIS-256 supports both 16 and 32 byte tags, we use the 16 byte variant, it is faster
-    // and provides sufficient security for our use case.
-    const TAG_SIZE: usize = 16;
     fn new(key: &EncryptionKey) -> Self {
         Self { key: key.clone() }
     }
+}
 
-    fn encrypt(&self, plaintext: &[u8], ad: &[u8]) -> Result<(Vec<u8>, [u8; 32])> {
+impl AeadCipher for Aegis256Cipher {
+    fn encrypt(&self, plaintext: &[u8], ad: &[u8]) -> Result<(Vec<u8>, Vec<u8>)> {
         let nonce = generate_secure_nonce();
         let (ciphertext, tag) =
-            Aegis256::<16>::new(self.key.as_bytes(), &nonce).encrypt(plaintext, ad);
+            Aegis256::<AEGIS_TAG_SIZE>::new(self.key.as_bytes(), &nonce).encrypt(plaintext, ad);
+
         let mut result = ciphertext;
         result.extend_from_slice(&tag);
-        Ok((result, nonce))
+        Ok((result, nonce.to_vec()))
     }
 
-    fn decrypt(&self, ciphertext: &[u8], nonce: &[u8; 32], ad: &[u8]) -> Result<Vec<u8>> {
-        if ciphertext.len() < Self::TAG_SIZE {
-            return Err(LimboError::InternalError(
-                "Ciphertext too short for AEGIS-256".into(),
-            ));
+    fn decrypt(&self, ciphertext: &[u8], nonce: &[u8], ad: &[u8]) -> Result<Vec<u8>> {
+        if ciphertext.len() < AEGIS_TAG_SIZE {
+            return Err(LimboError::InternalError("Ciphertext too short".into()));
         }
-        let (ct, tag) = ciphertext.split_at(ciphertext.len() - Self::TAG_SIZE);
-        let tag_array: [u8; 16] = tag
-            .try_into()
-            .map_err(|_| LimboError::InternalError("Invalid tag size for AEGIS-256".into()))?;
+        let (ct, tag) = ciphertext.split_at(ciphertext.len() - AEGIS_TAG_SIZE);
+        let tag_array: [u8; AEGIS_TAG_SIZE] = tag.try_into().map_err(|_| {
+            LimboError::InternalError(format!("Invalid tag size for AEGIS-256 {AEGIS_TAG_SIZE}"))
+        })?;
 
-        let plaintext = Aegis256::<16>::new(self.key.as_bytes(), nonce)
+        let nonce_array: [u8; 32] = nonce
+            .try_into()
+            .map_err(|_| LimboError::InternalError("Invalid nonce size for AEGIS-256".into()))?;
+
+        Aegis256::<AEGIS_TAG_SIZE>::new(self.key.as_bytes(), &nonce_array)
             .decrypt(ct, &tag_array, ad)
-            .map_err(|_| {
-                LimboError::InternalError("AEGIS-256 decryption failed: invalid tag".into())
-            })?;
-        Ok(plaintext)
+            .map_err(|_| LimboError::InternalError("AEGIS-256 decryption failed".into()))
+    }
+
+    fn encrypt_detached(&self, plaintext: &[u8], ad: &[u8]) -> Result<(Vec<u8>, Vec<u8>, Vec<u8>)> {
+        let nonce = generate_secure_nonce();
+        let (ciphertext, tag) =
+            Aegis256::<AEGIS_TAG_SIZE>::new(self.key.as_bytes(), &nonce).encrypt(plaintext, ad);
+
+        Ok((ciphertext, tag.to_vec(), nonce.to_vec()))
+    }
+
+    fn decrypt_detached(
+        &self,
+        ciphertext: &[u8],
+        nonce: &[u8],
+        tag: &[u8],
+        ad: &[u8],
+    ) -> Result<Vec<u8>> {
+        let tag_array: [u8; AEGIS_TAG_SIZE] = tag.try_into().map_err(|_| {
+            LimboError::InternalError(format!("Invalid tag size for AEGIS-256 {AEGIS_TAG_SIZE}"))
+        })?;
+        let nonce_array: [u8; 32] = nonce
+            .try_into()
+            .map_err(|_| LimboError::InternalError("Invalid nonce size for AEGIS-256".into()))?;
+
+        Aegis256::<AEGIS_TAG_SIZE>::new(self.key.as_bytes(), &nonce_array)
+            .decrypt(ciphertext, &tag_array, ad)
+            .map_err(|_| LimboError::InternalError("AEGIS-256 decrypt_detached failed".into()))
+    }
+}
+
+#[derive(Clone)]
+pub struct Aes256GcmCipher {
+    key: EncryptionKey,
+}
+
+impl Aes256GcmCipher {
+    fn new(key: &EncryptionKey) -> Self {
+        Self { key: key.clone() }
+    }
+}
+
+impl AeadCipher for Aes256GcmCipher {
+    fn encrypt(&self, plaintext: &[u8], _ad: &[u8]) -> Result<(Vec<u8>, Vec<u8>)> {
+        use aes_gcm::aead::{AeadInPlace, KeyInit};
+        use aes_gcm::Aes256Gcm;
+
+        let cipher = Aes256Gcm::new_from_slice(self.key.as_bytes())
+            .map_err(|_| LimboError::InternalError("Bad AES key".into()))?;
+        let nonce = Aes256Gcm::generate_nonce(&mut rand::thread_rng());
+        let mut buffer = plaintext.to_vec();
+
+        let tag = cipher
+            .encrypt_in_place_detached(&nonce, b"", &mut buffer)
+            .map_err(|_| LimboError::InternalError("AES-GCM encrypt failed".into()))?;
+
+        buffer.extend_from_slice(&tag[..AES256GCM_TAG_SIZE]);
+        Ok((buffer, nonce.to_vec()))
+    }
+
+    fn decrypt(&self, ciphertext: &[u8], nonce: &[u8], ad: &[u8]) -> Result<Vec<u8>> {
+        use aes_gcm::aead::{AeadInPlace, KeyInit};
+        use aes_gcm::{Aes256Gcm, Nonce};
+
+        if ciphertext.len() < AES256GCM_TAG_SIZE {
+            return Err(LimboError::InternalError("Ciphertext too short".into()));
+        }
+        let (ct, tag) = ciphertext.split_at(ciphertext.len() - AES256GCM_TAG_SIZE);
+
+        let cipher = Aes256Gcm::new_from_slice(self.key.as_bytes())
+            .map_err(|_| LimboError::InternalError("Bad AES key".into()))?;
+        let nonce = Nonce::from_slice(nonce);
+
+        let mut buffer = ct.to_vec();
+        cipher
+            .decrypt_in_place_detached(nonce, ad, &mut buffer, tag.into())
+            .map_err(|_| LimboError::InternalError("AES-GCM decrypt failed".into()))?;
+
+        Ok(buffer)
+    }
+
+    fn encrypt_detached(&self, plaintext: &[u8], ad: &[u8]) -> Result<(Vec<u8>, Vec<u8>, Vec<u8>)> {
+        use aes_gcm::aead::{AeadInPlace, KeyInit};
+        use aes_gcm::Aes256Gcm;
+
+        let cipher = Aes256Gcm::new_from_slice(self.key.as_bytes())
+            .map_err(|_| LimboError::InternalError("Bad AES key".into()))?;
+        let nonce = Aes256Gcm::generate_nonce(&mut rand::thread_rng());
+
+        let mut buffer = plaintext.to_vec();
+        let tag = cipher
+            .encrypt_in_place_detached(&nonce, ad, &mut buffer)
+            .map_err(|_| LimboError::InternalError("AES-GCM encrypt_detached failed".into()))?;
+
+        Ok((buffer, nonce.to_vec(), tag.to_vec()))
+    }
+
+    fn decrypt_detached(
+        &self,
+        ciphertext: &[u8],
+        nonce: &[u8],
+        tag: &[u8],
+        ad: &[u8],
+    ) -> Result<Vec<u8>> {
+        use aes_gcm::aead::{AeadInPlace, KeyInit};
+        use aes_gcm::{Aes256Gcm, Nonce};
+
+        let cipher = Aes256Gcm::new_from_slice(self.key.as_bytes())
+            .map_err(|_| LimboError::InternalError("Bad AES key".into()))?;
+        let nonce = Nonce::from_slice(nonce);
+
+        let mut buffer = ciphertext.to_vec();
+        cipher
+            .decrypt_in_place_detached(nonce, ad, &mut buffer, tag.into())
+            .map_err(|_| LimboError::InternalError("AES-GCM decrypt_detached failed".into()))?;
+
+        Ok(buffer)
     }
 }
 
@@ -180,8 +309,8 @@ impl CipherMode {
     /// Returns the authentication tag size for this cipher mode.
     pub fn tag_size(&self) -> usize {
         match self {
-            CipherMode::Aes256Gcm => 16,
-            CipherMode::Aegis256 => 16,
+            CipherMode::Aes256Gcm => AES256GCM_TAG_SIZE,
+            CipherMode::Aegis256 => AEGIS_TAG_SIZE,
         }
     }
 
@@ -193,8 +322,17 @@ impl CipherMode {
 
 #[derive(Clone)]
 pub enum Cipher {
-    Aes256Gcm(Box<Aes256Gcm>),
-    Aegis256(Box<Aegis256Cipher>),
+    Aes256Gcm(Aes256GcmCipher),
+    Aegis256(Aegis256Cipher),
+}
+
+impl Cipher {
+    fn as_aead(&self) -> &dyn AeadCipher {
+        match self {
+            Cipher::Aes256Gcm(c) => c,
+            Cipher::Aegis256(c) => c,
+        }
+    }
 }
 
 impl std::fmt::Debug for Cipher {
@@ -210,10 +348,11 @@ impl std::fmt::Debug for Cipher {
 pub struct EncryptionContext {
     cipher_mode: CipherMode,
     cipher: Cipher,
+    page_size: usize,
 }
 
 impl EncryptionContext {
-    pub fn new(cipher_mode: CipherMode, key: &EncryptionKey) -> Result<Self> {
+    pub fn new(cipher_mode: CipherMode, key: &EncryptionKey, page_size: usize) -> Result<Self> {
         let required_size = cipher_mode.required_key_size();
         if key.as_slice().len() != required_size {
             return Err(crate::LimboError::InvalidArgument(format!(
@@ -225,15 +364,13 @@ impl EncryptionContext {
         }
 
         let cipher = match cipher_mode {
-            CipherMode::Aes256Gcm => {
-                let cipher_key: &Key<Aes256Gcm> = key.as_ref().into();
-                Cipher::Aes256Gcm(Box::new(Aes256Gcm::new(cipher_key)))
-            }
-            CipherMode::Aegis256 => Cipher::Aegis256(Box::new(Aegis256Cipher::new(key))),
+            CipherMode::Aes256Gcm => Cipher::Aes256Gcm(Aes256GcmCipher::new(key)),
+            CipherMode::Aegis256 => Cipher::Aegis256(Aegis256Cipher::new(key)),
         };
         Ok(Self {
             cipher_mode,
             cipher,
+            page_size,
         })
     }
 
@@ -255,36 +392,38 @@ impl EncryptionContext {
         tracing::debug!("encrypting page {}", page_id);
         assert_eq!(
             page.len(),
-            ENCRYPTED_PAGE_SIZE,
-            "Page data must be exactly {ENCRYPTED_PAGE_SIZE} bytes"
+            self.page_size,
+            "Page data must be exactly {} bytes",
+            self.page_size
         );
 
         let metadata_size = self.cipher_mode.metadata_size();
-        let reserved_bytes = &page[ENCRYPTED_PAGE_SIZE - metadata_size..];
+        let reserved_bytes = &page[self.page_size - metadata_size..];
         let reserved_bytes_zeroed = reserved_bytes.iter().all(|&b| b == 0);
         assert!(
             reserved_bytes_zeroed,
             "last reserved bytes must be empty/zero, but found non-zero bytes"
         );
 
-        let payload = &page[..ENCRYPTED_PAGE_SIZE - metadata_size];
+        let payload = &page[..self.page_size - metadata_size];
         let (encrypted, nonce) = self.encrypt_raw(payload)?;
 
         let nonce_size = self.cipher_mode.nonce_size();
         assert_eq!(
             encrypted.len(),
-            ENCRYPTED_PAGE_SIZE - nonce_size,
+            self.page_size - nonce_size,
             "Encrypted page must be exactly {} bytes",
-            ENCRYPTED_PAGE_SIZE - nonce_size
+            self.page_size - nonce_size
         );
 
-        let mut result = Vec::with_capacity(ENCRYPTED_PAGE_SIZE);
+        let mut result = Vec::with_capacity(self.page_size);
         result.extend_from_slice(&encrypted);
         result.extend_from_slice(&nonce);
         assert_eq!(
             result.len(),
-            ENCRYPTED_PAGE_SIZE,
-            "Encrypted page must be exactly {ENCRYPTED_PAGE_SIZE} bytes"
+            self.page_size,
+            "Encrypted page must be exactly {} bytes",
+            self.page_size
         );
         Ok(result)
     }
@@ -298,8 +437,9 @@ impl EncryptionContext {
         tracing::debug!("decrypting page {}", page_id);
         assert_eq!(
             encrypted_page.len(),
-            ENCRYPTED_PAGE_SIZE,
-            "Encrypted page data must be exactly {ENCRYPTED_PAGE_SIZE} bytes"
+            self.page_size,
+            "Encrypted page data must be exactly {} bytes",
+            self.page_size
         );
 
         let nonce_size = self.cipher_mode.nonce_size();
@@ -312,60 +452,40 @@ impl EncryptionContext {
         let metadata_size = self.cipher_mode.metadata_size();
         assert_eq!(
             decrypted_data.len(),
-            ENCRYPTED_PAGE_SIZE - metadata_size,
+            self.page_size - metadata_size,
             "Decrypted page data must be exactly {} bytes",
-            ENCRYPTED_PAGE_SIZE - metadata_size
+            self.page_size - metadata_size
         );
 
-        let mut result = Vec::with_capacity(ENCRYPTED_PAGE_SIZE);
+        let mut result = Vec::with_capacity(self.page_size);
         result.extend_from_slice(&decrypted_data);
-        result.resize(ENCRYPTED_PAGE_SIZE, 0);
+        result.resize(self.page_size, 0);
         assert_eq!(
             result.len(),
-            ENCRYPTED_PAGE_SIZE,
-            "Decrypted page data must be exactly {ENCRYPTED_PAGE_SIZE} bytes"
+            self.page_size,
+            "Decrypted page data must be exactly {} bytes",
+            self.page_size
         );
         Ok(result)
     }
 
     /// encrypts raw data using the configured cipher, returns ciphertext and nonce
     fn encrypt_raw(&self, plaintext: &[u8]) -> Result<(Vec<u8>, Vec<u8>)> {
-        match &self.cipher {
-            Cipher::Aes256Gcm(cipher) => {
-                let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
-                let ciphertext = cipher
-                    .encrypt(&nonce, plaintext)
-                    .map_err(|e| LimboError::InternalError(format!("Encryption failed: {e:?}")))?;
-                Ok((ciphertext, nonce.to_vec()))
-            }
-            Cipher::Aegis256(cipher) => {
-                let ad = b"";
-                let (ciphertext, nonce) = cipher.encrypt(plaintext, ad)?;
-                Ok((ciphertext, nonce.to_vec()))
-            }
-        }
+        self.cipher.as_aead().encrypt(plaintext, b"")
     }
 
     fn decrypt_raw(&self, ciphertext: &[u8], nonce: &[u8]) -> Result<Vec<u8>> {
-        match &self.cipher {
-            Cipher::Aes256Gcm(cipher) => {
-                let nonce = Nonce::from_slice(nonce);
-                let plaintext = cipher.decrypt(nonce, ciphertext).map_err(|e| {
-                    crate::LimboError::InternalError(format!("Decryption failed: {e:?}"))
-                })?;
-                Ok(plaintext)
-            }
-            Cipher::Aegis256(cipher) => {
-                let nonce_array: [u8; 32] = nonce.try_into().map_err(|_| {
-                    LimboError::InternalError(format!(
-                        "Invalid nonce size for AEGIS-256: expected 32, got {}",
-                        nonce.len()
-                    ))
-                })?;
-                let ad = b"";
-                cipher.decrypt(ciphertext, &nonce_array, ad)
-            }
-        }
+        self.cipher.as_aead().decrypt(ciphertext, nonce, b"")
+    }
+
+    fn encrypt_raw_detached(&self, plaintext: &[u8]) -> Result<(Vec<u8>, Vec<u8>, Vec<u8>)> {
+        self.cipher.as_aead().encrypt_detached(plaintext, b"")
+    }
+
+    fn decrypt_raw_detached(&self, ciphertext: &[u8], nonce: &[u8], tag: &[u8]) -> Result<Vec<u8>> {
+        self.cipher
+            .as_aead()
+            .decrypt_detached(ciphertext, nonce, tag, b"")
     }
 
     #[cfg(not(feature = "encryption"))]
@@ -391,10 +511,12 @@ fn generate_secure_nonce() -> [u8; 32] {
     nonce
 }
 
+#[cfg(feature = "encryption")]
 #[cfg(test)]
 mod tests {
     use super::*;
     use rand::Rng;
+    const DEFAULT_ENCRYPTED_PAGE_SIZE: usize = 4096;
 
     fn generate_random_hex_key() -> String {
         let mut rng = rand::thread_rng();
@@ -404,15 +526,14 @@ mod tests {
     }
 
     #[test]
-    #[cfg(feature = "encryption")]
     fn test_aes_encrypt_decrypt_round_trip() {
         let mut rng = rand::thread_rng();
         let cipher_mode = CipherMode::Aes256Gcm;
         let metadata_size = cipher_mode.metadata_size();
-        let data_size = ENCRYPTED_PAGE_SIZE - metadata_size;
+        let data_size = DEFAULT_ENCRYPTED_PAGE_SIZE - metadata_size;
 
         let page_data = {
-            let mut page = vec![0u8; ENCRYPTED_PAGE_SIZE];
+            let mut page = vec![0u8; DEFAULT_ENCRYPTED_PAGE_SIZE];
             page.iter_mut()
                 .take(data_size)
                 .for_each(|byte| *byte = rng.gen());
@@ -420,21 +541,21 @@ mod tests {
         };
 
         let key = EncryptionKey::from_hex_string(&generate_random_hex_key()).unwrap();
-        let ctx = EncryptionContext::new(CipherMode::Aes256Gcm, &key).unwrap();
+        let ctx = EncryptionContext::new(CipherMode::Aes256Gcm, &key, DEFAULT_ENCRYPTED_PAGE_SIZE)
+            .unwrap();
 
         let page_id = 42;
         let encrypted = ctx.encrypt_page(&page_data, page_id).unwrap();
-        assert_eq!(encrypted.len(), ENCRYPTED_PAGE_SIZE);
+        assert_eq!(encrypted.len(), DEFAULT_ENCRYPTED_PAGE_SIZE);
         assert_ne!(&encrypted[..data_size], &page_data[..data_size]);
         assert_ne!(&encrypted[..], &page_data[..]);
 
         let decrypted = ctx.decrypt_page(&encrypted, page_id).unwrap();
-        assert_eq!(decrypted.len(), ENCRYPTED_PAGE_SIZE);
+        assert_eq!(decrypted.len(), DEFAULT_ENCRYPTED_PAGE_SIZE);
         assert_eq!(decrypted, page_data);
     }
 
     #[test]
-    #[cfg(feature = "encryption")]
     fn test_aegis256_cipher_wrapper() {
         let key = EncryptionKey::from_hex_string(&generate_random_hex_key()).unwrap();
         let cipher = Aegis256Cipher::new(&key);
@@ -451,10 +572,10 @@ mod tests {
     }
 
     #[test]
-    #[cfg(feature = "encryption")]
     fn test_aegis256_raw_encryption() {
         let key = EncryptionKey::from_hex_string(&generate_random_hex_key()).unwrap();
-        let ctx = EncryptionContext::new(CipherMode::Aegis256, &key).unwrap();
+        let ctx = EncryptionContext::new(CipherMode::Aegis256, &key, DEFAULT_ENCRYPTED_PAGE_SIZE)
+            .unwrap();
 
         let plaintext = b"Hello, AEGIS-256!";
         let (ciphertext, nonce) = ctx.encrypt_raw(plaintext).unwrap();
@@ -467,15 +588,14 @@ mod tests {
     }
 
     #[test]
-    #[cfg(feature = "encryption")]
     fn test_aegis256_encrypt_decrypt_round_trip() {
         let mut rng = rand::thread_rng();
         let cipher_mode = CipherMode::Aegis256;
         let metadata_size = cipher_mode.metadata_size();
-        let data_size = ENCRYPTED_PAGE_SIZE - metadata_size;
+        let data_size = DEFAULT_ENCRYPTED_PAGE_SIZE - metadata_size;
 
         let page_data = {
-            let mut page = vec![0u8; ENCRYPTED_PAGE_SIZE];
+            let mut page = vec![0u8; DEFAULT_ENCRYPTED_PAGE_SIZE];
             page.iter_mut()
                 .take(data_size)
                 .for_each(|byte| *byte = rng.gen());
@@ -483,15 +603,16 @@ mod tests {
         };
 
         let key = EncryptionKey::from_hex_string(&generate_random_hex_key()).unwrap();
-        let ctx = EncryptionContext::new(CipherMode::Aegis256, &key).unwrap();
+        let ctx = EncryptionContext::new(CipherMode::Aegis256, &key, DEFAULT_ENCRYPTED_PAGE_SIZE)
+            .unwrap();
 
         let page_id = 42;
         let encrypted = ctx.encrypt_page(&page_data, page_id).unwrap();
-        assert_eq!(encrypted.len(), ENCRYPTED_PAGE_SIZE);
+        assert_eq!(encrypted.len(), DEFAULT_ENCRYPTED_PAGE_SIZE);
         assert_ne!(&encrypted[..data_size], &page_data[..data_size]);
 
         let decrypted = ctx.decrypt_page(&encrypted, page_id).unwrap();
-        assert_eq!(decrypted.len(), ENCRYPTED_PAGE_SIZE);
+        assert_eq!(decrypted.len(), DEFAULT_ENCRYPTED_PAGE_SIZE);
         assert_eq!(decrypted, page_data);
     }
 }