fix(bootloader): self encryption workflow in bootloader not working on C5

Added explicit wait for key manager state to be idle before configuring
the register for flash encryption key usage from efuse. This now ensures
that flash contents are encrypted using efuse programmed key.

Also refactored code a bit to move into target specific directory.

Author: mahavirj

components/bootloader_support/include/esp_flash_encrypt.h

@@ -180,9 +180,17 @@ void esp_flash_encryption_init_checks(void);
 /** @brief Set all secure eFuse features related to flash encryption
  *
  * @return
- *  - ESP_OK - Successfully
+ *  - ESP_OK - On success
  */
 esp_err_t esp_flash_encryption_enable_secure_features(void);
+
+/** @brief Enable the key manager for flash encryption
+ *
+ * @return
+ *  - ESP_OK - On success
+ */
+esp_err_t esp_flash_encryption_enable_key_mgr(void);
+
 #endif /* BOOTLOADER_BUILD && CONFIG_SECURE_FLASH_ENC_ENABLED */
 
 /** @brief Returns the verification status for all physical security features of flash encryption in release mode

components/bootloader_support/src/esp32c5/flash_encryption_secure_features.c

@@ -11,6 +11,9 @@
 #include "esp_efuse_table.h"
 #include "esp_log.h"
 #include "sdkconfig.h"
+#include "soc/keymng_reg.h"
+#include "soc/pcr_reg.h"
+#include "soc/pcr_struct.h"
 
 static __attribute__((unused)) const char *TAG = "flash_encrypt";
 
@@ -58,3 +61,31 @@ esp_err_t esp_flash_encryption_enable_secure_features(void)
 
     return ESP_OK;
 }
+
+// TODO: Update to use LL APIs once key manager support added in IDF-8621
+esp_err_t esp_flash_encryption_enable_key_mgr(void)
+{
+    // Set the force power down bit to 0 to enable key manager
+    PCR.km_pd_ctrl.km_mem_force_pd = 0;
+    // Reset the key manager
+    PCR.km_conf.km_clk_en = 1;
+    PCR.km_conf.km_rst_en = 1;
+    PCR.km_conf.km_rst_en = 0;
+
+    // Wait for key manager to be ready
+    while (!PCR.km_conf.km_ready) {
+    };
+
+    // Wait for key manager state machine to be idle
+    while (REG_READ(KEYMNG_STATE_REG) != 0) {
+    };
+
+    // Set the key manager to use efuse key
+    REG_SET_FIELD(KEYMNG_STATIC_REG, KEYMNG_USE_EFUSE_KEY, 2);
+
+    // Reset MSPI to re-load the flash encryption key
+    REG_SET_BIT(PCR_MSPI_CLK_CONF_REG, PCR_MSPI_AXI_RST_EN);
+    REG_CLR_BIT(PCR_MSPI_CLK_CONF_REG, PCR_MSPI_AXI_RST_EN);
+
+    return ESP_OK;
+}

components/bootloader_support/src/esp32p4/flash_encryption_secure_features.c

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2022-2023 Espressif Systems (Shanghai) CO LTD
+ * SPDX-FileCopyrightText: 2022-2024 Espressif Systems (Shanghai) CO LTD
  *
  * SPDX-License-Identifier: Apache-2.0
  */
@@ -11,6 +11,8 @@
 #include "esp_efuse_table.h"
 #include "esp_log.h"
 #include "sdkconfig.h"
+#include "hal/key_mgr_ll.h"
+#include "hal/mspi_timing_tuning_ll.h"
 
 static __attribute__((unused)) const char *TAG = "flash_encrypt";
 
@@ -48,3 +50,22 @@ esp_err_t esp_flash_encryption_enable_secure_features(void)
 
     return ESP_OK;
 }
+
+esp_err_t esp_flash_encryption_enable_key_mgr(void)
+{
+    // Enable and reset key manager
+    // To suppress build errors about spinlock's __DECLARE_RCC_ATOMIC_ENV
+    int __DECLARE_RCC_ATOMIC_ENV __attribute__ ((unused));
+    key_mgr_ll_enable_bus_clock(true);
+    key_mgr_ll_enable_peripheral_clock(true);
+    key_mgr_ll_reset_register();
+
+    while (key_mgr_ll_get_state() != ESP_KEY_MGR_STATE_IDLE) {
+    };
+
+    // Force Key Manager to use eFuse key for XTS-AES operation
+    key_mgr_ll_set_key_usage(ESP_KEY_MGR_XTS_AES_128_KEY, ESP_KEY_MGR_USE_EFUSE_KEY);
+    _mspi_timing_ll_reset_mspi();
+
+    return ESP_OK;
+}

components/bootloader_support/src/flash_encryption/flash_encrypt.c

@@ -15,17 +15,7 @@
 #include "esp_efuse_table.h"
 #include "esp_log.h"
 #include "hal/wdt_hal.h"
-
-// Need to remove check and merge accordingly for ESP32C5 once key manager support added in IDF-8621
-#if SOC_KEY_MANAGER_FE_KEY_DEPLOY || CONFIG_IDF_TARGET_ESP32C5
-#if CONFIG_IDF_TARGET_ESP32C5
-#include "soc/keymng_reg.h"
-#include "soc/pcr_reg.h"
-#else /* CONFIG_IDF_TARGET_ESP32C5 */
-#include "hal/key_mgr_ll.h"
-#include "hal/mspi_timing_tuning_ll.h"
-#endif /* !CONFIG_IDF_TARGET_ESP32C5 */
-#endif /* SOC_KEY_MANAGER_FE_KEY_DEPLOY */
+#include "sdkconfig.h"
 
 #ifdef CONFIG_SOC_EFUSE_CONSISTS_OF_ONE_KEY_BLOCK
 #include "soc/sensitive_reg.h"
@@ -221,26 +211,6 @@ static esp_err_t check_and_generate_encryption_keys(void)
         }
         ESP_LOGI(TAG, "Using pre-loaded flash encryption key in efuse");
     }
-// Need to remove check for ESP32C5 and merge accordingly once key manager support added in IDF-8621
-#if SOC_KEY_MANAGER_FE_KEY_DEPLOY || CONFIG_IDF_TARGET_ESP32C5
-#if CONFIG_IDF_TARGET_ESP32C5
-    REG_SET_FIELD(KEYMNG_STATIC_REG, KEYMNG_USE_EFUSE_KEY, 2);
-    REG_SET_BIT(PCR_MSPI_CLK_CONF_REG, PCR_MSPI_AXI_RST_EN);
-    REG_CLR_BIT(PCR_MSPI_CLK_CONF_REG, PCR_MSPI_AXI_RST_EN);
-#else /* CONFIG_IDF_TARGET_ESP32C5 */
-    // Enable and reset key manager
-    // To suppress build errors about spinlock's __DECLARE_RCC_ATOMIC_ENV
-    int __DECLARE_RCC_ATOMIC_ENV __attribute__ ((unused));
-    key_mgr_ll_enable_bus_clock(true);
-    key_mgr_ll_enable_peripheral_clock(true);
-    key_mgr_ll_reset_register();
-    while (key_mgr_ll_get_state() != ESP_KEY_MGR_STATE_IDLE) {
-    };
-    // Force Key Manager to use eFuse key for XTS-AES operation
-    key_mgr_ll_set_key_usage(ESP_KEY_MGR_XTS_AES_128_KEY, ESP_KEY_MGR_USE_EFUSE_KEY);
-    _mspi_timing_ll_reset_mspi();
-#endif /* !CONFIG_IDF_TARGET_ESP32C5 */
-#endif /* SOC_KEY_MANAGER_FE_KEY_DEPLOY */
 
     return ESP_OK;
 }
@@ -288,6 +258,11 @@ esp_err_t esp_flash_encrypt_contents(void)
     REG_WRITE(SENSITIVE_XTS_AES_KEY_UPDATE_REG, 1);
 #endif
 
+// TODO: Remove C5 target config after key manager LL support- see IDF-8621
+#if CONFIG_SOC_KEY_MANAGER_FE_KEY_DEPLOY || CONFIG_IDF_TARGET_ESP32C5
+    esp_flash_encryption_enable_key_mgr();
+#endif
+
     err = encrypt_bootloader();
     if (err != ESP_OK) {
         return err;