From 2e08c5d7978915de40bd913fe3bf3794dfa5dec4 Mon Sep 17 00:00:00 2001
From: John Doe <ficeto@Hristos-Mac-mini-2.local>
Date: Wed, 1 Jul 2015 15:38:30 +0300
Subject: [PATCH] Add an option to reboot on failed update

---
 cores/esp8266/Esp.cpp | 13 ++++++++++---
 cores/esp8266/Esp.h   |  2 +-
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/cores/esp8266/Esp.cpp b/cores/esp8266/Esp.cpp
index cda089b04..bf374814e 100644
--- a/cores/esp8266/Esp.cpp
+++ b/cores/esp8266/Esp.cpp
@@ -358,10 +358,12 @@ uint32_t EspClass::getFreeSketchSpace() {
     return freeSpaceEnd - freeSpaceStart;
 }
 
-bool EspClass::updateSketch(Stream& in, uint32_t size) {
+bool EspClass::updateSketch(Stream& in, uint32_t size, bool restartOnFail) {
 
-    if (size > getFreeSketchSpace())
+    if (size > getFreeSketchSpace()){
+        if(restartOnFail) ESP.restart();
         return false;
+    }
 
     uint32_t usedSize = getSketchSize();
     uint32_t freeSpaceStart = (usedSize + FLASH_SECTOR_SIZE - 1) & (~(FLASH_SECTOR_SIZE - 1));
@@ -374,8 +376,10 @@ bool EspClass::updateSketch(Stream& in, uint32_t size) {
     noInterrupts();
     int rc = SPIEraseAreaEx(freeSpaceStart, roundedSize);
     interrupts();
-    if (rc)
+    if (rc){
+        if(restartOnFail) ESP.restart();
         return false;
+    }
 
 #ifdef DEBUG_SERIAL
     DEBUG_SERIAL.println("erase done");
@@ -397,12 +401,14 @@ bool EspClass::updateSketch(Stream& in, uint32_t size) {
 #ifdef DEBUG_SERIAL
             DEBUG_SERIAL.println("stream read failed");
 #endif
+            if(restartOnFail) ESP.restart();
             return false;
         }
 
         if(addr == freeSpaceStart) {
             // check for valid first magic byte
             if(*((uint8 *) buffer.get()) != 0xE9) {
+                if(restartOnFail) ESP.restart();
                 return false;
             }
         }
@@ -414,6 +420,7 @@ bool EspClass::updateSketch(Stream& in, uint32_t size) {
 #ifdef DEBUG_SERIAL
             DEBUG_SERIAL.println("write failed");
 #endif            
+            if(restartOnFail) ESP.restart();
             return false;
         }
 
diff --git a/cores/esp8266/Esp.h b/cores/esp8266/Esp.h
index 80555d591..8021ed198 100644
--- a/cores/esp8266/Esp.h
+++ b/cores/esp8266/Esp.h
@@ -116,7 +116,7 @@ class EspClass {
 
         uint32_t getSketchSize();
         uint32_t getFreeSketchSpace();
-        bool updateSketch(Stream& in, uint32_t size);
+        bool updateSketch(Stream& in, uint32_t size, bool restartOnFail = false);
 
         String getResetInfo();
         struct rst_info * getResetInfoPtr();