From e609b6fb10f45dbdd3d9a54c3d00e2066dc85979 Mon Sep 17 00:00:00 2001
From: Ashish Sharma <ashish.sharma@espressif.com>
Date: Tue, 8 Jul 2025 15:02:07 +0800
Subject: [PATCH] fix(esp_tls): init DS peripheral RSA ALT context on heap

---
 components/esp-tls/esp_tls_mbedtls.c | 28 +++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/components/esp-tls/esp_tls_mbedtls.c b/components/esp-tls/esp_tls_mbedtls.c
index cfd2bf5672..2eb8b0f16d 100644
--- a/components/esp-tls/esp_tls_mbedtls.c
+++ b/components/esp-tls/esp_tls_mbedtls.c
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2019-2024 Espressif Systems (Shanghai) CO LTD
+ * SPDX-FileCopyrightText: 2019-2025 Espressif Systems (Shanghai) CO LTD
  *
  * SPDX-License-Identifier: Apache-2.0
  */
@@ -39,6 +39,7 @@ static esp_err_t esp_set_atecc608a_pki_context(esp_tls_t *tls, const void *pki);
 #endif /* CONFIG_ESP_TLS_USE_SECURE_ELEMENT */
 
 #if defined(CONFIG_ESP_TLS_USE_DS_PERIPHERAL)
+#include <pk_wrap.h>
 #include "rsa_sign_alt.h"
 static esp_err_t esp_mbedtls_init_pk_ctx_for_ds(const void *pki);
 #endif /* CONFIG_ESP_TLS_USE_DS_PERIPHERAL */
@@ -359,6 +360,18 @@ void esp_mbedtls_cleanup(esp_tls_t *tls)
 #endif
     mbedtls_x509_crt_free(&tls->cacert);
     mbedtls_x509_crt_free(&tls->clientcert);
+
+#ifdef CONFIG_ESP_TLS_USE_DS_PERIPHERAL
+    if (mbedtls_pk_get_type(&tls->clientkey) == MBEDTLS_PK_RSA_ALT) {
+        mbedtls_rsa_alt_context *rsa_alt = tls->clientkey.MBEDTLS_PRIVATE(pk_ctx);
+        if (rsa_alt && rsa_alt->key != NULL) {
+            mbedtls_rsa_free(rsa_alt->key);
+            mbedtls_free(rsa_alt->key);
+            rsa_alt->key = NULL;
+        }
+    }
+#endif
+
     mbedtls_pk_free(&tls->clientkey);
     mbedtls_entropy_free(&tls->entropy);
     mbedtls_ssl_config_free(&tls->conf);
@@ -1097,12 +1110,18 @@ static esp_err_t esp_mbedtls_init_pk_ctx_for_ds(const void *pki)
 {
     int ret = -1;
     /* initialize the mbedtls pk context with rsa context */
-    mbedtls_rsa_context rsakey;
-    mbedtls_rsa_init(&rsakey);
-    if ((ret = mbedtls_pk_setup_rsa_alt(((const esp_tls_pki_t*)pki)->pk_key, &rsakey, NULL, esp_ds_rsa_sign,
+    mbedtls_rsa_context *rsakey = calloc(1, sizeof(mbedtls_rsa_context));
+    if (rsakey == NULL) {
+        ESP_LOGE(TAG, "Failed to allocate memory for mbedtls_rsa_context");
+        return ESP_ERR_NO_MEM;
+    }
+    mbedtls_rsa_init(rsakey);
+    if ((ret = mbedtls_pk_setup_rsa_alt(((const esp_tls_pki_t*)pki)->pk_key, rsakey, NULL, esp_ds_rsa_sign,
                                         esp_ds_get_keylen )) != 0) {
         ESP_LOGE(TAG, "Error in mbedtls_pk_setup_rsa_alt, returned -0x%04X", -ret);
         mbedtls_print_error_msg(ret);
+        mbedtls_rsa_free(rsakey);
+        free(rsakey);
         ret = ESP_FAIL;
         goto exit;
     }
@@ -1113,7 +1132,6 @@ static esp_err_t esp_mbedtls_init_pk_ctx_for_ds(const void *pki)
     }
     ESP_LOGD(TAG, "DS peripheral params initialized.");
 exit:
-    mbedtls_rsa_free(&rsakey);
     return ret;
 }
 #endif /* CONFIG_ESP_TLS_USE_DS_PERIPHERAL */