From 063d608e18dc3df81dcb59b37c3e6844442bdc3b Mon Sep 17 00:00:00 2001
From: me-no-dev <hristo@espressif.com>
Date: Fri, 2 Jun 2017 18:24:26 +0300
Subject: [PATCH] Optimize Update class

- add roll back API to switch the running partition
- do not write the partition magic until the end to prevent booting
into partially written update
---
 libraries/Update/src/Update.h    |  9 +++++
 libraries/Update/src/Updater.cpp | 68 ++++++++++++++++++++++++++++----
 2 files changed, 70 insertions(+), 7 deletions(-)

diff --git a/libraries/Update/src/Update.h b/libraries/Update/src/Update.h
index 6f733d41..6cdcff24 100644
--- a/libraries/Update/src/Update.h
+++ b/libraries/Update/src/Update.h
@@ -137,6 +137,15 @@ class UpdateClass {
       return written;
     }
 
+    /*
+      check if there is a firmware on the other OTA partition that you can bootinto
+    */
+    bool canRollBack();
+    /*
+      set the other OTA partition as bootable (reboot to enable)
+    */
+    bool rollBack();
+
   private:
     void _reset();
     void _abort(uint8_t err);
diff --git a/libraries/Update/src/Updater.cpp b/libraries/Update/src/Updater.cpp
index 11676d27..7f9b4c57 100644
--- a/libraries/Update/src/Updater.cpp
+++ b/libraries/Update/src/Updater.cpp
@@ -35,6 +35,34 @@ static const char * _err2str(uint8_t _error){
     return ("UNKNOWN");
 }
 
+static bool _partitionIsBootable(const esp_partition_t* partition){
+    uint8_t buf[4];
+    if(!partition){
+        return false;
+    }
+    if(!ESP.flashRead(partition->address, (uint32_t*)buf, 4)) {
+        return false;
+    }
+
+    if(buf[0] != ESP_IMAGE_HEADER_MAGIC) {
+        return false;
+    }
+    return true;
+}
+
+static bool _enablePartition(const esp_partition_t* partition){
+    uint8_t buf[4];
+    if(!partition){
+        return false;
+    }
+    if(!ESP.flashRead(partition->address, (uint32_t*)buf, 4)) {
+        return false;
+    }
+    buf[0] = ESP_IMAGE_HEADER_MAGIC;
+
+    return ESP.flashWrite(partition->address, (uint32_t*)buf, 4);
+}
+
 UpdateClass::UpdateClass()
 : _error(0)
 , _buffer(0)
@@ -56,6 +84,22 @@ void UpdateClass::_reset() {
     _command = U_FLASH;
 }
 
+bool UpdateClass::canRollBack(){
+    if(_buffer){ //Update is running
+        return false;
+    }
+    const esp_partition_t* partition = esp_ota_get_next_update_partition(NULL);
+    return _partitionIsBootable(partition);
+}
+
+bool UpdateClass::rollBack(){
+    if(_buffer){ //Update is running
+        return false;
+    }
+    const esp_partition_t* partition = esp_ota_get_next_update_partition(NULL);
+    return _partitionIsBootable(partition) && !esp_ota_set_boot_partition(partition);
+}
+
 bool UpdateClass::begin(size_t size, int command) {
     if(_size > 0){
         log_w("already running");
@@ -121,6 +165,17 @@ void UpdateClass::abort(){
 }
 
 bool UpdateClass::_writeBuffer(){
+    //first bytes of new firmware
+    if(!_progress && _command == U_FLASH){
+        //check magic
+        if(_buffer[0] != ESP_IMAGE_HEADER_MAGIC){
+            _abort(UPDATE_ERROR_MAGIC_BYTE);
+            return false;
+        }
+        //remove magic byte from the firmware now and write it upon success
+        //this ensures that partially written firmware will not be bootable
+        _buffer[0] = 0xFF;
+    }
     if(!ESP.flashEraseSector((_partition->address + _progress)/SPI_FLASH_SEC_SIZE)){
         _abort(UPDATE_ERROR_ERASE);
         return false;
@@ -129,6 +184,10 @@ bool UpdateClass::_writeBuffer(){
         _abort(UPDATE_ERROR_WRITE);
         return false;
     }
+    //restore magic or md5 will fail
+    if(!_progress && _command == U_FLASH){
+        _buffer[0] = ESP_IMAGE_HEADER_MAGIC;
+    }
     _md5.add(_buffer, _bufferLen);
     _progress += _bufferLen;
     _bufferLen = 0;
@@ -150,17 +209,11 @@ bool UpdateClass::_verifyHeader(uint8_t data) {
 
 bool UpdateClass::_verifyEnd() {
     if(_command == U_FLASH) {
-        uint8_t buf[4];
-        if(!ESP.flashRead(_partition->address, (uint32_t*)buf, 4)) {
+        if(!_partitionIsBootable(_partition)) {
             _abort(UPDATE_ERROR_READ);
             return false;
         }
 
-        if(buf[0] != ESP_IMAGE_HEADER_MAGIC) {
-            _abort(UPDATE_ERROR_MAGIC_BYTE);
-            return false;
-        }
-
         if(esp_ota_set_boot_partition(_partition)){
             _abort(UPDATE_ERROR_ACTIVATE);
             return false;
@@ -168,6 +221,7 @@ bool UpdateClass::_verifyEnd() {
         _reset();
         return true;
     } else if(_command == U_SPIFFS) {
+        _reset();
         return true;
     }
     return false;