From f11e90c94db1b3c629c003f0e6a4bcdf17a1ceb0 Mon Sep 17 00:00:00 2001
From: bit-aloo <sshourya17@gmail.com>
Date: Mon, 1 Sep 2025 16:18:22 +0530
Subject: [PATCH] refactor Aegis256Cipher to implement AeadCipher

---
 core/storage/encryption.rs | 36 +++++++++++++++++++-----------------
 1 file changed, 19 insertions(+), 17 deletions(-)

diff --git a/core/storage/encryption.rs b/core/storage/encryption.rs
index a37c2eaaf..e48d5dafb 100644
--- a/core/storage/encryption.rs
+++ b/core/storage/encryption.rs
@@ -90,7 +90,7 @@ pub trait AeadCipher {
 // 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 {
@@ -98,34 +98,36 @@ 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())