DMA-enabled work-queue based AES driver (#3897)

* AES-DMA work queue backend Co-authored-by: Dominic Fischer <14130965+Dominaezzz@users.noreply.github.com> * Clean up manual subslicing * Don't copy in write_iv * Add more to why we can use copy_from_nonoverlapping * Replace buffers with NoBuffer * Move buffers to backend * Volatile-ly zero the key * Make saving state the user's responsibility * Ensure data is aligned on esp32 * Also make sure the DMA finishes in AesTransfer * Deduplicate * Fix paperwork * Use the handler attribute * Remove unused method * Update esp-hal/MIGRATING-1.0.0-rc.0.md Co-authored-by: Dominic Fischer <14130965+Dominaezzz@users.noreply.github.com> * Fix build after rebase * Add empty Final to NoBuffer * Use the move api internally * Make () the view type --------- Co-authored-by: Dominic Fischer <14130965+Dominaezzz@users.noreply.github.com>
2025-09-26 20:00:32 +00:00 · 2025-08-12 14:19:24 +02:00 · 2025-08-12 14:19:24 +02:00 · fab1221894
commit fab1221894
parent ca4e28da4a
17 changed files with 1305 additions and 135 deletions
--- a/esp-hal/CHANGELOG.md
+++ b/esp-hal/CHANGELOG.md
@ -17,6 +17,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `aes::cipher_modes`, `aes::CipherModeState` for constructing `AesContext`s (#3895)
 - `DmaTxBuffer` and `DmaRxBuffer` now have a `Final` associated type. (#3923)
 - `RsaBackend, RsaContext`: Work-queue based RSA driver (#3910)
+- `aes::{AesBackend, AesContext, dma::AesDmaBackend}`: Work-queue based AES driver (#3880, #3897)
+- `aes::cipher_modes`, `aes::CipherState` for constructing `AesContext`s (#3895)
+- `aes::dma::DmaCipherState` so that `AesDma` can properly support cipher modes that require state (IV, nonce, etc.) (#3897)

 ### Changed

@ -24,6 +27,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Update `embassy-usb` to v0.5.0 (#3848)
 - `aes::Key` variants have been renamed from bytes to bits (e.g. `Key16 -> Key128`) (#3845)
 - `aes::Mode` has been replaced by `Operation`. The key length is now solely determined by the key. (#3882)
+- `AesDma::process` now takes `DmaCipherState` instead of `CipherMode`. (#3897)
 - `Aes::process` has been split into `Aes::encrypt` and `Aes::decrypt` (#3882)
 - Blocking RMT transactions can now be `poll`ed without blocking, returning whether they have completed. (#3716)
 - RISC-V: Interrupt handler don't get a TrapFrame passed in anymore (#3903)
--- a/esp-hal/MIGRATING-1.0.0-rc.0.md
+++ b/esp-hal/MIGRATING-1.0.0-rc.0.md
@ -27,11 +27,12 @@ The previous way to obtain RNG object has changed like so:

 ## AES changes

-The `esp_hal::aes::Aes` driver has been slightly reworked:
+The `esp_hal::aes::Aes` and `esp_hal::aes::AesDma` drivers has been slightly reworked:

 - `Mode` has been replaced by `Operation`. Operation has `Encrypt` and `Decrypt` variants, but the key length is no longer part of the enum. The key length is specified by the key. AesDma now takes this `Operation`.
 - `Aes::process` has been split into `encrypt` and `decrypt`. These functions no longer take a mode parameter.
- `AesDma::write_block` has been removed.
+- `AesDma::write_block` and `AesDma::write_key` have been removed.
+- `AesDma::process` now takes `DmaCipherState` which includes information to initialize the block cipher mode of operation.

 ```diff
 -aes.process(block, Mode::Encrypt128, key);
@ -40,6 +41,26 @@ The `esp_hal::aes::Aes` driver has been slightly reworked:
 -aes.process(block, Mode::Decrypt256, key);
 +aes.decrypt(block, key_32_bytes);
 ```
+
+```diff
+use esp_hal::aes::dma::DmaCipherState;
+use esp_hal::aes::cipher_modes::Ecb;
+
+ let transfer = aes_dma
+     .process(
+         1,
+         output,
+         input,
+         Operation::Encrypt,
+-        CipherMode::Ecb,
+        &DmaCipherState::from(Ecb),
+         key,
+     )
+     .map_err(|e| e.0)
+     .unwrap();
+ (aes_dma, output, input) = transfer.wait();
+```
+
 ## ISR Callback Changes

 Previously callbacks were of type `extern "C" fn()`, now they are `IsrCallback`. In most places no changes are needed but when using `bind_interrupt` directly
--- a/esp-hal/src/aes/cipher_modes.rs
+++ b/esp-hal/src/aes/cipher_modes.rs
@ -52,6 +52,14 @@ impl UnsafeCryptoBuffers {
    pub fn in_place(&self) -> bool {
        self.input.addr() == self.output.addr()
    }
+
+    #[cfg(aes_dma)]
+    pub(crate) unsafe fn byte_add(self, bytes: usize) -> Self {
+        UnsafeCryptoBuffers {
+            input: unsafe { self.input.byte_add(bytes) },
+            output: unsafe { self.output.byte_add(bytes) },
+        }
+    }
 }

 /// Electronic codebook mode.
@ -173,6 +181,20 @@ impl Ofb {
        });
        self.offset = offset;
    }
+
+    #[cfg(aes_dma)]
+    pub(super) fn flush(&mut self, buffer: UnsafeCryptoBuffers) -> usize {
+        let mut offset = self.offset;
+        buffer
+            .first_n((BLOCK_SIZE - offset) % BLOCK_SIZE)
+            .for_data_chunks(1, |input, output, _| {
+                unsafe { output.write(input.read() ^ self.iv[offset]) };
+                offset += 1;
+            });
+        let flushed = offset - self.offset;
+        self.offset = offset % BLOCK_SIZE;
+        flushed
+    }
 }

 /// Counter mode.
@ -224,6 +246,20 @@ impl Ctr {
        });
        self.offset = offset;
    }
+
+    #[cfg(aes_dma)]
+    pub(super) fn flush(&mut self, buffer: UnsafeCryptoBuffers) -> usize {
+        let mut offset = self.offset;
+        buffer
+            .first_n((BLOCK_SIZE - offset) % BLOCK_SIZE)
+            .for_data_chunks(1, |plaintext, ciphertext, _| {
+                unsafe { ciphertext.write(plaintext.read() ^ self.buffer[offset]) };
+                offset += 1;
+            });
+        let flushed = offset - self.offset;
+        self.offset = offset % BLOCK_SIZE;
+        flushed
+    }
 }

 /// Cipher feedback with 8-bit shift mode.
@ -332,6 +368,41 @@ impl Cfb128 {
        });
        self.offset = offset;
    }
+
+    #[cfg(aes_dma)]
+    pub(super) fn flush_encrypt(&mut self, buffer: UnsafeCryptoBuffers) -> usize {
+        let mut offset = self.offset;
+        buffer
+            .first_n((BLOCK_SIZE - offset) % BLOCK_SIZE)
+            .for_data_chunks(1, |plaintext, ciphertext, _| {
+                unsafe {
+                    self.iv[offset] ^= plaintext.read();
+                    ciphertext.write(self.iv[offset]);
+                }
+                offset += 1;
+            });
+        let flushed = offset - self.offset;
+        self.offset = offset % BLOCK_SIZE;
+        flushed
+    }
+
+    #[cfg(aes_dma)]
+    pub(super) fn flush_decrypt(&mut self, buffer: UnsafeCryptoBuffers) -> usize {
+        let mut offset = self.offset;
+        buffer
+            .first_n((BLOCK_SIZE - offset) % BLOCK_SIZE)
+            .for_data_chunks(1, |ciphertext, plaintext, _| {
+                unsafe {
+                    let c = ciphertext.read();
+                    plaintext.write(self.iv[offset] ^ c);
+                    self.iv[offset] = c;
+                }
+                offset += 1;
+            });
+        let flushed = offset - self.offset;
+        self.offset = offset % BLOCK_SIZE;
+        flushed
+    }
 }

 // Utilities
@ -352,6 +423,15 @@ impl UnsafeCryptoBuffers {
            cb(input, output, len)
        }
    }
+
+    #[cfg(aes_dma)]
+    fn first_n(self, n: usize) -> UnsafeCryptoBuffers {
+        let len = n.min(self.input.len());
+        Self {
+            input: NonNull::slice_from_raw_parts(self.input.cast(), len),
+            output: NonNull::slice_from_raw_parts(self.output.cast(), len),
+        }
+    }
 }

 fn pointer_chunks<T>(
--- a/esp-hal/src/aes/mod.rs
+++ b/esp-hal/src/aes/mod.rs
--- a/esp-hal/src/dma/buffers.rs
+++ b/esp-hal/src/dma/buffers.rs
@ -1,6 +1,8 @@
+#[cfg(psram_dma)]
+use core::ops::Range;
 use core::{
    ops::{Deref, DerefMut},
-    ptr::null_mut,
+    ptr::{NonNull, null_mut},
 };

 use super::*;
@ -1580,3 +1582,368 @@ impl DerefMut for DmaLoopBuf {
        self.buffer
    }
 }
+
+/// A Preparation that masks itself as a DMA buffer.
+///
+/// Fow low level use, where none of the pre-made buffers really fit.
+///
+/// This type likely never should be visible outside of esp-hal.
+pub(crate) struct NoBuffer(Preparation);
+impl NoBuffer {
+    fn prep(&self) -> Preparation {
+        Preparation {
+            start: self.0.start,
+            direction: self.0.direction,
+            #[cfg(psram_dma)]
+            accesses_psram: self.0.accesses_psram,
+            burst_transfer: self.0.burst_transfer,
+            check_owner: self.0.check_owner,
+            auto_write_back: self.0.auto_write_back,
+        }
+    }
+}
+unsafe impl DmaTxBuffer for NoBuffer {
+    type View = ();
+    type Final = ();
+
+    fn prepare(&mut self) -> Preparation {
+        self.prep()
+    }
+
+    fn into_view(self) -> Self::View {}
+    fn from_view(_view: Self::View) {}
+}
+unsafe impl DmaRxBuffer for NoBuffer {
+    type View = ();
+    type Final = ();
+
+    fn prepare(&mut self) -> Preparation {
+        self.prep()
+    }
+
+    fn into_view(self) -> Self::View {}
+    fn from_view(_view: Self::View) {}
+}
+
+/// Prepares data unsafely to be transmitted via DMA.
+///
+/// `block_size` is the requirement imposed by the peripheral that receives the data. It
+/// ensures that the DMA will not try to copy a partial block, which would cause the RX DMA (that
+/// moves results back into RAM) to never complete.
+///
+/// The function returns the DMA buffer, and the number of bytes that will be transferred.
+///
+/// # Safety
+///
+/// The caller must keep all its descriptors and the buffers they
+/// point to valid while the buffer is being transferred.
+#[cfg_attr(not(aes_dma), expect(unused))]
+pub(crate) unsafe fn prepare_for_tx(
+    descriptors: &mut [DmaDescriptor],
+    mut data: NonNull<[u8]>,
+    block_size: usize,
+) -> Result<(NoBuffer, usize), DmaError> {
+    let alignment =
+        BurstConfig::DEFAULT.min_alignment(unsafe { data.as_ref() }, TransferDirection::Out);
+
+    if !data.addr().get().is_multiple_of(alignment) {
+        // ESP32 has word alignment requirement on the TX descriptors, too.
+        return Err(DmaError::InvalidAlignment(DmaAlignmentError::Address));
+    }
+
+    // Whichever is stricter, data location or peripheral requirements.
+    //
+    // This ensures that the RX DMA, if used, can transfer the returned number of bytes using at
+    // most N+2 descriptors. While the hardware doesn't require this on the TX DMA side, (the TX DMA
+    // can, except on the ESP32, transfer any amount of data), it makes usage MUCH simpler.
+    let alignment = alignment.max(block_size);
+    let chunk_size = 4096 - alignment;
+
+    let data_len = data.len().min(chunk_size * descriptors.len());
+
+    cfg_if::cfg_if! {
+        if #[cfg(psram_dma)] {
+            let data_addr = data.addr().get();
+            let data_in_psram = crate::psram::psram_range().contains(&data_addr);
+
+            // Make sure input data is in PSRAM instead of cache
+            if data_in_psram {
+                unsafe { crate::soc::cache_writeback_addr(data_addr as u32, data_len as u32) };
+            }
+        }
+    }
+
+    let mut descriptors = unwrap!(DescriptorSet::new(descriptors));
+    // TODO: it would be best if this function returned the amount of data that could be linked
+    // up.
+    unwrap!(descriptors.link_with_buffer(unsafe { data.as_mut() }, chunk_size));
+    unwrap!(descriptors.set_tx_length(data_len, chunk_size));
+
+    for desc in descriptors.linked_iter_mut() {
+        desc.reset_for_tx(desc.next.is_null());
+    }
+
+    Ok((
+        NoBuffer(Preparation {
+            start: descriptors.head(),
+            direction: TransferDirection::Out,
+            burst_transfer: BurstConfig::DEFAULT,
+            check_owner: None,
+            auto_write_back: true,
+            #[cfg(psram_dma)]
+            accesses_psram: data_in_psram,
+        }),
+        data_len,
+    ))
+}
+
+/// Prepare buffers to receive data from DMA.
+///
+/// The function returns the DMA buffer, and the number of bytes that will be transferred.
+///
+/// # Safety
+///
+/// The caller must keep all its descriptors and the buffers they
+/// point to valid while the buffer is being transferred.
+#[cfg_attr(not(aes_dma), expect(unused))]
+pub(crate) unsafe fn prepare_for_rx(
+    descriptors: &mut [DmaDescriptor],
+    #[cfg(psram_dma)] align_buffers: &mut [Option<ManualWritebackBuffer>; 2],
+    mut data: NonNull<[u8]>,
+) -> (NoBuffer, usize) {
+    let chunk_size =
+        BurstConfig::DEFAULT.max_chunk_size_for(unsafe { data.as_ref() }, TransferDirection::In);
+
+    // The data we have to process may not be appropriate for the DMA:
+    // - it may be improperly aligned for PSRAM
+    // - it may not have a length that is a multiple of the external memory block size
+
+    cfg_if::cfg_if! {
+        if #[cfg(psram_dma)] {
+            let data_addr = data.addr().get();
+            let data_in_psram = crate::psram::psram_range().contains(&data_addr);
+        } else {
+            let data_in_psram = false;
+        }
+    }
+
+    let mut descriptors = unwrap!(DescriptorSet::new(descriptors));
+    let data_len = if data_in_psram {
+        cfg_if::cfg_if! {
+            if #[cfg(psram_dma)] {
+                // This could use a better API, but right now we'll have to build the descriptor list by
+                // hand.
+                let consumed_bytes = build_descriptor_list_for_psram(
+                    &mut descriptors,
+                    align_buffers,
+                    data,
+                );
+
+                // Invalidate data written by the DMA
+                unsafe {
+                    crate::soc::cache_invalidate_addr(data_addr as u32, consumed_bytes as u32);
+                }
+
+                consumed_bytes
+            } else {
+                unreachable!()
+            }
+        }
+    } else {
+        // Just set up descriptors as usual
+        let data_len = data.len();
+        unwrap!(descriptors.link_with_buffer(unsafe { data.as_mut() }, chunk_size));
+        unwrap!(descriptors.set_tx_length(data_len, chunk_size));
+
+        data_len
+    };
+
+    for desc in descriptors.linked_iter_mut() {
+        desc.reset_for_rx();
+    }
+
+    (
+        NoBuffer(Preparation {
+            start: descriptors.head(),
+            direction: TransferDirection::In,
+            burst_transfer: BurstConfig::DEFAULT,
+            check_owner: None,
+            auto_write_back: true,
+            #[cfg(psram_dma)]
+            accesses_psram: data_in_psram,
+        }),
+        data_len,
+    )
+}
+
+#[cfg(psram_dma)]
+fn build_descriptor_list_for_psram(
+    descriptors: &mut DescriptorSet<'_>,
+    copy_buffers: &mut [Option<ManualWritebackBuffer>; 2],
+    data: NonNull<[u8]>,
+) -> usize {
+    let data_len = data.len();
+    let data_addr = data.addr().get();
+
+    let min_alignment = ExternalBurstConfig::DEFAULT.min_psram_alignment(TransferDirection::In);
+    let chunk_size = 4096 - min_alignment;
+
+    let mut desciptor_iter = DescriptorChainingIter::new(descriptors.descriptors);
+    let mut copy_buffer_iter = copy_buffers.iter_mut();
+
+    // MIN_LAST_DMA_LEN could make this really annoying, so we're just allocating a bit larger
+    // buffer and shove edge cases into a single one. If we have >24 bytes on the S2, the 2-buffer
+    // alignment algo works fine as one of them can steal 16 bytes, the other will have
+    // MIN_LAST_DMA_LEN data to work with.
+    let has_aligned_data = data_len > BUF_LEN;
+
+    // Calculate byte offset to the start of the buffer
+    let offset = data_addr % min_alignment;
+    let head_to_copy = min_alignment - offset;
+    let head_to_copy = if !has_aligned_data {
+        BUF_LEN
+    } else if head_to_copy > 0 && head_to_copy < MIN_LAST_DMA_LEN {
+        head_to_copy + min_alignment
+    } else {
+        head_to_copy
+    };
+    let head_to_copy = head_to_copy.min(data_len);
+
+    // Calculate last unaligned part
+    let tail_to_copy = (data_len - head_to_copy) % min_alignment;
+    let tail_to_copy = if tail_to_copy > 0 && tail_to_copy < MIN_LAST_DMA_LEN {
+        tail_to_copy + min_alignment
+    } else {
+        tail_to_copy
+    };
+
+    let mut consumed = 0;
+
+    // Align beginning
+    if head_to_copy > 0 {
+        let copy_buffer = unwrap!(copy_buffer_iter.next());
+        let buffer =
+            copy_buffer.insert(ManualWritebackBuffer::new(get_range(data, 0..head_to_copy)));
+
+        let Some(descriptor) = desciptor_iter.next() else {
+            return consumed;
+        };
+        descriptor.set_size(head_to_copy);
+        descriptor.buffer = buffer.buffer_ptr();
+        consumed += head_to_copy;
+    };
+
+    // Chain up descriptors for the main aligned data part.
+    let mut aligned_data = get_range(data, head_to_copy..data.len() - tail_to_copy);
+    while !aligned_data.is_empty() {
+        let Some(descriptor) = desciptor_iter.next() else {
+            return consumed;
+        };
+        let chunk = aligned_data.len().min(chunk_size);
+
+        descriptor.set_size(chunk);
+        descriptor.buffer = aligned_data.cast::<u8>().as_ptr();
+        consumed += chunk;
+        aligned_data = get_range(aligned_data, chunk..aligned_data.len());
+    }
+
+    // Align end
+    if tail_to_copy > 0 {
+        let copy_buffer = unwrap!(copy_buffer_iter.next());
+        let buffer = copy_buffer.insert(ManualWritebackBuffer::new(get_range(
+            data,
+            data.len() - tail_to_copy..data.len(),
+        )));
+
+        let Some(descriptor) = desciptor_iter.next() else {
+            return consumed;
+        };
+        descriptor.set_size(tail_to_copy);
+        descriptor.buffer = buffer.buffer_ptr();
+        consumed += tail_to_copy;
+    }
+
+    consumed
+}
+
+#[cfg(psram_dma)]
+fn get_range(ptr: NonNull<[u8]>, range: Range<usize>) -> NonNull<[u8]> {
+    let len = range.end - range.start;
+    NonNull::slice_from_raw_parts(unsafe { ptr.cast().byte_add(range.start) }, len)
+}
+
+#[cfg(psram_dma)]
+struct DescriptorChainingIter<'a> {
+    /// index of the next element to emit
+    index: usize,
+    descriptors: &'a mut [DmaDescriptor],
+}
+#[cfg(psram_dma)]
+impl<'a> DescriptorChainingIter<'a> {
+    fn new(descriptors: &'a mut [DmaDescriptor]) -> Self {
+        Self {
+            descriptors,
+            index: 0,
+        }
+    }
+
+    fn next(&mut self) -> Option<&'_ mut DmaDescriptor> {
+        if self.index == 0 {
+            self.index += 1;
+            self.descriptors.get_mut(0)
+        } else if self.index < self.descriptors.len() {
+            let index = self.index;
+            self.index += 1;
+
+            // Grab a pointer to the current descriptor.
+            let ptr = &raw mut self.descriptors[index];
+
+            // Link the descriptor to the previous one.
+            self.descriptors[index - 1].next = ptr;
+
+            // Reborrow the pointer so that it doesn't get invalidated by our continued use of the
+            // descriptor reference.
+            Some(unsafe { &mut *ptr })
+        } else {
+            None
+        }
+    }
+}
+
+#[cfg(psram_dma)]
+const MIN_LAST_DMA_LEN: usize = if cfg!(esp32s2) { 5 } else { 1 };
+#[cfg(psram_dma)]
+const BUF_LEN: usize = 16 + 2 * (MIN_LAST_DMA_LEN - 1); // 2x makes aligning short buffers simpler
+
+/// PSRAM helper. DMA can write data of any alignment into this buffer, and it can be written by
+/// the CPU back to PSRAM.
+#[cfg(psram_dma)]
+pub(crate) struct ManualWritebackBuffer {
+    dst_address: NonNull<u8>,
+    buffer: [u8; BUF_LEN],
+    n_bytes: u8,
+}
+
+#[cfg(psram_dma)]
+impl ManualWritebackBuffer {
+    pub fn new(ptr: NonNull<[u8]>) -> Self {
+        assert!(ptr.len() <= BUF_LEN);
+        Self {
+            dst_address: ptr.cast(),
+            buffer: [0; BUF_LEN],
+            n_bytes: ptr.len() as u8,
+        }
+    }
+
+    pub fn write_back(self) {
+        unsafe {
+            self.dst_address
+                .as_ptr()
+                .copy_from(self.buffer.as_ptr(), self.n_bytes as usize);
+        }
+    }
+
+    pub fn buffer_ptr(&self) -> *mut u8 {
+        self.buffer.as_ptr().cast_mut()
+    }
+}
--- a/esp-hal/src/dma/gdma.rs
+++ b/esp-hal/src/dma/gdma.rs
@ -33,6 +33,16 @@ pub struct AnyGdmaChannel<'d> {
    _lifetime: PhantomData<&'d mut ()>,
 }

+impl AnyGdmaChannel<'_> {
+    #[cfg_attr(esp32c2, expect(unused))]
+    pub(crate) unsafe fn clone_unchecked(&self) -> Self {
+        Self {
+            channel: self.channel,
+            _lifetime: PhantomData,
+        }
+    }
+}
+
 impl crate::private::Sealed for AnyGdmaChannel<'_> {}
 impl<'d> DmaChannel for AnyGdmaChannel<'d> {
    type Rx = AnyGdmaRxChannel<'d>;
--- a/esp-metadata-generated/src/_generated_esp32c3.rs
+++ b/esp-metadata-generated/src/_generated_esp32c3.rs
@ -440,9 +440,10 @@ macro_rules! for_each_peripheral {
        _for_each_inner!((GPIO16 <= virtual())); _for_each_inner!((GPIO17 <= virtual()));
        _for_each_inner!((GPIO18 <= virtual())); _for_each_inner!((GPIO19 <= virtual()));
        _for_each_inner!((GPIO20 <= virtual())); _for_each_inner!((GPIO21 <= virtual()));
-        _for_each_inner!((AES <= AES() (unstable))); _for_each_inner!((APB_CTRL <=
-        APB_CTRL() (unstable))); _for_each_inner!((APB_SARADC <= APB_SARADC()
-        (unstable))); _for_each_inner!((ASSIST_DEBUG <= ASSIST_DEBUG() (unstable)));
+        _for_each_inner!((AES <= AES(AES : { bind_peri_interrupt, enable_peri_interrupt,
+        disable_peri_interrupt }) (unstable))); _for_each_inner!((APB_CTRL <= APB_CTRL()
+        (unstable))); _for_each_inner!((APB_SARADC <= APB_SARADC() (unstable)));
+        _for_each_inner!((ASSIST_DEBUG <= ASSIST_DEBUG() (unstable)));
        _for_each_inner!((BB <= BB() (unstable))); _for_each_inner!((DMA <= DMA()
        (unstable))); _for_each_inner!((DS <= DS() (unstable))); _for_each_inner!((EFUSE
        <= EFUSE() (unstable))); _for_each_inner!((EXTMEM <= EXTMEM() (unstable)));
@ -487,9 +488,10 @@ macro_rules! for_each_peripheral {
        (GPIO12 <= virtual()), (GPIO13 <= virtual()), (GPIO14 <= virtual()), (GPIO15 <=
        virtual()), (GPIO16 <= virtual()), (GPIO17 <= virtual()), (GPIO18 <= virtual()),
        (GPIO19 <= virtual()), (GPIO20 <= virtual()), (GPIO21 <= virtual()), (AES <=
-        AES() (unstable)), (APB_CTRL <= APB_CTRL() (unstable)), (APB_SARADC <=
-        APB_SARADC() (unstable)), (ASSIST_DEBUG <= ASSIST_DEBUG() (unstable)), (BB <=
-        BB() (unstable)), (DMA <= DMA() (unstable)), (DS <= DS() (unstable)), (EFUSE <=
+        AES(AES : { bind_peri_interrupt, enable_peri_interrupt, disable_peri_interrupt })
+        (unstable)), (APB_CTRL <= APB_CTRL() (unstable)), (APB_SARADC <= APB_SARADC()
+        (unstable)), (ASSIST_DEBUG <= ASSIST_DEBUG() (unstable)), (BB <= BB()
+        (unstable)), (DMA <= DMA() (unstable)), (DS <= DS() (unstable)), (EFUSE <=
        EFUSE() (unstable)), (EXTMEM <= EXTMEM() (unstable)), (FE <= FE() (unstable)),
        (FE2 <= FE2() (unstable)), (GPIO <= GPIO() (unstable)), (GPIO_SD <= GPIO_SD()
        (unstable)), (HMAC <= HMAC() (unstable)), (I2C_ANA_MST <= I2C_ANA_MST()
--- a/esp-metadata-generated/src/_generated_esp32c6.rs
+++ b/esp-metadata-generated/src/_generated_esp32c6.rs
@ -456,7 +456,8 @@ macro_rules! for_each_peripheral {
        _for_each_inner!((GPIO24 <= virtual())); _for_each_inner!((GPIO25 <= virtual()));
        _for_each_inner!((GPIO26 <= virtual())); _for_each_inner!((GPIO27 <= virtual()));
        _for_each_inner!((GPIO28 <= virtual())); _for_each_inner!((GPIO29 <= virtual()));
-        _for_each_inner!((GPIO30 <= virtual())); _for_each_inner!((AES <= AES()
+        _for_each_inner!((GPIO30 <= virtual())); _for_each_inner!((AES <= AES(AES : {
+        bind_peri_interrupt, enable_peri_interrupt, disable_peri_interrupt })
        (unstable))); _for_each_inner!((APB_SARADC <= APB_SARADC() (unstable)));
        _for_each_inner!((ASSIST_DEBUG <= ASSIST_DEBUG() (unstable)));
        _for_each_inner!((ATOMIC <= ATOMIC() (unstable))); _for_each_inner!((DMA <= DMA()
@ -536,7 +537,8 @@ macro_rules! for_each_peripheral {
        (GPIO20 <= virtual()), (GPIO21 <= virtual()), (GPIO22 <= virtual()), (GPIO23 <=
        virtual()), (GPIO24 <= virtual()), (GPIO25 <= virtual()), (GPIO26 <= virtual()),
        (GPIO27 <= virtual()), (GPIO28 <= virtual()), (GPIO29 <= virtual()), (GPIO30 <=
-        virtual()), (AES <= AES() (unstable)), (APB_SARADC <= APB_SARADC() (unstable)),
+        virtual()), (AES <= AES(AES : { bind_peri_interrupt, enable_peri_interrupt,
+        disable_peri_interrupt }) (unstable)), (APB_SARADC <= APB_SARADC() (unstable)),
        (ASSIST_DEBUG <= ASSIST_DEBUG() (unstable)), (ATOMIC <= ATOMIC() (unstable)),
        (DMA <= DMA() (unstable)), (DS <= DS() (unstable)), (ECC <= ECC() (unstable)),
        (EFUSE <= EFUSE() (unstable)), (EXTMEM <= EXTMEM() (unstable)), (GPIO <= GPIO()
--- a/esp-metadata-generated/src/_generated_esp32h2.rs
+++ b/esp-metadata-generated/src/_generated_esp32h2.rs
@ -438,7 +438,8 @@ macro_rules! for_each_peripheral {
        _for_each_inner!((GPIO14 <= virtual())); _for_each_inner!((GPIO22 <= virtual()));
        _for_each_inner!((GPIO23 <= virtual())); _for_each_inner!((GPIO24 <= virtual()));
        _for_each_inner!((GPIO25 <= virtual())); _for_each_inner!((GPIO26 <= virtual()));
-        _for_each_inner!((GPIO27 <= virtual())); _for_each_inner!((AES <= AES()
+        _for_each_inner!((GPIO27 <= virtual())); _for_each_inner!((AES <= AES(AES : {
+        bind_peri_interrupt, enable_peri_interrupt, disable_peri_interrupt })
        (unstable))); _for_each_inner!((APB_SARADC <= APB_SARADC() (unstable)));
        _for_each_inner!((ASSIST_DEBUG <= ASSIST_DEBUG() (unstable)));
        _for_each_inner!((DMA <= DMA() (unstable))); _for_each_inner!((DS <= DS()
@ -509,9 +510,10 @@ macro_rules! for_each_peripheral {
        virtual()), (GPIO10 <= virtual()), (GPIO11 <= virtual()), (GPIO12 <= virtual()),
        (GPIO13 <= virtual()), (GPIO14 <= virtual()), (GPIO22 <= virtual()), (GPIO23 <=
        virtual()), (GPIO24 <= virtual()), (GPIO25 <= virtual()), (GPIO26 <= virtual()),
-        (GPIO27 <= virtual()), (AES <= AES() (unstable)), (APB_SARADC <= APB_SARADC()
-        (unstable)), (ASSIST_DEBUG <= ASSIST_DEBUG() (unstable)), (DMA <= DMA()
-        (unstable)), (DS <= DS() (unstable)), (ECC <= ECC() (unstable)), (EFUSE <=
+        (GPIO27 <= virtual()), (AES <= AES(AES : { bind_peri_interrupt,
+        enable_peri_interrupt, disable_peri_interrupt }) (unstable)), (APB_SARADC <=
+        APB_SARADC() (unstable)), (ASSIST_DEBUG <= ASSIST_DEBUG() (unstable)), (DMA <=
+        DMA() (unstable)), (DS <= DS() (unstable)), (ECC <= ECC() (unstable)), (EFUSE <=
        EFUSE() (unstable)), (GPIO <= GPIO() (unstable)), (GPIO_SD <= GPIO_SD()
        (unstable)), (HMAC <= HMAC() (unstable)), (HP_APM <= HP_APM() (unstable)),
        (HP_SYS <= HP_SYS() (unstable)), (I2C_ANA_MST <= I2C_ANA_MST() (unstable)), (I2C0
--- a/esp-metadata-generated/src/_generated_esp32s2.rs
+++ b/esp-metadata-generated/src/_generated_esp32s2.rs
@ -479,7 +479,8 @@ macro_rules! for_each_peripheral {
        _for_each_inner!((GPIO40 <= virtual())); _for_each_inner!((GPIO41 <= virtual()));
        _for_each_inner!((GPIO42 <= virtual())); _for_each_inner!((GPIO43 <= virtual()));
        _for_each_inner!((GPIO44 <= virtual())); _for_each_inner!((GPIO45 <= virtual()));
-        _for_each_inner!((GPIO46 <= virtual())); _for_each_inner!((AES <= AES()
+        _for_each_inner!((GPIO46 <= virtual())); _for_each_inner!((AES <= AES(AES : {
+        bind_peri_interrupt, enable_peri_interrupt, disable_peri_interrupt })
        (unstable))); _for_each_inner!((APB_SARADC <= APB_SARADC() (unstable)));
        _for_each_inner!((DEDICATED_GPIO <= DEDICATED_GPIO() (unstable)));
        _for_each_inner!((DS <= DS() (unstable))); _for_each_inner!((EFUSE <= EFUSE()
@ -541,7 +542,8 @@ macro_rules! for_each_peripheral {
        (GPIO34 <= virtual()), (GPIO35 <= virtual()), (GPIO36 <= virtual()), (GPIO37 <=
        virtual()), (GPIO38 <= virtual()), (GPIO39 <= virtual()), (GPIO40 <= virtual()),
        (GPIO41 <= virtual()), (GPIO42 <= virtual()), (GPIO43 <= virtual()), (GPIO44 <=
-        virtual()), (GPIO45 <= virtual()), (GPIO46 <= virtual()), (AES <= AES()
+        virtual()), (GPIO45 <= virtual()), (GPIO46 <= virtual()), (AES <= AES(AES : {
+        bind_peri_interrupt, enable_peri_interrupt, disable_peri_interrupt })
        (unstable)), (APB_SARADC <= APB_SARADC() (unstable)), (DEDICATED_GPIO <=
        DEDICATED_GPIO() (unstable)), (DS <= DS() (unstable)), (EFUSE <= EFUSE()
        (unstable)), (EXTMEM <= EXTMEM() (unstable)), (FE <= FE() (unstable)), (FE2 <=
--- a/esp-metadata-generated/src/_generated_esp32s3.rs
+++ b/esp-metadata-generated/src/_generated_esp32s3.rs
@ -475,7 +475,8 @@ macro_rules! for_each_peripheral {
        _for_each_inner!((GPIO42 <= virtual())); _for_each_inner!((GPIO43 <= virtual()));
        _for_each_inner!((GPIO44 <= virtual())); _for_each_inner!((GPIO45 <= virtual()));
        _for_each_inner!((GPIO46 <= virtual())); _for_each_inner!((GPIO47 <= virtual()));
-        _for_each_inner!((GPIO48 <= virtual())); _for_each_inner!((AES <= AES()
+        _for_each_inner!((GPIO48 <= virtual())); _for_each_inner!((AES <= AES(AES : {
+        bind_peri_interrupt, enable_peri_interrupt, disable_peri_interrupt })
        (unstable))); _for_each_inner!((APB_CTRL <= APB_CTRL() (unstable)));
        _for_each_inner!((APB_SARADC <= APB_SARADC() (unstable)));
        _for_each_inner!((ASSIST_DEBUG <= ASSIST_DEBUG() (unstable)));
@ -546,7 +547,8 @@ macro_rules! for_each_peripheral {
        (GPIO36 <= virtual()), (GPIO37 <= virtual()), (GPIO38 <= virtual()), (GPIO39 <=
        virtual()), (GPIO40 <= virtual()), (GPIO41 <= virtual()), (GPIO42 <= virtual()),
        (GPIO43 <= virtual()), (GPIO44 <= virtual()), (GPIO45 <= virtual()), (GPIO46 <=
-        virtual()), (GPIO47 <= virtual()), (GPIO48 <= virtual()), (AES <= AES()
+        virtual()), (GPIO47 <= virtual()), (GPIO48 <= virtual()), (AES <= AES(AES : {
+        bind_peri_interrupt, enable_peri_interrupt, disable_peri_interrupt })
        (unstable)), (APB_CTRL <= APB_CTRL() (unstable)), (APB_SARADC <= APB_SARADC()
        (unstable)), (ASSIST_DEBUG <= ASSIST_DEBUG() (unstable)), (DMA <= DMA()
        (unstable)), (DS <= DS() (unstable)), (EFUSE <= EFUSE() (unstable)), (EXTMEM <=
--- a/esp-metadata/devices/esp32c3.toml
+++ b/esp-metadata/devices/esp32c3.toml
@ -13,7 +13,7 @@ cores  = 1
 trm    = "https://www.espressif.com/sites/default/files/documentation/esp32-c3_technical_reference_manual_en.pdf"

 peripherals = [
-    { name = "AES" },
+    { name = "AES", interrupts = { peri = "AES" } },
    { name = "APB_CTRL" },
    { name = "APB_SARADC" },
    { name = "ASSIST_DEBUG" },
--- a/esp-metadata/devices/esp32c6.toml
+++ b/esp-metadata/devices/esp32c6.toml
@ -13,7 +13,7 @@ cores  = 1
 trm    = "https://www.espressif.com/sites/default/files/documentation/esp32-c6_technical_reference_manual_en.pdf"

 peripherals = [
-    { name = "AES" },
+    { name = "AES", interrupts = { peri = "AES" } },
    { name = "APB_SARADC" },
    { name = "ASSIST_DEBUG" },
    { name = "ATOMIC" },
--- a/esp-metadata/devices/esp32h2.toml
+++ b/esp-metadata/devices/esp32h2.toml
@ -13,7 +13,7 @@ cores  = 1
 trm    = "https://www.espressif.com/sites/default/files/documentation/esp32-h2_technical_reference_manual_en.pdf"

 peripherals = [
-    { name = "AES" },
+    { name = "AES", interrupts = { peri = "AES" } },
    { name = "APB_SARADC" },
    { name = "ASSIST_DEBUG" },
    { name = "DMA" },
--- a/esp-metadata/devices/esp32s2.toml
+++ b/esp-metadata/devices/esp32s2.toml
@ -13,7 +13,7 @@ cores  = 1
 trm    = "https://www.espressif.com/sites/default/files/documentation/esp32-s2_technical_reference_manual_en.pdf"

 peripherals = [
-    { name = "AES" },
+    { name = "AES", interrupts = { peri = "AES" } },
    { name = "APB_SARADC" },
    { name = "DEDICATED_GPIO" },
    { name = "DS" },
--- a/esp-metadata/devices/esp32s3.toml
+++ b/esp-metadata/devices/esp32s3.toml
@ -13,7 +13,7 @@ cores  = 2
 trm    = "https://www.espressif.com/sites/default/files/documentation/esp32-s3_technical_reference_manual_en.pdf"

 peripherals = [
-    { name = "AES" },
+    { name = "AES", interrupts = { peri = "AES" } },
    { name = "APB_CTRL" },
    { name = "APB_SARADC" },
    { name = "ASSIST_DEBUG" },
--- a/hil-test/tests/aes.rs
+++ b/hil-test/tests/aes.rs
@ -1,20 +1,29 @@
 //! AES Tests

-//% CHIPS: esp32 esp32c3 esp32c6 esp32h2 esp32s2 esp32s3
-//% FEATURES: unstable
+//% CHIPS(quad): esp32s2
+// The S3 dev kit in the HIL-tester has octal PSRAM.
+//% CHIPS(octal): esp32s3
+// ESP32 has no AES-DMA, no point in setting up PSRAM
+//% CHIPS(no_psram): esp32 esp32c3 esp32c6 esp32h2
+
+//% ENV(octal): ESP_HAL_CONFIG_PSRAM_MODE=octal
+//% FEATURES(quad, octal): psram
+//% FEATURES: unstable esp-alloc/nightly

 #![no_std]
 #![no_main]

 use embassy_executor::Spawner;
 use embassy_sync::{blocking_mutex::raw::CriticalSectionRawMutex, signal::Signal};
+#[cfg(aes_dma)]
+use esp_hal::aes::dma::AesDmaBackend;
 use esp_hal::{
    Config,
    aes::{
        Aes,
        AesBackend,
        AesContext,
-        CipherModeState,
+        CipherState,
        Key,
        Operation,
        cipher_modes::{Cbc, Cfb8, Cfb128, Ctr, Ecb, Ofb},
@ -225,6 +234,9 @@ const CIPHERTEXT_CFB128: [u8; PLAINTEXT_BUF_SIZE] = [

 esp_bootloader_esp_idf::esp_app_desc!();

+#[cfg(aes_dma)]
+extern crate alloc;
+
 const fn pad_to<const K: usize>(input: &[u8]) -> [u8; K] {
    let mut out = [0; K];

@ -254,7 +266,7 @@ where

 fn aes_roundtrip<const K: usize>(
    tag: &'static str,
-    block_ctx: impl Into<CipherModeState>,
+    block_ctx: impl Into<CipherState>,
    plaintext: &[u8],
    ciphertext: &[u8],
    buffer: &mut [u8],
@ -284,6 +296,16 @@ fn run_cipher_tests(buffer: &mut [u8]) {
    let mut plaintext = [0; PLAINTEXT_BUF_SIZE];
    fill_with_plaintext(&mut plaintext);

+    // Let's use ECB as a test case for short unaligned DMA transfers.
+    let short_len = 16;
+    aes_roundtrip::<16>(
+        "Short ECB",
+        Ecb,
+        &plaintext[0..short_len],
+        &CIPHERTEXT_ECB_128[0..short_len],
+        &mut buffer[0..short_len],
+    );
+
    aes_roundtrip::<16>("ECB", Ecb, &plaintext, &CIPHERTEXT_ECB_128, buffer);
    #[cfg(esp32s2)]
    aes_roundtrip::<24>("ECB", Ecb, &plaintext, &CIPHERTEXT_ECB_192, buffer);
@ -307,7 +329,7 @@ fn run_cipher_tests(buffer: &mut [u8]) {
 }

 #[cfg(test)]
-#[embedded_test::tests(default_timeout = 3, executor = hil_test::Executor::new())]
+#[embedded_test::tests(default_timeout = 10, executor = hil_test::Executor::new())] // defmt slows the tests down a bit
 mod tests {
    use super::*;

@ -340,7 +362,7 @@ mod tests {
    #[cfg(not(esp32))]
    fn test_aes_dma_ecb() {
        use esp_hal::{
-            aes::dma::{AesDma, CipherMode},
+            aes::dma::AesDma,
            dma::{DmaRxBuf, DmaTxBuf},
            dma_buffers,
        };
@ -353,6 +375,8 @@ mod tests {
        where
            Key: From<[u8; K]>,
        {
+            use esp_hal::aes::dma::DmaCipherState;
+
            const DMA_BUFFER_SIZE: usize = 16;

            let (output, rx_descriptors, input, tx_descriptors) = dma_buffers!(DMA_BUFFER_SIZE);
@ -367,7 +391,7 @@ mod tests {
                    output,
                    input,
                    Operation::Encrypt,
-                    CipherMode::Ecb,
+                    &DmaCipherState::from(Ecb),
                    pad_to::<K>(KEY),
                )
                .map_err(|e| e.0)
@ -383,7 +407,7 @@ mod tests {
                    output,
                    input,
                    Operation::Decrypt,
-                    CipherMode::Ecb,
+                    &DmaCipherState::from(Ecb),
                    pad_to::<K>(KEY),
                )
                .map_err(|e| e.0)
@ -434,6 +458,66 @@ mod tests {
        run_cipher_tests(&mut buffer);
    }

+    #[test]
+    #[cfg(aes_dma)]
+    fn test_aes_dma_work_queue() {
+        use allocator_api2::vec::Vec;
+
+        let p = esp_hal::init(Config::default().with_cpu_clock(CpuClock::max()));
+
+        esp_alloc::heap_allocator!(size: 32 * 1024);
+
+        cfg_if::cfg_if! {
+            if #[cfg(esp32s2)] {
+                let mut aes = AesDmaBackend::new(p.AES, p.DMA_CRYPTO);
+            } else {
+                let mut aes = AesDmaBackend::new(p.AES, p.DMA_CH0);
+            }
+        }
+        let _backend = aes.start();
+
+        let mut internal_memory =
+            Vec::with_capacity_in(PLAINTEXT_BUF_SIZE + 15, esp_alloc::InternalMemory);
+        internal_memory.resize(PLAINTEXT_BUF_SIZE + 15, 0);
+
+        // Different alignments in internal memory
+        for shift in 0..15 {
+            let buffer = &mut internal_memory[shift..][..PLAINTEXT_BUF_SIZE];
+            run_cipher_tests(buffer);
+        }
+    }
+
+    #[test]
+    #[cfg(all(aes_dma, psram))]
+    fn test_aes_dma_work_queue_psram() {
+        use allocator_api2::vec::Vec;
+
+        let p = esp_hal::init(Config::default().with_cpu_clock(CpuClock::max()));
+        esp_alloc::psram_allocator!(p.PSRAM, esp_hal::psram);
+
+        cfg_if::cfg_if! {
+            if #[cfg(esp32s2)] {
+                let mut aes = AesDmaBackend::new(p.AES, p.DMA_CRYPTO);
+            } else {
+                let mut aes = AesDmaBackend::new(p.AES, p.DMA_CH0);
+            }
+        }
+        let _backend = aes.start();
+
+        let mut plaintext = [0; PLAINTEXT_BUF_SIZE];
+        fill_with_plaintext(&mut plaintext);
+
+        // Different alignments in external memory
+        let mut external_memory =
+            Vec::with_capacity_in(PLAINTEXT_BUF_SIZE + 15, esp_alloc::ExternalMemory);
+        external_memory.resize(PLAINTEXT_BUF_SIZE + 15, 0);
+
+        for shift in 0..15 {
+            let buffer = &mut external_memory[shift..][..PLAINTEXT_BUF_SIZE];
+            run_cipher_tests(buffer);
+        }
+    }
+
    #[test]
    fn test_aes_work_queue_work_posted_before_queue_started() {
        let p = esp_hal::init(Config::default().with_cpu_clock(CpuClock::max()));