From c0b796532bb86a112562b7b3be12b3c2d37dad95 Mon Sep 17 00:00:00 2001
From: Harshit Malpani <harshit.malpani@espressif.com>
Date: Mon, 10 Jan 2022 18:32:37 +0530
Subject: [PATCH] esp_https_ota: fix for checking chip id at start of OTA

---
 components/esp_https_ota/src/esp_https_ota.c  | 53 ++++++++++++++-----
 .../ota/advanced_https_ota/example_test.py    |  2 +-
 2 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/components/esp_https_ota/src/esp_https_ota.c b/components/esp_https_ota/src/esp_https_ota.c
index ede3ff9b2d..612b523d23 100644
--- a/components/esp_https_ota/src/esp_https_ota.c
+++ b/components/esp_https_ota/src/esp_https_ota.c
@@ -286,17 +286,8 @@ failure:
     return err;
 }
 
-esp_err_t esp_https_ota_get_img_desc(esp_https_ota_handle_t https_ota_handle, esp_app_desc_t *new_app_info)
+static esp_err_t read_header(esp_https_ota_t *handle)
 {
-    esp_https_ota_t *handle = (esp_https_ota_t *)https_ota_handle;
-    if (handle == NULL || new_app_info == NULL)  {
-        ESP_LOGE(TAG, "esp_https_ota_read_img_desc: Invalid argument");
-        return ESP_ERR_INVALID_ARG;
-    }
-    if (handle->state < ESP_HTTPS_OTA_BEGIN) {
-        ESP_LOGE(TAG, "esp_https_ota_read_img_desc: Invalid state");
-        return ESP_FAIL;
-    }
     /*
      * `data_read_size` holds number of bytes needed to read complete header.
      * `bytes_read` holds number of bytes read.
@@ -327,10 +318,37 @@ esp_err_t esp_https_ota_get_img_desc(esp_https_ota_handle_t https_ota_handle, es
         return ESP_FAIL;
     }
     handle->binary_file_len = bytes_read;
+    return ESP_OK;
+}
+
+esp_err_t esp_https_ota_get_img_desc(esp_https_ota_handle_t https_ota_handle, esp_app_desc_t *new_app_info)
+{
+    esp_https_ota_t *handle = (esp_https_ota_t *)https_ota_handle;
+    if (handle == NULL || new_app_info == NULL)  {
+        ESP_LOGE(TAG, "esp_https_ota_read_img_desc: Invalid argument");
+        return ESP_ERR_INVALID_ARG;
+    }
+    if (handle->state < ESP_HTTPS_OTA_BEGIN) {
+        ESP_LOGE(TAG, "esp_https_ota_read_img_desc: Invalid state");
+        return ESP_FAIL;
+    }
+    if (read_header(handle) != ESP_OK) {
+        return ESP_FAIL;
+    }
     memcpy(new_app_info, &handle->ota_upgrade_buf[sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t)], sizeof(esp_app_desc_t));
     return ESP_OK;
 }
 
+static esp_err_t esp_ota_verify_chip_id(void *arg)
+{
+    esp_image_header_t *data = (esp_image_header_t*)(arg);
+    if (data->chip_id != CONFIG_IDF_FIRMWARE_CHIP_ID) {
+        ESP_LOGE(TAG, "Mismatch chip id, expected %d, found %d", CONFIG_IDF_FIRMWARE_CHIP_ID, data->chip_id);
+        return ESP_ERR_INVALID_VERSION;
+    }
+    return ESP_OK;
+}
+
 esp_err_t esp_https_ota_perform(esp_https_ota_handle_t https_ota_handle)
 {
     esp_https_ota_t *handle = (esp_https_ota_t *)https_ota_handle;
@@ -357,16 +375,25 @@ esp_err_t esp_https_ota_perform(esp_https_ota_handle_t https_ota_handle)
             /* In case `esp_https_ota_read_img_desc` was invoked first,
                then the image data read there should be written to OTA partition
                */
+            int binary_file_len = 0;
             if (handle->binary_file_len) {
                 /*
                  * Header length gets added to handle->binary_file_len in _ota_write
                  * Clear handle->binary_file_len to avoid additional 289 bytes in binary_file_len
                  */
-                int binary_file_len = handle->binary_file_len;
+                binary_file_len = handle->binary_file_len;
                 handle->binary_file_len = 0;
-                return _ota_write(handle, (const void *)handle->ota_upgrade_buf, binary_file_len);
+            } else {
+                if (read_header(handle) != ESP_OK) {
+                    return ESP_FAIL;
+                }
+                binary_file_len = IMAGE_HEADER_SIZE;
             }
-            /* falls through */
+            err = esp_ota_verify_chip_id(handle->ota_upgrade_buf);
+            if (err != ESP_OK) {
+                return err;
+            }
+            return _ota_write(handle, (const void *)handle->ota_upgrade_buf, binary_file_len);
         case ESP_HTTPS_OTA_IN_PROGRESS:
             data_read = esp_http_client_read(handle->http_client,
                                              handle->ota_upgrade_buf,
diff --git a/examples/system/ota/advanced_https_ota/example_test.py b/examples/system/ota/advanced_https_ota/example_test.py
index 6c3fa358d4..589706f2c0 100644
--- a/examples/system/ota/advanced_https_ota/example_test.py
+++ b/examples/system/ota/advanced_https_ota/example_test.py
@@ -301,7 +301,7 @@ def test_examples_protocol_advanced_https_ota_example_random(env, extra_data):
 
     print('writing to device: {}'.format('https://' + host_ip + ':' + str(server_port) + '/' + random_bin_name))
     dut1.write('https://' + host_ip + ':' + str(server_port) + '/' + random_bin_name)
-    dut1.expect('esp_ota_ops: OTA image has invalid magic byte', timeout=10)
+    dut1.expect(re.compile(r'esp_https_ota: Mismatch chip id, expected 0, found \d'), timeout=10)
     os.remove(binary_file)
     thread1.terminate()