From 2e9cb80033069ec2dc04f278eff79c56ff05611c Mon Sep 17 00:00:00 2001
From: Dong Heng <dongheng@espressif.com>
Date: Thu, 24 Jan 2019 15:32:24 +0800
Subject: [PATCH] feat(ota): Add OTA binary link address verify of ESP8285 or
 ESP8266 + 1MB flash

---
 components/app_update/esp_ota_ops.c         | 29 +++++++++++++++++++++
 components/esp8266/Makefile.projbuild       | 10 +++----
 examples/system/ota/main/ota_example_main.c |  2 +-
 3 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/components/app_update/esp_ota_ops.c b/components/app_update/esp_ota_ops.c
index 9f41cda7..b9072340 100644
--- a/components/app_update/esp_ota_ops.c
+++ b/components/app_update/esp_ota_ops.c
@@ -69,6 +69,28 @@ static ota_select s_ota_select[2];
 
 const static char *TAG = "esp_ota_ops";
 
+#ifndef CONFIG_ESP8266_BOOT_COPY_APP
+static inline int esp_ota_verify_binary(const esp_partition_pos_t *pos, esp_image_header_t *image)
+{
+    const int32_t entry = image->entry_addr - 0x40200010;
+
+    ESP_LOGD(TAG, "OTA binary start entry 0x%x, partition start from 0x%x to 0x%x\n", entry, pos->offset,
+                        pos->offset + pos->size);
+
+    if (pos->offset + pos->size <= 0x100000) {
+        if (entry <= 0 || entry <= pos->offset || entry >= pos->offset + pos->size) {
+            const char *doc_str = "<<ESP8266_RTOS_SDK/examples/system/ota/README.md>>";
+
+            ESP_LOGE(TAG, "**Important**: The OTA binary link data is error, "
+                          "please refer to document %s for how to generate OTA binaries", doc_str);
+            return ESP_ERR_INVALID_ARG;
+        }
+    }
+
+    return ESP_OK;
+}
+#endif
+
 /* Return true if this is an OTA app partition */
 static bool is_ota_partition(const esp_partition_t *p)
 {
@@ -244,6 +266,13 @@ esp_err_t esp_ota_end(esp_ota_handle_t handle)
         goto cleanup;
     }
 
+#ifndef CONFIG_ESP8266_BOOT_COPY_APP
+    if (esp_ota_verify_binary(&part_pos, &data.image) != ESP_OK) {
+        ret = ESP_ERR_OTA_VALIDATE_FAILED;
+        goto cleanup;
+    }
+#endif
+
 #ifdef CONFIG_SECURE_BOOT_ENABLED
     ret = esp_secure_boot_verify_signature(it->part->address, data.image_len);
     if (ret != ESP_OK) {
diff --git a/components/esp8266/Makefile.projbuild b/components/esp8266/Makefile.projbuild
index 735d0ae9..d90d587b 100644
--- a/components/esp8266/Makefile.projbuild
+++ b/components/esp8266/Makefile.projbuild
@@ -94,7 +94,7 @@ OTA_V2_TO_V3_BIN  := ./build/$(PROJECT_NAME).v2_to_v3.ota.bin
 CONFIG_APP2_OFFSET ?= $(CONFIG_APP1_OFFSET)
 CONFIG_APP2_SIZE ?= $(CONFIG_APP1_SIZE)
 
-OTA1_OFFSET := CONFIG_APP1_OFFSET
+OTA1_OFFSET := $(CONFIG_APP1_OFFSET)
 ifdef CONFIG_ESP8266_BOOT_COPY_APP
 OTA2_LINK_OFFSET := $(CONFIG_APP1_OFFSET)
 else
@@ -102,7 +102,7 @@ OTA2_LINK_OFFSET := $(CONFIG_APP2_OFFSET)
 endif
 
 $(OTA2_BIN): all_binaries
-ifeq ($(CONFIG_ESPTOOLPY_FLASHSIZE), "1MB")
+ifneq ($(OTA1_OFFSET), $(OTA2_LINK_OFFSET))
 	@rm -f ./build/esp8266/esp8266_out.ld
 	@make APP_OFFSET=$(OTA2_LINK_OFFSET) APP_SIZE=$(CONFIG_APP2_SIZE) CFLAGS= CXXFLAGS=
 endif
@@ -110,16 +110,16 @@ endif
 	@echo [GEN] $(OTA2_BIN)
 
 $(OTA1_BIN): all_binaries
-ifeq ($(CONFIG_ESPTOOLPY_FLASHSIZE), "1MB")
+ifneq ($(OTA1_OFFSET), $(OTA2_LINK_OFFSET))
 	@rm -f ./build/esp8266/esp8266_out.ld
 endif
-	@make APP_OFFSET=$(CONFIG_APP1_OFFSET) APP_SIZE=$(CONFIG_APP1_SIZE) CFLAGS= CXXFLAGS=
+	@make APP_OFFSET=$(OTA1_OFFSET) APP_SIZE=$(CONFIG_APP1_SIZE) CFLAGS= CXXFLAGS=
 	@cp $(RAW_BIN) $(OTA1_BIN)
 	@echo [GEN] $(OTA1_BIN)
 
 $(OTA_BIN): $(OTA1_BIN) $(OTA2_BIN)
 	@cp $(OTA1_BIN) $(OTA_BIN)
-ifeq ($(CONFIG_ESPTOOLPY_FLASHSIZE), "1MB")
+ifneq ($(OTA1_OFFSET), $(OTA2_LINK_OFFSET))
 	@cat $(OTA2_BIN) >> $(OTA_BIN)
 endif
 	@cp $(OTA1_BIN) $(RAW_BIN)
diff --git a/examples/system/ota/main/ota_example_main.c b/examples/system/ota/main/ota_example_main.c
index b5137287..27e58646 100644
--- a/examples/system/ota/main/ota_example_main.c
+++ b/examples/system/ota/main/ota_example_main.c
@@ -185,7 +185,7 @@ bool _esp_ota_firm_parse_http(esp_ota_firm_t *ota_firm, const char *text, size_t
             memset(length_str, 0, sizeof(length_str));
             memcpy(length_str, ptr, ptr2 - ptr);
             ota_firm->content_len = atoi(length_str);
-#ifdef CONFIG_ESPTOOLPY_FLASHSIZE_1MB
+#if defined(CONFIG_ESPTOOLPY_FLASHSIZE_1MB) && !defined(CONFIG_ESP8266_BOOT_COPY_APP)
             ota_firm->ota_size = ota_firm->content_len / ota_firm->ota_num;
             ota_firm->ota_offset = ota_firm->ota_size * ota_firm->update_ota_num;
 #else