diff --git a/rust/k210-shared/src/soc/aes.rs b/rust/k210-shared/src/soc/aes.rs
index 24950da..079c879 100644
--- a/rust/k210-shared/src/soc/aes.rs
+++ b/rust/k210-shared/src/soc/aes.rs
@@ -30,30 +30,16 @@ fn write4pad(arr: &mut [u8], ofs: usize, val: u32) {
     arr[ofs..ofs+n].copy_from_slice(&val.to_le_bytes()[0..n]);
 }
 
-/** AES operation (encrypt or decrypt) using hardware engine. Takes a &mut
- * AES as only one operation can be active at a time.
- *
- * Supported modes:
- *
- * Mode Keybits           Extra input              Extra output
- * ---- ---------------   ------------------------ ----------------
- * ECB  128 / 192 / 256   no IV
- * CBC  128 / 192 / 256   128 bit IV
- * GCM  128 / 192 / 256   96 bit IV + ? bytes AAD) 128 bit tag (optional)
- */
-pub fn run(
+fn setup(
         aes: &mut pac::AES,
         cipher_mode: cipher_mode,
         encrypt_sel: encrypt_sel,
         key: &[u8],
         iv: &[u8],
         aad: &[u8],
-        ind: &[u8],
-        outd: &mut [u8],
-        tag: &mut [u8],
-    )
-{
-    assert!(ind.len() != 0);
+        len: usize,
+    ) {
+    assert!(len != 0);
     match cipher_mode {
         cipher_mode::ECB => assert!(iv.len() == 0 && aad.len() == 0),
         cipher_mode::CBC => assert!(iv.len() == 16 && aad.len() == 0),
@@ -89,7 +75,7 @@ pub fn run(
         aes.encrypt_sel.write(|w|
             w.encrypt_sel().variant(encrypt_sel));
         aes.aad_num.write(|w| w.bits((aad.len() as u32).wrapping_sub(1)));
-        aes.pc_num.write(|w| w.bits((ind.len() as u32).wrapping_sub(1)));
+        aes.pc_num.write(|w| w.bits((len as u32).wrapping_sub(1)));
 
         // Turn on engine
         aes.en.write(|w|
@@ -103,21 +89,16 @@ pub fn run(
             aes.aad_data.write(|w| w.bits(read4pad(aad, i * 4)));
         }
 
-        // Send and receive plaintext/ciphertext 
-        let mut iptr = 0;
-        let mut optr = 0;
-        while optr < ind.len() {
-            while iptr < ind.len() && aes.data_in_flag.read().data_in_flag() == DATA_IN_FLAG_A::CAN_INPUT {
-                aes.text_data.write(|w| w.bits(read4pad(ind, iptr)));
-                iptr += 4;
-            }
-            while aes.data_out_flag.read().data_out_flag() == DATA_OUT_FLAG_A::CAN_OUTPUT {
-                write4pad(outd, optr, aes.out_data.read().bits());
-                optr += 4;
-            }
-        }
+    }
+}
 
-        if cipher_mode == cipher_mode::GCM && tag.len() != 0 {
+fn finish(
+        aes: &mut pac::AES,
+        cipher_mode: cipher_mode,
+        tag: &mut [u8],
+    ) {
+    unsafe {
+        if cipher_mode == cipher_mode::GCM {
             // Read and store tag, if requested
             // TODO: the engine also supports writing a tag through gcm_in_tag
             // and verifying it, presumably in linear time.
@@ -134,9 +115,11 @@ pub fn run(
                 atomic::compiler_fence(Ordering::SeqCst)
             }
 
-            for i in 0..4 {
-                let val = aes.gcm_out_tag[3 - i].read().bits();
-                tag[i*4..i*4+4].copy_from_slice(&val.to_be_bytes());
+            if tag.len() != 0 {
+                for i in 0..4 {
+                    let val = aes.gcm_out_tag[3 - i].read().bits();
+                    tag[i*4..i*4+4].copy_from_slice(&val.to_be_bytes());
+                }
             }
         }
 
@@ -146,3 +129,48 @@ pub fn run(
         }
     }
 }
+
+/** AES operation (encrypt or decrypt) using hardware engine. Takes a &mut
+ * AES as only one operation can be active at a time.
+ *
+ * Supported modes:
+ *
+ * Mode Keybits           Extra input              Extra output
+ * ---- ---------------   ------------------------ ----------------
+ * ECB  128 / 192 / 256   no IV
+ * CBC  128 / 192 / 256   128 bit IV
+ * GCM  128 / 192 / 256   96 bit IV + ? bytes AAD) 128 bit tag (optional)
+ */
+pub fn run(
+        aes: &mut pac::AES,
+        cipher_mode: cipher_mode,
+        encrypt_sel: encrypt_sel,
+        key: &[u8],
+        iv: &[u8],
+        aad: &[u8],
+        ind: &[u8],
+        outd: &mut [u8],
+        tag: &mut [u8],
+    )
+{
+    setup(aes, cipher_mode, encrypt_sel, key, iv, aad, ind.len());
+
+    unsafe {
+        // Send and receive plaintext/ciphertext
+        let mut iptr = 0;
+        let mut optr = 0;
+        while optr < ind.len() {
+            while iptr < ind.len() && aes.data_in_flag.read().data_in_flag() == DATA_IN_FLAG_A::CAN_INPUT {
+                aes.text_data.write(|w| w.bits(read4pad(ind, iptr)));
+                iptr += 4;
+            }
+            while aes.data_out_flag.read().data_out_flag() == DATA_OUT_FLAG_A::CAN_OUTPUT {
+                write4pad(outd, optr, aes.out_data.read().bits());
+                optr += 4;
+            }
+        }
+
+    }
+
+    finish(aes, cipher_mode, tag);
+}