From 98c53234fc7c6b75d77728ccae752584a0419996 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        | 56 +++++++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/components/mbedtls/port/aes/dma/esp_aes.c b/components/mbedtls/port/aes/dma/esp_aes.c
index b5e4befd38..6bf2bee26d 100644
--- a/components/mbedtls/port/aes/dma/esp_aes.c
+++ b/components/mbedtls/port/aes/dma/esp_aes.c
@@ -362,8 +362,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 5f63444c02..a5ccf69f2f 100644
--- a/components/mbedtls/test/test_aes.c
+++ b/components/mbedtls/test/test_aes.c
@@ -1,3 +1,9 @@
+/*
+ * SPDX-FileCopyrightText: 2022-2023 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Unlicense OR CC0-1.0
+ */
+
 /* mbedTLS AES test
 */
 #include <string.h>
@@ -15,6 +21,7 @@
 #include "freertos/FreeRTOS.h"
 #include "freertos/task.h"
 #include "freertos/semphr.h"
+#include "soc/lldesc.h"
 
 static const uint8_t key_256[] = {
     0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
@@ -108,6 +115,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;