Support 192 and 256-bit keys for AES

playfulFence · Mar 20, 2024 · b4c05b8 · b4c05b8
1 parent 6d9048a
commit b4c05b8
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 11 deletions.
diff --git a/esp-hal/src/aes/mod.rs b/esp-hal/src/aes/mod.rs
@@ -124,8 +124,10 @@ const ALIGN_SIZE: usize = core::mem::size_of::<u32>();
 
 pub enum Mode {
     Encryption128 = 0,
+    Encryption192 = 1,
     Encryption256 = 2,
     Decryption128 = 4,
+    Decryption192 = 5,
     Decryption256 = 6,
 }
 
@@ -150,7 +152,14 @@ impl<'d> Aes<'d> {
     }
 
     /// Encrypts/Decrypts the given buffer based on `mode` parameter
-    pub fn process(&mut self, block: &mut [u8; 16], mode: Mode, key: &[u8; 16]) {
+    pub fn process(&mut self, block: &mut [u8; 16], mode: Mode, key: &[u8]) {
+        if key.len() != 16
+            || (cfg!(any(feature = "esp32", feature = "esp32s2")) && key.len() != 24)
+            || key.len() != 32
+        {
+            panic!("Invalid key size");
+        }
+
         self.write_key(key);
         self.set_mode(mode as u8);
         self.set_block(block);
@@ -392,7 +401,7 @@ pub mod dma {
             read_buffer: &'t mut RXBUF,
             mode: Mode,
             cipher_mode: CipherMode,
-            key: [u8; 16],
+            key: &[u8],
         ) -> Result<AesDmaTransferRxTx<'t, 'd, C>, crate::dma::DmaError>
         where
             TXBUF: ReadBuffer<Word = u8>,
@@ -423,7 +432,7 @@ pub mod dma {
             read_buffer_len: usize,
             mode: Mode,
             cipher_mode: CipherMode,
-            key: [u8; 16],
+            key: &[u8],
         ) -> Result<(), crate::dma::DmaError> {
             // AES has to be restarted after each calculation
             self.reset_aes();

diff --git a/examples/src/bin/aes.rs b/examples/src/bin/aes.rs
@@ -73,14 +73,14 @@ fn main() -> ! {
     );
     let sw_decrypted = block;
 
-    assert!(eq(&sw_encrypted.into(), &hw_encrypted));
-    assert!(eq(&sw_decrypted.into(), &hw_decrypted));
+    assert!(eq(&sw_encrypted, &hw_encrypted));
+    assert!(eq(&sw_decrypted, &hw_decrypted));
 
     println!("done");
 
     loop {}
 }
 
-fn eq(slice1: &[u8; 16], slice2: &[u8; 16]) -> bool {
+fn eq(slice1: &[u8], slice2: &[u8]) -> bool {
     slice1.iter().zip(slice2.iter()).all(|(a, b)| a == b)
 }
diff --git a/examples/src/bin/aes_dma.rs b/examples/src/bin/aes_dma.rs
@@ -50,7 +50,7 @@ fn main() -> ! {
             &mut output,
             Mode::Encryption128,
             CipherMode::Ecb,
-            keytext,
+            &keytext,
         )
         .unwrap();
     transfer.wait().unwrap();
@@ -72,7 +72,7 @@ fn main() -> ! {
             &mut output,
             Mode::Decryption128,
             CipherMode::Ecb,
-            keytext,
+            &keytext,
         )
         .unwrap();
     transfer.wait().unwrap();
@@ -109,13 +109,13 @@ fn main() -> ! {
     );
     let sw_decrypted = block.clone();
 
-    assert!(eq(&sw_encrypted.into(), &hw_encrypted));
-    assert!(eq(&sw_decrypted.into(), &hw_decrypted));
+    assert!(eq(&sw_encrypted, &hw_encrypted));
+    assert!(eq(&sw_decrypted, &hw_decrypted));
 
     println!("done");
     loop {}
 }
 
-fn eq(slice1: &[u8; 16], slice2: &[u8; 16]) -> bool {
+fn eq(slice1: &[u8], slice2: &[u8]) -> bool {
     slice1.iter().zip(slice2.iter()).all(|(a, b)| a == b)
 }