From 1e731872787d90a0a9c8496537470c9f30df629d Mon Sep 17 00:00:00 2001
From: Sergio Gasquez <sergio.gasquez@gmail.com>
Date: Thu, 11 Apr 2024 14:42:28 +0200
Subject: [PATCH] test: Cover all aes modes

---
 esp-hal/src/aes/mod.rs |  4 ++
 hil-test/tests/aes.rs  | 83 +++++++++++++++++++++++++++++++++++++++---
 2 files changed, 82 insertions(+), 5 deletions(-)

diff --git a/esp-hal/src/aes/mod.rs b/esp-hal/src/aes/mod.rs
index acae07e9702..d2f12ac6a70 100644
--- a/esp-hal/src/aes/mod.rs
+++ b/esp-hal/src/aes/mod.rs
@@ -124,8 +124,12 @@ const ALIGN_SIZE: usize = core::mem::size_of::<u32>();
 
 pub enum Mode {
     Encryption128 = 0,
+    #[cfg(any(esp32, esp32s2))]
+    Encryption192 = 1,
     Encryption256 = 2,
     Decryption128 = 4,
+    #[cfg(any(esp32, esp32s2))]
+    Decryption192 = 5,
     Decryption256 = 6,
 }
 
diff --git a/hil-test/tests/aes.rs b/hil-test/tests/aes.rs
index 139772dbb8b..2d6e3d6f1a9 100644
--- a/hil-test/tests/aes.rs
+++ b/hil-test/tests/aes.rs
@@ -37,7 +37,7 @@ mod tests {
     }
 
     #[test]
-    fn test_aes_encryption(mut ctx: Context<'static>) {
+    fn test_aes_128_encryption(mut ctx: Context<'static>) {
         let keytext = "SUp4SeCp@sSw0rd".as_bytes();
         let plaintext = "message".as_bytes();
         let encrypted_message = [
@@ -45,11 +45,9 @@ mod tests {
             0x31, 0x96,
         ];
 
-        // create an array with aes128 key size
         let mut keybuf = [0_u8; 16];
         keybuf[..keytext.len()].copy_from_slice(keytext);
 
-        // create an array with aes block size
         let mut block_buf = [0_u8; 16];
         block_buf[..plaintext.len()].copy_from_slice(plaintext);
 
@@ -59,7 +57,7 @@ mod tests {
     }
 
     #[test]
-    fn test_aes_decryption(mut ctx: Context<'static>) {
+    fn test_aes_128_decryption(mut ctx: Context<'static>) {
         let keytext = "SUp4SeCp@sSw0rd".as_bytes();
         let plaintext = "message".as_bytes();
         let mut encrypted_message = [
@@ -67,7 +65,6 @@ mod tests {
             0x31, 0x96,
         ];
 
-        // create an array with aes128 key size
         let mut keybuf = [0_u8; 16];
         keybuf[..keytext.len()].copy_from_slice(keytext);
 
@@ -75,4 +72,80 @@ mod tests {
             .process(&mut encrypted_message, Mode::Decryption128, &keybuf);
         assert_eq!(&encrypted_message[..plaintext.len()], plaintext);
     }
+
+    #[test]
+    #[cfg(any(feature = "esp32", feature = "esp32s2"))]
+    fn test_aes_192_encryption(mut ctx: Context<'static>) {
+        let keytext = "SUp4SeCp@sSw0rd".as_bytes();
+        let plaintext = "message".as_bytes();
+        let encrypted_message = [
+            0x79, 0x88, 0x3f, 0x9d, 0x67, 0x27, 0xf4, 0x18, 0x3, 0xe3, 0xc6, 0x6a, 0x2e, 0x76,
+            0xb6, 0xf7,
+        ];
+
+        let mut keybuf = [0_u8; 16];
+        keybuf[..keytext.len()].copy_from_slice(keytext);
+
+        let mut block_buf = [0_u8; 16];
+        block_buf[..plaintext.len()].copy_from_slice(plaintext);
+
+        let mut block = block_buf.clone();
+        ctx.aes.process(&mut block, Mode::Encryption192, &keybuf);
+        assert_eq!(block, encrypted_message);
+    }
+
+    #[test]
+    #[cfg(any(feature = "esp32", feature = "esp32s2"))]
+    fn test_aes_192_decryption(mut ctx: Context<'static>) {
+        let keytext = "SUp4SeCp@sSw0rd".as_bytes();
+        let plaintext = "message".as_bytes();
+        let mut encrypted_message = [
+            0x79, 0x88, 0x3f, 0x9d, 0x67, 0x27, 0xf4, 0x18, 0x3, 0xe3, 0xc6, 0x6a, 0x2e, 0x76,
+            0xb6, 0xf7,
+        ];
+
+        let mut keybuf = [0_u8; 16];
+        keybuf[..keytext.len()].copy_from_slice(keytext);
+
+        ctx.aes
+            .process(&mut encrypted_message, Mode::Decryption192, &keybuf);
+        assert_eq!(&encrypted_message[..plaintext.len()], plaintext);
+    }
+
+    #[test]
+    fn test_aes_256_encryption(mut ctx: Context<'static>) {
+        let keytext = "SUp4SeCp@sSw0rd".as_bytes();
+        let plaintext = "message".as_bytes();
+        let encrypted_message = [
+            0x0, 0x63, 0x3f, 0x2, 0xa4, 0x53, 0x9, 0x72, 0x20, 0x6d, 0xc9, 0x8, 0x7c, 0xe5, 0xfd,
+            0xc,
+        ];
+
+        let mut keybuf = [0_u8; 16];
+        keybuf[..keytext.len()].copy_from_slice(keytext);
+
+        let mut block_buf = [0_u8; 16];
+        block_buf[..plaintext.len()].copy_from_slice(plaintext);
+
+        let mut block = block_buf.clone();
+        ctx.aes.process(&mut block, Mode::Encryption256, &keybuf);
+        assert_eq!(block, encrypted_message);
+    }
+
+    #[test]
+    fn test_aes_256_decryption(mut ctx: Context<'static>) {
+        let keytext = "SUp4SeCp@sSw0rd".as_bytes();
+        let plaintext = "message".as_bytes();
+        let mut encrypted_message = [
+            0x0, 0x63, 0x3f, 0x2, 0xa4, 0x53, 0x9, 0x72, 0x20, 0x6d, 0xc9, 0x8, 0x7c, 0xe5, 0xfd,
+            0xc,
+        ];
+
+        let mut keybuf = [0_u8; 16];
+        keybuf[..keytext.len()].copy_from_slice(keytext);
+
+        ctx.aes
+            .process(&mut encrypted_message, Mode::Decryption256, &keybuf);
+        assert_eq!(&encrypted_message[..plaintext.len()], plaintext);
+    }
 }