From 65436262d5d2becb09f094ef8b54e926a5d9f5b4 Mon Sep 17 00:00:00 2001
From: Laukik Hase <laukik.hase@espressif.com>
Date: Fri, 26 Sep 2025 19:24:04 +0530
Subject: [PATCH] feat(esp_tee): Add support for TEE secure storage encryption
 for ESP32-C61

---
 .../tee_sec_storage/tee_sec_storage.c         | 61 ++++++++++++++++---
 .../mbedtls/esp_tee/esp_tee_mbedtls_config.h  |  4 ++
 2 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/components/esp_tee/subproject/components/tee_sec_storage/tee_sec_storage.c b/components/esp_tee/subproject/components/tee_sec_storage/tee_sec_storage.c
index d42c20c0cf0..0ad68fd8699 100644
--- a/components/esp_tee/subproject/components/tee_sec_storage/tee_sec_storage.c
+++ b/components/esp_tee/subproject/components/tee_sec_storage/tee_sec_storage.c
@@ -16,6 +16,8 @@
 #if SOC_HMAC_SUPPORTED
 #include "esp_hmac.h"
 #include "esp_hmac_pbkdf2.h"
+#else
+#include "mbedtls/md.h"
 #endif
 
 #include "mbedtls/aes.h"
@@ -139,25 +141,60 @@ static int rand_func(void *rng_state, unsigned char *output, size_t len)
 }
 
 #if CONFIG_SECURE_TEE_SEC_STG_MODE_RELEASE
-static esp_err_t compute_nvs_keys_with_hmac(hmac_key_id_t hmac_key_id, nvs_sec_cfg_t *cfg)
+static esp_err_t compute_nvs_keys_with_hmac(esp_efuse_block_t key_blk, nvs_sec_cfg_t *cfg)
 {
-    uint32_t ekey_seed[8] = {[0 ... 7] = EKEY_SEED};
-    uint32_t tkey_seed[8] = {[0 ... 7] = TKEY_SEED};
+    const uint32_t ekey_seed[8] = {[0 ... 7] = EKEY_SEED};
+    const uint32_t tkey_seed[8] = {[0 ... 7] = TKEY_SEED};
+    esp_err_t err = ESP_FAIL;
 
     memset(cfg, 0x00, sizeof(nvs_sec_cfg_t));
 
-    esp_err_t err = ESP_FAIL;
+#if SOC_HMAC_SUPPORTED
+    hmac_key_id_t hmac_key_id = (hmac_key_id_t)(key_blk - EFUSE_BLK_KEY0);
+    if (hmac_key_id < 0 || hmac_key_id >= HMAC_KEY_MAX) {
+        return ESP_ERR_INVALID_ARG;
+    }
+
     err = esp_hmac_calculate(hmac_key_id, ekey_seed, sizeof(ekey_seed), (uint8_t *)cfg->eky);
     err |= esp_hmac_calculate(hmac_key_id, tkey_seed, sizeof(tkey_seed), (uint8_t *)cfg->tky);
     if (err != ESP_OK) {
+        memset(cfg, 0x00, sizeof(nvs_sec_cfg_t));
+        ESP_LOGE(TAG, "Failed to calculate seed HMAC");
+        return err;
+    }
+    ESP_FAULT_ASSERT(err == ESP_OK);
+#else
+    const mbedtls_md_info_t *md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
+    if (md_info == NULL) {
+        return ESP_FAIL;
+    }
+
+    uint8_t key_buf[SHA256_DIGEST_SZ] = {0};
+    /* NOTE: The eFuse key for SoCs without HMAC support should NOT be
+     * read-protected when burning in the eFuse
+     */
+    err = esp_efuse_read_block(key_blk, key_buf, 0, sizeof(key_buf) * 8);
+    if (err != ESP_OK) {
+        return err;
+    }
+
+    int ret = mbedtls_md_hmac(md_info, key_buf, sizeof(key_buf),
+                              (const uint8_t *)ekey_seed, sizeof(ekey_seed),
+                              cfg->eky);
+    ret |= mbedtls_md_hmac(md_info, key_buf, sizeof(key_buf),
+                           (const uint8_t *)tkey_seed, sizeof(tkey_seed),
+                           cfg->tky);
+    if (ret != 0) {
+        memset(cfg, 0x00, sizeof(nvs_sec_cfg_t));
         ESP_LOGE(TAG, "Failed to calculate seed HMAC");
         return ESP_FAIL;
     }
-    ESP_FAULT_ASSERT(err == ESP_OK);
+    ESP_FAULT_ASSERT(ret == 0);
+    memset(key_buf, 0x00, sizeof(key_buf));
+#endif
 
     /* NOTE: If the XTS E-key and T-key are the same, we have a hash collision */
     ESP_FAULT_ASSERT(memcmp(cfg->eky, cfg->tky, NVS_KEY_SIZE) != 0);
-
     return ESP_OK;
 }
 #endif
@@ -170,13 +207,17 @@ static esp_err_t read_security_cfg_hmac(nvs_sec_cfg_t *cfg)
 
 #if CONFIG_SECURE_TEE_SEC_STG_MODE_RELEASE
     esp_efuse_block_t hmac_key_blk = (esp_efuse_block_t)(EFUSE_BLK_KEY0 + (esp_efuse_block_t)CONFIG_SECURE_TEE_SEC_STG_EFUSE_HMAC_KEY_ID);
-    esp_efuse_purpose_t hmac_efuse_blk_purpose = esp_efuse_get_key_purpose(hmac_key_blk);
-    if (hmac_efuse_blk_purpose != ESP_EFUSE_KEY_PURPOSE_HMAC_UP) {
-        ESP_LOGE(TAG, "HMAC key is not burnt in eFuse block");
+    esp_efuse_purpose_t purpose = esp_efuse_get_key_purpose(hmac_key_blk);
+#if SOC_HMAC_SUPPORTED
+    if (purpose != ESP_EFUSE_KEY_PURPOSE_HMAC_UP) {
+#else
+    if (purpose != ESP_EFUSE_KEY_PURPOSE_USER) {
+#endif
+        ESP_LOGE(TAG, "Key is not burnt in eFuse block with expected purpose");
         return ESP_ERR_NOT_FOUND;
     }
 
-    esp_err_t err = compute_nvs_keys_with_hmac((hmac_key_id_t)CONFIG_SECURE_TEE_SEC_STG_EFUSE_HMAC_KEY_ID, cfg);
+    esp_err_t err = compute_nvs_keys_with_hmac(hmac_key_blk, cfg);
     if (err != ESP_OK) {
         return err;
     }
diff --git a/components/mbedtls/esp_tee/esp_tee_mbedtls_config.h b/components/mbedtls/esp_tee/esp_tee_mbedtls_config.h
index b30a8493e12..d8ea8ff8460 100644
--- a/components/mbedtls/esp_tee/esp_tee_mbedtls_config.h
+++ b/components/mbedtls/esp_tee/esp_tee_mbedtls_config.h
@@ -80,6 +80,10 @@
 #define MBEDTLS_ECP_VERIFY_ALT
 #endif
 
+#if !SOC_HMAC_SUPPORTED
+#define MBEDTLS_MD_C
+#endif
+
 #define MBEDTLS_ENTROPY_C
 
 #endif /* ESP_TEE_MBEDTLS_CONFIG_H */