From 714a59ab64c4a713c38ced157d78dacd88242880 Mon Sep 17 00:00:00 2001
From: Paul Sokolovsky <pfalcon@users.sourceforge.net>
Date: Tue, 3 Jan 2017 00:02:00 +0300
Subject: [PATCH] esp8266/modesp: Make check_fw() work with OTA firmware.

---
 esp8266/modesp.c | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/esp8266/modesp.c b/esp8266/modesp.c
index 1cd231f79f..e162981ffb 100644
--- a/esp8266/modesp.c
+++ b/esp8266/modesp.c
@@ -628,10 +628,12 @@ STATIC mp_obj_t esp_flash_size(void) {
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_0(esp_flash_size_obj, esp_flash_size);
 
+// If there's just 1 loadable segment at the start of flash,
+// we assume there's a yaota8266 bootloader.
+#define IS_OTA_FIRMWARE() ((*(uint32_t*)0x40200000 & 0xff00) == 0x100)
+
 STATIC mp_obj_t esp_flash_user_start(void) {
-    if ((*(uint32_t*)0x40200000 & 0xff00) == 0x100) {
-        // If there's just 1 loadable segment at the start of flash,
-        // we assume there's a yaota8266 bootloader.
+    if (IS_OTA_FIRMWARE()) {
         return MP_OBJ_NEW_SMALL_INT(0x3c000 + 0x90000);
     } else {
         return MP_OBJ_NEW_SMALL_INT(0x90000);
@@ -641,10 +643,20 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_0(esp_flash_user_start_obj, esp_flash_user_start)
 
 STATIC mp_obj_t esp_check_fw(void) {
     MD5_CTX ctx;
-    uint32_t *sz_p = (uint32_t*)0x40208ffc;
-    printf("size: %d\n", *sz_p);
+    char *fw_start = (char*)0x40200000;
+    if (IS_OTA_FIRMWARE()) {
+        // Skip yaota8266 bootloader
+        fw_start += 0x3c000;
+    }
+
+    uint32_t size = *(uint32_t*)(fw_start + 0x8ffc);
+    printf("size: %d\n", size);
+    if (size > 1024 * 1024) {
+        print("Invalid size\n");
+        return mp_const_false;
+    }
     MD5Init(&ctx);
-    MD5Update(&ctx, (char*)0x40200004, *sz_p - 4);
+    MD5Update(&ctx, fw_start + 4, size - 4);
     unsigned char digest[16];
     MD5Final(digest, &ctx);
     printf("md5: ");
@@ -652,7 +664,7 @@ STATIC mp_obj_t esp_check_fw(void) {
         printf("%02x", digest[i]);
     }
     printf("\n");
-    return mp_obj_new_bool(memcmp(digest, (void*)(0x40200000 + *sz_p), sizeof(digest)) == 0);
+    return mp_obj_new_bool(memcmp(digest, fw_start + size, sizeof(digest)) == 0);
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_0(esp_check_fw_obj, esp_check_fw);