From fa7f7bde81884542bc315dc1ba8042d24eb41e40 Mon Sep 17 00:00:00 2001
From: Markus Sattler <help.markus+git@gmail.com>
Date: Sun, 20 Dec 2015 12:45:36 +0100
Subject: [PATCH] add _verifyHeader and _verifyEnd

_verifyHeader is called before the beginning of the update progress to verify the first byte using peek
_verifyEnd is called on the end before the eboot command is written to verify first byte + flash config

add missing _reset() on timeout
---
 cores/esp8266/Updater.cpp | 71 ++++++++++++++++++++++++++++++++++++---
 cores/esp8266/Updater.h   | 23 ++++++++-----
 2 files changed, 80 insertions(+), 14 deletions(-)

diff --git a/cores/esp8266/Updater.cpp b/cores/esp8266/Updater.cpp
index c890a13af..131d29ab2 100644
--- a/cores/esp8266/Updater.cpp
+++ b/cores/esp8266/Updater.cpp
@@ -172,6 +172,14 @@ bool UpdaterClass::end(bool evenIfRemaining){
 #endif
   }
 
+  if(!_verifyEnd()) {
+#ifdef DEBUG_UPDATER
+    printError(DEBUG_UPDATER);
+#endif
+    _reset();
+    return false;
+  }
+
   if (_command == U_FLASH) {
     eboot_command ebcmd;
     ebcmd.action = ACTION_COPY_RAW;
@@ -246,19 +254,67 @@ size_t UpdaterClass::write(uint8_t *data, size_t len) {
   return len;
 }
 
+bool UpdaterClass::_verifyHeader(uint8_t data) {
+    if(_command == U_FLASH) {
+        // check for valid first magic byte (is always 0xE9)
+        if(data != 0xE9) {
+            _error = UPDATE_ERROR_MAGIC_BYTE;
+            _currentAddress = (_startAddress + _size);
+            return false;
+        }
+        return true;
+    } else if(_command == U_SPIFFS) {
+        // no check of SPIFFS possible with first byte.
+        return true;
+    }
+    return false;
+}
+
+bool UpdaterClass::_verifyEnd() {
+    if(_command == U_FLASH) {
+
+        uint8_t buf[4];
+        if(!ESP.flashRead(_startAddress, (uint32_t *) &buf[0], 4)) {
+            _error = UPDATE_ERROR_READ;
+            _currentAddress = (_startAddress);
+            return false;
+        }
+
+        // check for valid first magic byte
+        if(buf[0] != 0xE9) {
+            _error = UPDATE_ERROR_MAGIC_BYTE;
+            _currentAddress = (_startAddress);
+            return false;
+        }
+
+        uint32_t bin_flash_size = ESP.magicFlashChipSize((buf[3] & 0xf0) >> 4);
+
+        // check if new bin fits to SPI flash
+        if(bin_flash_size > ESP.getFlashChipRealSize()) {
+            _error = UPDATE_ERROR_NEW_FLASH_CONFIG;
+            _currentAddress = (_startAddress);
+            return false;
+        }
+
+        return true;
+    } else if(_command == U_SPIFFS) {
+        // SPIFFS is already over written checks make no sense any more.
+        return true;
+    }
+    return false;
+}
+
 size_t UpdaterClass::writeStream(Stream &data) {
     size_t written = 0;
     size_t toRead = 0;
     if(hasError() || !isRunning())
         return 0;
 
-    // check for valid first magic byte (is always 0xE9)
-    if(data.peek() != 0xE9) {
-        _error = UPDATE_ERROR_MAGIC_BYTE;
-        _currentAddress = (_startAddress + _size);
+    if(!_verifyHeader(data.peek())) {
 #ifdef DEBUG_UPDATER
         printError(DEBUG_UPDATER);
 #endif
+        _reset();
         return 0;
     }
 
@@ -273,8 +329,9 @@ size_t UpdaterClass::writeStream(Stream &data) {
 #ifdef DEBUG_UPDATER
                 printError(DEBUG_UPDATER);
 #endif
+                _reset();
+                return written;
             }
-            return written;
         }
         _bufferLen += toRead;
         if((_bufferLen == remaining() || _bufferLen == FLASH_SECTOR_SIZE) && !_writeBuffer())
@@ -293,6 +350,8 @@ void UpdaterClass::printError(Stream &out){
     out.println("Flash Write Failed");
   } else if(_error == UPDATE_ERROR_ERASE){
     out.println("Flash Erase Failed");
+  } else if(_error == UPDATE_ERROR_READ){
+    out.println("Flash Read Failed");
   } else if(_error == UPDATE_ERROR_SPACE){
     out.println("Not Enough Space");
   } else if(_error == UPDATE_ERROR_SIZE){
@@ -303,6 +362,8 @@ void UpdaterClass::printError(Stream &out){
     out.println("MD5 Check Failed");
   } else if(_error == UPDATE_ERROR_FLASH_CONFIG){
     out.printf("Flash config wrong real: %d IDE: %d\n", ESP.getFlashChipRealSize(), ESP.getFlashChipSize());
+  } else if(_error == UPDATE_ERROR_NEW_FLASH_CONFIG){
+    out.printf("new Flash config wrong real: %d\n", ESP.getFlashChipRealSize());
   } else if(_error == UPDATE_ERROR_MAGIC_BYTE){
     out.println("Magic byte is wrong, not 0xE9");
   } else {
diff --git a/cores/esp8266/Updater.h b/cores/esp8266/Updater.h
index 94ea3b241..48fe8560f 100644
--- a/cores/esp8266/Updater.h
+++ b/cores/esp8266/Updater.h
@@ -5,15 +5,17 @@
 #include "flash_utils.h"
 #include "MD5Builder.h"
 
-#define UPDATE_ERROR_OK             (0)
-#define UPDATE_ERROR_WRITE          (1)
-#define UPDATE_ERROR_ERASE          (2)
-#define UPDATE_ERROR_SPACE          (3)
-#define UPDATE_ERROR_SIZE           (4)
-#define UPDATE_ERROR_STREAM         (5)
-#define UPDATE_ERROR_MD5            (6)
-#define UPDATE_ERROR_FLASH_CONFIG   (7)
-#define UPDATE_ERROR_MAGIC_BYTE     (8)
+#define UPDATE_ERROR_OK                 (0)
+#define UPDATE_ERROR_WRITE              (1)
+#define UPDATE_ERROR_ERASE              (2)
+#define UPDATE_ERROR_READ               (3)
+#define UPDATE_ERROR_SPACE              (4)
+#define UPDATE_ERROR_SIZE               (5)
+#define UPDATE_ERROR_STREAM             (6)
+#define UPDATE_ERROR_MD5                (7)
+#define UPDATE_ERROR_FLASH_CONFIG       (8)
+#define UPDATE_ERROR_NEW_FLASH_CONFIG   (9)
+#define UPDATE_ERROR_MAGIC_BYTE         (10)
 
 
 #define U_FLASH   0
@@ -134,6 +136,9 @@ class UpdaterClass {
     void _reset();
     bool _writeBuffer();
 
+    bool _verifyHeader(uint8_t data);
+    bool _verifyEnd();
+
     uint8_t _error;
     uint8_t *_buffer;
     size_t _bufferLen;