From 771a02da9339d3c1d43211ecac33aaa0896ba9eb 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/flash_ops.c              | 35 ++++---------------
 tools/ci/config/target-test.yml               |  2 +-
 .../configs/flash_encryption_psram            |  2 +-
 3 files changed, 8 insertions(+), 31 deletions(-)

diff --git a/components/spi_flash/flash_ops.c b/components/spi_flash/flash_ops.c
index da7da957e1..158c4feb51 100644
--- a/components/spi_flash/flash_ops.c
+++ b/components/spi_flash/flash_ops.c
@@ -450,16 +450,14 @@ out:
 }
 #endif // CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 
-#if !CONFIG_SPI_FLASH_USE_LEGACY_IMPL
-extern void spi_common_set_dummy_output(esp_rom_spiflash_read_mode_t mode);
+#ifndef CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 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;
+
 #if CONFIG_ESPTOOLPY_FLASHFREQ_80M
     dummy_bit = ESP_ROM_SPIFLASH_DUMMY_LEN_PLUS_80M;
 #elif CONFIG_ESPTOOLPY_FLASHFREQ_40M
@@ -469,7 +467,6 @@ static void IRAM_ATTR flash_rom_init(void)
 #elif CONFIG_ESPTOOLPY_FLASHFREQ_20M
     dummy_bit = ESP_ROM_SPIFLASH_DUMMY_LEN_PLUS_20M;
 #endif
-#endif//CONFIG_IDF_TARGET_ESP32
 
 #if CONFIG_ESPTOOLPY_FLASHFREQ_80M
     freqdiv = 1;
@@ -481,37 +478,17 @@ static void IRAM_ATTR flash_rom_init(void)
     freqdiv = 4;
 #endif
 
-#if CONFIG_ESPTOOLPY_FLASHMODE_QIO
-    read_mode = ESP_ROM_SPIFLASH_QIO_MODE;
-#elif CONFIG_ESPTOOLPY_FLASHMODE_QOUT
-    read_mode = ESP_ROM_SPIFLASH_QOUT_MODE;
-#elif CONFIG_ESPTOOLPY_FLASHMODE_DIO
-    read_mode = ESP_ROM_SPIFLASH_DIO_MODE;
-#elif CONFIG_ESPTOOLPY_FLASHMODE_DOUT
-    read_mode = ESP_ROM_SPIFLASH_DOUT_MODE;
-#endif
-
-#if CONFIG_IDF_TARGET_ESP32
     g_rom_spiflash_dummy_len_plus[1] = dummy_bit;
-#else
-    spi_dummy_len_fix(1, freqdiv);
-#endif //CONFIG_IDF_TARGET_ESP32
-
-#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;
 }
 #endif // !CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 esp_err_t IRAM_ATTR spi_flash_write_encrypted(size_t dest_addr, const void *src, size_t size)
 {
-    flash_rom_init();
     CHECK_WRITE_ADDRESS(dest_addr, size);
     const uint8_t *ssrc = (const uint8_t *)src;
     if ((dest_addr % 16) != 0) {
@@ -558,6 +535,7 @@ esp_err_t IRAM_ATTR spi_flash_write_encrypted(size_t dest_addr, const void *src,
             }
 
             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) {
@@ -732,7 +710,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 44db632a15..86450ac21b 100644
--- a/tools/ci/config/target-test.yml
+++ b/tools/ci/config/target-test.yml
@@ -331,7 +331,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/flash_encryption_psram b/tools/unit-test-app/configs/flash_encryption_psram
index 13648e1d12..12ce59a17e 100644
--- a/tools/unit-test-app/configs/flash_encryption_psram
+++ b/tools/unit-test-app/configs/flash_encryption_psram
@@ -1,4 +1,4 @@
-# This config is for ESP32 only (no ESP32-S2/S3 flash encryption support yet, ESP32-C3 has no psram)
+# This config is for ESP32 only
 CONFIG_IDF_TARGET="esp32"
 TEST_COMPONENTS=spi_flash
 TEST_GROUPS=flash_encryption