From 94e139e639c849bddc8540965486f5488cedc7d5 Mon Sep 17 00:00:00 2001
From: Mahavir Jain <mahavir@espressif.com>
Date: Mon, 29 May 2023 14:39:17 +0530
Subject: [PATCH] aes: fix DMA descriptor calculation for the alignment case

The number of the DMA descriptors allocated for certain length (e.g.,
8176) were not sufficient (off by 1 error). This used to result in the
dynamic memory corruption as the region was modified beyond the
allocated range.

This change fixes the DMA descriptor calculation part and allocates
sufficient DMA descriptors based on the data length alignment considerations.

Test has also been added to cover the specific scenario in the CI.

Closes https://github.com/espressif/esp-idf/issues/11310
---
 components/mbedtls/port/aes/dma/esp_aes.c |  4 +-
 components/mbedtls/test/test_aes.c        | 50 +++++++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/components/mbedtls/port/aes/dma/esp_aes.c b/components/mbedtls/port/aes/dma/esp_aes.c
index 8347a4a08e..0ceded7b0c 100644
--- a/components/mbedtls/port/aes/dma/esp_aes.c
+++ b/components/mbedtls/port/aes/dma/esp_aes.c
@@ -366,8 +366,8 @@ static int esp_aes_process_dma(esp_aes_context *ctx, const unsigned char *input,
             return esp_aes_process_dma_ext_ram(ctx, input, output, len, stream_out, input_needs_realloc, output_needs_realloc);
         }
 
-        /* Set up dma descriptors for input and output */
-        lldesc_num = lldesc_get_required_num(block_bytes);
+        /* Set up dma descriptors for input and output considering the 16 byte alignment requirement for EDMA */
+        lldesc_num = lldesc_get_required_num_constrained(block_bytes, LLDESC_MAX_NUM_PER_DESC_16B_ALIGNED);
 
         /* Allocate both in and out descriptors to save a malloc/free per function call */
         block_desc = heap_caps_calloc(lldesc_num * 2, sizeof(lldesc_t), MALLOC_CAP_DMA);
diff --git a/components/mbedtls/test/test_aes.c b/components/mbedtls/test/test_aes.c
index 4d5e4872d9..b2aa81e00f 100644
--- a/components/mbedtls/test/test_aes.c
+++ b/components/mbedtls/test/test_aes.c
@@ -16,6 +16,7 @@
 #include "freertos/task.h"
 #include "freertos/semphr.h"
 #include "esp_memory_utils.h"
+#include "soc/lldesc.h"
 
 
 static const uint8_t key_256[] = {
@@ -110,6 +111,55 @@ TEST_CASE("mbedtls CBC AES-256 test", "[aes]")
     free(decryptedtext);
 }
 
+TEST_CASE("mbedtls CBC AES-256 DMA buffer align test", "[aes]")
+{
+#define ALIGN_DOWN(val, align) ((val) & ~((align) - 1))
+    // Size is taken considering the maximum DMA buffer size
+    const unsigned SZ = ALIGN_DOWN((2*LLDESC_MAX_NUM_PER_DESC), 16);
+    mbedtls_aes_context ctx;
+    uint8_t nonce[16];
+
+    const uint8_t expected_cipher_end[] = {
+        0x9e, 0xcb, 0x1d, 0x24, 0x01, 0xc8, 0x3f, 0xba,
+        0xde, 0x76, 0xea, 0x9c, 0xf3, 0x64, 0x23, 0x19,
+        0x8c, 0x67, 0xd4, 0x1a, 0xd1, 0xe0, 0xbf, 0xc3,
+        0xd2, 0xb8, 0x40, 0x95, 0x89, 0x41, 0x09, 0xdb,
+    };
+
+    memcpy(nonce, iv, 16);
+
+    // allocate internal memory
+    uint8_t *chipertext = heap_caps_malloc(SZ, MALLOC_CAP_8BIT|MALLOC_CAP_INTERNAL);
+    uint8_t *plaintext = heap_caps_malloc(SZ, MALLOC_CAP_8BIT|MALLOC_CAP_INTERNAL);
+    uint8_t *decryptedtext = heap_caps_malloc(SZ, MALLOC_CAP_8BIT|MALLOC_CAP_INTERNAL);
+
+    TEST_ASSERT_NOT_NULL(chipertext);
+    TEST_ASSERT_NOT_NULL(plaintext);
+    TEST_ASSERT_NOT_NULL(decryptedtext);
+
+    mbedtls_aes_init(&ctx);
+    mbedtls_aes_setkey_enc(&ctx, key_256, 256);
+
+    memset(plaintext, 0x3A, SZ);
+    memset(decryptedtext, 0x0, SZ);
+
+    // Encrypt
+    mbedtls_aes_crypt_cbc(&ctx, MBEDTLS_AES_ENCRYPT, SZ, nonce, plaintext, chipertext);
+    TEST_ASSERT_EQUAL_HEX8_ARRAY(expected_cipher_end, chipertext + SZ - 32, 32);
+
+    // Decrypt
+    memcpy(nonce, iv, 16);
+    mbedtls_aes_setkey_dec(&ctx, key_256, 256);
+    mbedtls_aes_crypt_cbc(&ctx, MBEDTLS_AES_DECRYPT, SZ, nonce, chipertext, decryptedtext);
+
+    TEST_ASSERT_EQUAL_HEX8_ARRAY(plaintext, decryptedtext, SZ);
+
+    mbedtls_aes_free(&ctx);
+    free(plaintext);
+    free(chipertext);
+    free(decryptedtext);
+}
+
 TEST_CASE("mbedtls CTR AES-256 test", "[aes]")
 {
     const unsigned SZ = 1000;