From eb6126382a62be655a34a208e91a48b7cd5f6dd0 Mon Sep 17 00:00:00 2001
From: rahulKadam1264 <rahulkadam1264@gmail.com>
Date: Sat, 25 Jan 2025 19:44:10 +0530
Subject: [PATCH] Fixes #532

---
 fastcrypto/Cargo.toml             |  3 +-
 fastcrypto/src/aes.rs             | 58 +++++++++++++++++++++++++++++++
 fastcrypto/src/lib.rs             |  8 +++++
 fastcrypto/src/tests/aes_tests.rs | 31 +++++++++++++++++
 4 files changed, 99 insertions(+), 1 deletion(-)

diff --git a/fastcrypto/Cargo.toml b/fastcrypto/Cargo.toml
index 2908128570..a20a4df0bd 100644
--- a/fastcrypto/Cargo.toml
+++ b/fastcrypto/Cargo.toml
@@ -101,7 +101,8 @@ name = "hash"
 harness = false
 
 [features]
-default = []
+default = ["std"]
+no_std = []
 
 # Allow copying keys
 copy_key = []
diff --git a/fastcrypto/src/aes.rs b/fastcrypto/src/aes.rs
index 7b5bc5c0ae..08c89dd272 100644
--- a/fastcrypto/src/aes.rs
+++ b/fastcrypto/src/aes.rs
@@ -241,6 +241,64 @@ where
     }
 }
 
+/// Unified decryption interface for on-chain use.
+/// 
+/// # Parameters
+/// - `key`: The AES key as a byte slice.
+/// - `iv`: The initialization vector as a byte slice.
+/// - `ciphertext`: The encrypted data.
+/// - `mode`: The AES mode to use (e.g., "AES128_CTR", "AES256_CBC").
+/// 
+/// # Returns
+/// - `Ok(Vec<u8>)`: The decrypted plaintext.
+/// - `Err(FastCryptoError)`: If decryption fails.
+///
+/// # Example
+/// ```
+/// use fastcrypto::aes::aes_decrypt_on_chain;
+///
+/// let key = [0u8; 16]; // Example key
+/// let iv = [0u8; 16]; // Example IV
+/// let ciphertext = vec![0x1f, 0x2e, 0x3d]; // Example ciphertext
+///
+/// let plaintext = aes_decrypt_on_chain(&key, &iv, &ciphertext, "AES128_CTR").unwrap();
+/// ```
+
+pub fn aes_decrypt_on_chain(
+    key: &[u8],
+    iv: &[u8],
+    ciphertext: &[u8],
+    mode: &str,
+) -> Result<Vec<u8>, FastCryptoError> {
+    match mode {
+        "AES128_CTR" => {
+            let aes_key = AesKey::<typenum::U16>::from_bytes(key)?;
+            let iv = InitializationVector::<typenum::U16>::from_bytes(iv)?;
+            let cipher = Aes128Ctr::new(aes_key);
+            cipher.decrypt(&iv, ciphertext)
+        }
+        "AES256_CTR" => {
+            let aes_key = AesKey::<typenum::U32>::from_bytes(key)?;
+            let iv = InitializationVector::<typenum::U16>::from_bytes(iv)?;
+            let cipher = Aes256Ctr::new(aes_key);
+            cipher.decrypt(&iv, ciphertext)
+        }
+        "AES128_CBC" => {
+            let aes_key = AesKey::<typenum::U16>::from_bytes(key)?;
+            let iv = InitializationVector::<typenum::U16>::from_bytes(iv)?;
+            let cipher = Aes128CbcPkcs7::new(aes_key);
+            cipher.decrypt(&iv, ciphertext)
+        }
+        "AES256_CBC" => {
+            let aes_key = AesKey::<typenum::U32>::from_bytes(key)?;
+            let iv = InitializationVector::<typenum::U16>::from_bytes(iv)?;
+            let cipher = Aes256CbcPkcs7::new(aes_key);
+            cipher.decrypt(&iv, ciphertext)
+        }
+        _ => Err(FastCryptoError::InvalidInput),
+    }
+}
+
 /// AES128 in CBC-mode using PKCS #7 padding.
 pub type Aes128CbcPkcs7 = AesCbc<aes::Aes128, aes::cipher::block_padding::Pkcs7>;
 
diff --git a/fastcrypto/src/lib.rs b/fastcrypto/src/lib.rs
index df3e8967cf..b56684b1a3 100644
--- a/fastcrypto/src/lib.rs
+++ b/fastcrypto/src/lib.rs
@@ -8,6 +8,14 @@
     rust_2021_compatibility
 )]
 
+#![cfg_attr(not(feature = "std"), no_std)]
+
+#[cfg(not(feature = "std"))]
+extern crate alloc;
+
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
+
 #[cfg(test)]
 #[path = "tests/ed25519_tests.rs"]
 pub mod ed25519_tests;
diff --git a/fastcrypto/src/tests/aes_tests.rs b/fastcrypto/src/tests/aes_tests.rs
index 7703b029ab..0f8c8b2412 100644
--- a/fastcrypto/src/tests/aes_tests.rs
+++ b/fastcrypto/src/tests/aes_tests.rs
@@ -221,3 +221,34 @@ fn test_sk_zeroization_on_drop() {
         }
     }
 }
+
+#[test]
+    fn test_aes_decrypt_on_chain_valid() {
+        let plaintext = b"Hello, on-chain!";
+        let key = AesKey::<typenum::U32>::generate(&mut thread_rng());
+        let iv = InitializationVector::<typenum::U16>::generate(&mut thread_rng());
+
+        let cipher = Aes256Ctr::new(key.clone());
+        let ciphertext = cipher.encrypt(&iv, plaintext);
+
+        let decrypted = aes_decrypt_on_chain(
+            &key.as_bytes(),
+            &iv.as_bytes(),
+            &ciphertext,
+            "AES256_CTR",
+        )
+        .unwrap();
+
+        assert_eq!(decrypted, plaintext);
+    }
+
+    #[test]
+    fn test_aes_decrypt_on_chain_invalid_mode() {
+        let key = [0u8; 16];
+        let iv = [0u8; 16];
+        let ciphertext = vec![0x1f, 0x2e, 0x3d];
+
+        let result = aes_decrypt_on_chain(&key, &iv, &ciphertext, "INVALID_MODE");
+
+        assert!(result.is_err());
+    }