From 69bcbe97f4120dd6e50f01d4f4e5ae28db36ce60 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_apps/main/test_aes.c | 52 +++++++++++++++++++-
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/components/mbedtls/port/aes/dma/esp_aes.c b/components/mbedtls/port/aes/dma/esp_aes.c
index efc531a55a6..ec4f3ae739b 100644
--- a/components/mbedtls/port/aes/dma/esp_aes.c
+++ b/components/mbedtls/port/aes/dma/esp_aes.c
@@ -371,8 +371,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_apps/main/test_aes.c b/components/mbedtls/test_apps/main/test_aes.c
index 925e4a82f26..5ac71e82ec9 100644
--- a/components/mbedtls/test_apps/main/test_aes.c
+++ b/components/mbedtls/test_apps/main/test_aes.c
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2022 Espressif Systems (Shanghai) CO LTD
+ * SPDX-FileCopyrightText: 2022-2023 Espressif Systems (Shanghai) CO LTD
  *
  * SPDX-License-Identifier: Unlicense OR CC0-1.0
  */
@@ -21,6 +21,7 @@
 #include "freertos/task.h"
 #include "freertos/semphr.h"
 #include "esp_memory_utils.h"
+#include "soc/lldesc.h"
 
 
 static const uint8_t key_256[] = {
@@ -115,6 +116,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;