From 068dfcf1bac39f7a33f742445840b6d8cb301281 Mon Sep 17 00:00:00 2001
From: Cao Sen Miao <caosenmiao@espressif.com>
Date: Fri, 19 Feb 2021 11:50:27 +0800
Subject: [PATCH] ota: fix ota with flash encryption

---
 components/spi_flash/esp32/flash_ops_esp32.c     |  2 ++
 components/spi_flash/esp32s2/flash_ops_esp32s2.c |  2 ++
 components/spi_flash/flash_ops.c                 | 14 +++++++-------
 tools/ci/config/target-test.yml                  |  2 +-
 tools/unit-test-app/configs/psram_3              |  2 +-
 5 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/components/spi_flash/esp32/flash_ops_esp32.c b/components/spi_flash/esp32/flash_ops_esp32.c
index e23e1aa540..c7b7a07afe 100644
--- a/components/spi_flash/esp32/flash_ops_esp32.c
+++ b/components/spi_flash/esp32/flash_ops_esp32.c
@@ -32,6 +32,7 @@ static inline void IRAM_ATTR spi_flash_guard_end(void)
     }
 }
 
+extern void IRAM_ATTR flash_rom_init(void);
 esp_rom_spiflash_result_t IRAM_ATTR spi_flash_write_encrypted_chip(size_t dest_addr, const void *src, size_t size)
 {
     const uint8_t *ssrc = (const uint8_t *)src;
@@ -73,6 +74,7 @@ esp_rom_spiflash_result_t IRAM_ATTR spi_flash_write_encrypted_chip(size_t dest_a
         }
 
         spi_flash_guard_start();
+        flash_rom_init();
         rc = esp_rom_spiflash_write_encrypted(row_addr, (uint32_t *)encrypt_buf, 32);
         spi_flash_guard_end();
         if (rc != ESP_ROM_SPIFLASH_RESULT_OK) {
diff --git a/components/spi_flash/esp32s2/flash_ops_esp32s2.c b/components/spi_flash/esp32s2/flash_ops_esp32s2.c
index 6cbe9b4fa0..039845ba46 100644
--- a/components/spi_flash/esp32s2/flash_ops_esp32s2.c
+++ b/components/spi_flash/esp32s2/flash_ops_esp32s2.c
@@ -29,6 +29,7 @@ static const char *TAG = "spiflash_s2";
 #define SPICACHE SPIMEM0
 #define SPIFLASH SPIMEM1
 
+extern void IRAM_ATTR flash_rom_init(void);
 esp_rom_spiflash_result_t IRAM_ATTR spi_flash_write_encrypted_chip(size_t dest_addr, const void *src, size_t size)
 {
     const spi_flash_guard_funcs_t *ops = spi_flash_guard_get();
@@ -68,6 +69,7 @@ esp_rom_spiflash_result_t IRAM_ATTR spi_flash_write_encrypted_chip(size_t dest_a
         if (ops && ops->start) {
             ops->start();
         }
+        flash_rom_init();
         rc = SPI_Encrypt_Write(dest_addr, src, size);
         if (ops && ops->end) {
             ops->end();
diff --git a/components/spi_flash/flash_ops.c b/components/spi_flash/flash_ops.c
index c0a20104a3..f6fc2443bf 100644
--- a/components/spi_flash/flash_ops.c
+++ b/components/spi_flash/flash_ops.c
@@ -462,13 +462,13 @@ out:
 }
 #endif // CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 
-#if !CONFIG_SPI_FLASH_USE_LEGACY_IMPL
+#ifndef CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 extern void spi_common_set_dummy_output(esp_rom_spiflash_read_mode_t mode);
 extern void spi_dummy_len_fix(uint8_t spi, uint8_t freqdiv);
-static void IRAM_ATTR flash_rom_init(void)
+extern uint8_t g_rom_spiflash_dummy_len_plus[];
+void IRAM_ATTR flash_rom_init(void)
 {
     uint32_t freqdiv = 0;
-    esp_rom_spiflash_read_mode_t read_mode;
 
 #if CONFIG_IDF_TARGET_ESP32
     uint32_t dummy_bit = 0;
@@ -493,6 +493,8 @@ static void IRAM_ATTR flash_rom_init(void)
     freqdiv = 4;
 #endif
 
+#if !CONFIG_IDF_TARGET_ESP32S2 && !CONFIG_IDF_TARGET_ESP32
+    esp_rom_spiflash_read_mode_t read_mode;
 #if CONFIG_ESPTOOLPY_FLASHMODE_QIO
     read_mode = ESP_ROM_SPIFLASH_QIO_MODE;
 #elif CONFIG_ESPTOOLPY_FLASHMODE_QOUT
@@ -502,6 +504,7 @@ static void IRAM_ATTR flash_rom_init(void)
 #elif CONFIG_ESPTOOLPY_FLASHMODE_DOUT
     read_mode = ESP_ROM_SPIFLASH_DOUT_MODE;
 #endif
+#endif //!CONFIG_IDF_TARGET_ESP32S2 && !CONFIG_IDF_TARGET_ESP32
 
 #if CONFIG_IDF_TARGET_ESP32
     g_rom_spiflash_dummy_len_plus[1] = dummy_bit;
@@ -512,11 +515,10 @@ static void IRAM_ATTR flash_rom_init(void)
 #if !CONFIG_IDF_TARGET_ESP32S2 && !CONFIG_IDF_TARGET_ESP32
     spi_common_set_dummy_output(read_mode);
 #endif //!CONFIG_IDF_TARGET_ESP32S2
-    esp_rom_spiflash_config_readmode(read_mode);
     esp_rom_spiflash_config_clk(freqdiv, 1);
 }
 #else
-static void IRAM_ATTR flash_rom_init(void)
+void IRAM_ATTR flash_rom_init(void)
 {
     return;
 }
@@ -525,7 +527,6 @@ static void IRAM_ATTR flash_rom_init(void)
 esp_err_t IRAM_ATTR spi_flash_write_encrypted(size_t dest_addr, const void *src, size_t size)
 {
     esp_err_t err = ESP_OK;
-    flash_rom_init();
     CHECK_WRITE_ADDRESS(dest_addr, size);
     if ((dest_addr % 16) != 0) {
         return ESP_ERR_INVALID_ARG;
@@ -769,7 +770,6 @@ out:
 
 esp_err_t IRAM_ATTR spi_flash_read_encrypted(size_t src, void *dstv, size_t size)
 {
-    flash_rom_init();
     if (src + size > g_rom_flashchip.chip_size) {
         return ESP_ERR_INVALID_SIZE;
     }
diff --git a/tools/ci/config/target-test.yml b/tools/ci/config/target-test.yml
index cd90a2fee5..4c8096e77b 100644
--- a/tools/ci/config/target-test.yml
+++ b/tools/ci/config/target-test.yml
@@ -485,7 +485,7 @@ UT_020:
     - psram
 
 UT_021:
-  extends: .unit_test_32_template
+  extends: .unit_test_template
   parallel: 2
   tags:
     - ESP32_IDF
diff --git a/tools/unit-test-app/configs/psram_3 b/tools/unit-test-app/configs/psram_3
index 0e3ffdc193..af310d1ce7 100644
--- a/tools/unit-test-app/configs/psram_3
+++ b/tools/unit-test-app/configs/psram_3
@@ -1,5 +1,5 @@
 CONFIG_IDF_TARGET="esp32"
-TEST_COMPONENTS=driver freertos sdmmc
+TEST_COMPONENTS=driver
 CONFIG_ESP32_SPIRAM_SUPPORT=y
 CONFIG_ESP_INT_WDT_TIMEOUT_MS=800
 CONFIG_SPIRAM_OCCUPY_NO_HOST=y