From 9fd62d547c86f1bfa8bbee7d3f0b8219492ead43 Mon Sep 17 00:00:00 2001
From: wanlei <wanlei@espressif.com>
Date: Mon, 12 Dec 2022 15:09:22 +0800
Subject: [PATCH] spi_flash: fix flash function arg_check with spi host_id

---
 components/spi_flash/spi_flash_os_func_app.c | 58 +++++++++++---------
 components/spi_flash/test/test_esp_flash.c   | 18 ++++--
 2 files changed, 46 insertions(+), 30 deletions(-)

diff --git a/components/spi_flash/spi_flash_os_func_app.c b/components/spi_flash/spi_flash_os_func_app.c
index f9fc95fdc2..5dd958f9d9 100644
--- a/components/spi_flash/spi_flash_os_func_app.c
+++ b/components/spi_flash/spi_flash_os_func_app.c
@@ -243,33 +243,39 @@ esp_err_t esp_flash_init_os_functions(esp_flash_t *chip, int host_id, spi_bus_lo
         return ESP_ERR_INVALID_ARG;
     }
 
-    if (host_id == SPI1_HOST) {
-        //SPI1
-        chip->os_func = &esp_flash_spi1_default_os_functions;
-        chip->os_func_data = heap_caps_malloc(sizeof(spi1_app_func_arg_t),
-                                         MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT);
-        if (chip->os_func_data == NULL) {
-            return ESP_ERR_NO_MEM;
-        }
-        *(spi1_app_func_arg_t*) chip->os_func_data = (spi1_app_func_arg_t) {
-            .common_arg = {
+    switch (host_id)
+    {
+        case SPI1_HOST:
+            //SPI1
+            chip->os_func = &esp_flash_spi1_default_os_functions;
+            chip->os_func_data = heap_caps_malloc(sizeof(spi1_app_func_arg_t),
+                                            MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT);
+            if (chip->os_func_data == NULL) {
+                return ESP_ERR_NO_MEM;
+            }
+            *(spi1_app_func_arg_t*) chip->os_func_data = (spi1_app_func_arg_t) {
+                .common_arg = {
+                    .dev_lock = dev_handle,
+                },
+                .no_protect = true,
+            };
+            break;
+        case SPI2_HOST:
+#if SOC_SPI_PERIPH_NUM > 2
+        case SPI3_HOST:
+#endif
+            //SPI2, SPI3
+            chip->os_func = &esp_flash_spi23_default_os_functions;
+            chip->os_func_data = heap_caps_malloc(sizeof(app_func_arg_t),
+                                            MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT);
+            if (chip->os_func_data == NULL) {
+                return ESP_ERR_NO_MEM;
+            }
+            *(app_func_arg_t*) chip->os_func_data = (app_func_arg_t) {
                 .dev_lock = dev_handle,
-            },
-            .no_protect = true,
-        };
-    } else if (host_id == SPI2_HOST || host_id == SPI3_HOST) {
-        //SPI2, SPI3
-        chip->os_func = &esp_flash_spi23_default_os_functions;
-        chip->os_func_data = heap_caps_malloc(sizeof(app_func_arg_t),
-                                         MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT);
-        if (chip->os_func_data == NULL) {
-            return ESP_ERR_NO_MEM;
-        }
-        *(app_func_arg_t*) chip->os_func_data = (app_func_arg_t) {
-            .dev_lock = dev_handle,
-        };
-    } else {
-        return ESP_ERR_INVALID_ARG;
+            };
+            break;
+        default: return ESP_ERR_INVALID_ARG;
     }
 
     return ESP_OK;
diff --git a/components/spi_flash/test/test_esp_flash.c b/components/spi_flash/test/test_esp_flash.c
index d274608b1b..8cf7ee740e 100644
--- a/components/spi_flash/test/test_esp_flash.c
+++ b/components/spi_flash/test/test_esp_flash.c
@@ -349,7 +349,8 @@ static void setup_bus(spi_host_device_t host_id)
         gpio_iomux_out(hd_pin, spi_periph_signal[host_id].func, false);
 #endif //CONFIG_ESPTOOLPY_FLASHMODE_QIO || CONFIG_ESPTOOLPY_FLASHMODE_QOUT
         //currently the SPI bus for main flash chip is initialized through GPIO matrix
-    } else if (host_id == SPI2_HOST) {
+    }
+    else if (host_id == SPI2_HOST) {
         ESP_LOGI(TAG, "setup flash on SPI%d (FSPI) CS0...\n", host_id + 1);
         spi_bus_config_t fspi_bus_cfg = {
             .mosi_io_num = FSPI_PIN_NUM_MOSI,
@@ -361,7 +362,9 @@ static void setup_bus(spi_host_device_t host_id)
         };
         esp_err_t ret = spi_bus_initialize(host_id, &fspi_bus_cfg, 0);
         TEST_ESP_OK(ret);
-    } else if (host_id == SPI3_HOST) {
+    }
+#if SOC_SPI_PERIPH_NUM > 2
+    else if (host_id == SPI3_HOST) {
         ESP_LOGI(TAG, "setup flash on SPI%d (HSPI) CS0...\n", host_id + 1);
         spi_bus_config_t hspi_bus_cfg = {
             .mosi_io_num = HSPI_PIN_NUM_MOSI,
@@ -380,7 +383,9 @@ static void setup_bus(spi_host_device_t host_id)
 
         gpio_set_direction(HSPI_PIN_NUM_WP, GPIO_MODE_OUTPUT);
         gpio_set_level(HSPI_PIN_NUM_WP, 1);
-    } else {
+    }
+#endif
+    else {
         ESP_LOGE(TAG, "invalid bus");
     }
 }
@@ -389,7 +394,12 @@ static void setup_bus(spi_host_device_t host_id)
 static void release_bus(int host_id)
 {
     //SPI1 bus can't be deinitialized
-    if (host_id == SPI2_HOST || host_id == SPI3_HOST) {
+#if SOC_SPI_PERIPH_NUM > 2
+    if (host_id == SPI2_HOST || host_id == SPI3_HOST)
+#else
+    if (host_id == SPI2_HOST)
+#endif
+    {
         spi_bus_free(host_id);
     }
 }