From 98e16412a770da68dfb340f3b179fab705293c0e Mon Sep 17 00:00:00 2001
From: Laukik Hase <laukik.hase@espressif.com>
Date: Tue, 4 Mar 2025 16:17:35 +0530
Subject: [PATCH] refactor(esp_tee): Use the AES-GCM port layer for operations
 in the TEE

---
 .../mbedtls/esp_tee/esp_tee_mbedtls.cmake     | 27 +++++++++----------
 .../mbedtls/esp_tee/esp_tee_mbedtls_config.h  |  1 +
 .../mbedtls/port/aes/dma/esp_aes_dma_core.c   | 12 +++++++--
 3 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/components/mbedtls/esp_tee/esp_tee_mbedtls.cmake b/components/mbedtls/esp_tee/esp_tee_mbedtls.cmake
index b569a73a2b..9072a91c6c 100644
--- a/components/mbedtls/esp_tee/esp_tee_mbedtls.cmake
+++ b/components/mbedtls/esp_tee/esp_tee_mbedtls.cmake
@@ -13,20 +13,22 @@ list(APPEND include_dirs "${heap_dir}/include")
 set(srcs "${COMPONENT_DIR}/esp_tee/esp_tee_crypto_shared_gdma.c")
 
 # AES-SHA implementation
-list(APPEND include_dirs "${COMPONENT_DIR}/port/aes/include"
-                         "${COMPONENT_DIR}/port/aes/dma/include"
-                         "${COMPONENT_DIR}/port/sha/core/include")
+list(APPEND srcs "${COMPONENT_DIR}/port/aes/dma/esp_aes.c"
+                 "${COMPONENT_DIR}/port/aes/dma/esp_aes_dma_core.c")
 
 list(APPEND srcs "${COMPONENT_DIR}/port/aes/esp_aes_common.c"
-                 "${COMPONENT_DIR}/port/aes/dma/esp_aes.c"
-                 "${COMPONENT_DIR}/port/aes/dma/esp_aes_dma_core.c")
+                 "${COMPONENT_DIR}/port/aes/esp_aes_gcm.c")
 
 list(APPEND srcs "${COMPONENT_DIR}/port/sha/core/sha.c"
                  "${COMPONENT_DIR}/port/sha/esp_sha.c")
 
-idf_component_register(INCLUDE_DIRS "${include_dirs}"
-                       PRIV_REQUIRES "${priv_requires}"
-                       SRCS "${srcs}")
+list(APPEND include_dirs "${COMPONENT_DIR}/port/aes/include"
+                         "${COMPONENT_DIR}/port/aes/dma/include"
+                         "${COMPONENT_DIR}/port/sha/core/include")
+
+idf_component_register(SRCS "${srcs}"
+                       INCLUDE_DIRS "${include_dirs}"
+                       PRIV_REQUIRES "${priv_requires}")
 
 # Only build mbedtls libraries
 set(ENABLE_TESTING CACHE BOOL OFF)
@@ -50,9 +52,6 @@ endforeach()
 
 target_link_libraries(${COMPONENT_LIB} INTERFACE ${mbedtls_targets})
 
-if(CONFIG_MBEDTLS_HARDWARE_SHA)
-    target_sources(mbedcrypto PRIVATE  "${COMPONENT_DIR}/port/sha/core/esp_sha1.c"
-                                       "${COMPONENT_DIR}/port/sha/core/esp_sha256.c"
-                                       "${COMPONENT_DIR}/port/sha/core/esp_sha512.c"
-    )
-endif()
+target_sources(mbedcrypto PRIVATE "${COMPONENT_DIR}/port/sha/core/esp_sha1.c"
+                                  "${COMPONENT_DIR}/port/sha/core/esp_sha256.c"
+                                  "${COMPONENT_DIR}/port/sha/core/esp_sha512.c")
diff --git a/components/mbedtls/esp_tee/esp_tee_mbedtls_config.h b/components/mbedtls/esp_tee/esp_tee_mbedtls_config.h
index 665a009e1c..a676cd20a9 100644
--- a/components/mbedtls/esp_tee/esp_tee_mbedtls_config.h
+++ b/components/mbedtls/esp_tee/esp_tee_mbedtls_config.h
@@ -34,6 +34,7 @@
 #define MBEDTLS_CIPHER_C
 #define MBEDTLS_AES_C
 #define MBEDTLS_GCM_C
+#define MBEDTLS_GCM_ALT
 
 #define MBEDTLS_ASN1_WRITE_C
 #define MBEDTLS_ASN1_PARSE_C
diff --git a/components/mbedtls/port/aes/dma/esp_aes_dma_core.c b/components/mbedtls/port/aes/dma/esp_aes_dma_core.c
index 34152306d8..a4a3d7c43b 100644
--- a/components/mbedtls/port/aes/dma/esp_aes_dma_core.c
+++ b/components/mbedtls/port/aes/dma/esp_aes_dma_core.c
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2024 Espressif Systems (Shanghai) CO LTD
+ * SPDX-FileCopyrightText: 2024-2025 Espressif Systems (Shanghai) CO LTD
  *
  * SPDX-License-Identifier: Apache-2.0
  */
@@ -100,9 +100,11 @@ static IRAM_ATTR void esp_aes_complete_isr(void *arg)
         portYIELD_FROM_ISR();
     }
 }
+#endif
 
 void esp_aes_intr_alloc(void)
 {
+#if !ESP_TEE_BUILD
     if (op_complete_sem == NULL) {
         const int isr_flags = esp_intr_level_to_flags(CONFIG_MBEDTLS_AES_INTERRUPT_LEVEL);
 
@@ -120,8 +122,14 @@ void esp_aes_intr_alloc(void)
         // Static semaphore creation is unlikely to fail but still basic sanity
         assert(op_complete_sem != NULL);
     }
-}
+#else
+    // NOTE: Need to extern since the mbedtls component does not depend on
+    // the esp_tee (main) component
+    extern void esp_tee_aes_intr_alloc(void);
+    esp_tee_aes_intr_alloc();
 #endif
+}
+
 
 static esp_err_t esp_aes_isr_initialise( void )
 {