From 45d71ae4bd471284c60bb3b654535a94cdb3a8ed Mon Sep 17 00:00:00 2001
From: Robin Richtsfeld <robin.richtsfeld@gmail.com>
Date: Wed, 16 Oct 2019 10:11:15 +0200
Subject: [PATCH] Support FS update in two steps (#6505)

* Support FS update in two steps
---
 cores/esp8266/Updater.cpp                     | 47 +++++++++++++++----
 doc/ota_updates/readme.rst                    |  2 +
 .../src/ESP8266httpUpdate.cpp                 |  4 ++
 3 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/cores/esp8266/Updater.cpp b/cores/esp8266/Updater.cpp
index d6b6b621f..b028e19b9 100644
--- a/cores/esp8266/Updater.cpp
+++ b/cores/esp8266/Updater.cpp
@@ -23,6 +23,7 @@ extern "C" {
 }
 
 extern "C" uint32_t _FS_start;
+extern "C" uint32_t _FS_end;
 
 UpdaterClass::UpdaterClass()
 : _async(false)
@@ -105,15 +106,17 @@ bool UpdaterClass::begin(size_t size, int command, int ledPin, uint8_t ledOn) {
 
   wifi_set_sleep_type(NONE_SLEEP_T);
 
+  //address where we will start writing the update
   uintptr_t updateStartAddress = 0;
+  //size of current sketch rounded to a sector
+  size_t currentSketchSize = (ESP.getSketchSize() + FLASH_SECTOR_SIZE - 1) & (~(FLASH_SECTOR_SIZE - 1));
+  //size of the update rounded to a sector
+  size_t roundedSize = (size + FLASH_SECTOR_SIZE - 1) & (~(FLASH_SECTOR_SIZE - 1));
+
   if (command == U_FLASH) {
-    //size of current sketch rounded to a sector
-    size_t currentSketchSize = (ESP.getSketchSize() + FLASH_SECTOR_SIZE - 1) & (~(FLASH_SECTOR_SIZE - 1));
     //address of the end of the space available for sketch and update
     uintptr_t updateEndAddress = (uintptr_t)&_FS_start - 0x40200000;
-    //size of the update rounded to a sector
-    size_t roundedSize = (size + FLASH_SECTOR_SIZE - 1) & (~(FLASH_SECTOR_SIZE - 1));
-    //address where we will start writing the update
+
     updateStartAddress = (updateEndAddress > roundedSize)? (updateEndAddress - roundedSize) : 0;
 
 #ifdef DEBUG_UPDATER
@@ -129,7 +132,24 @@ bool UpdaterClass::begin(size_t size, int command, int ledPin, uint8_t ledOn) {
     }
   }
   else if (command == U_FS) {
-     updateStartAddress = (uintptr_t)&_FS_start - 0x40200000;
+    if((uintptr_t)&_FS_start + roundedSize > (uintptr_t)&_FS_end) {
+      _setError(UPDATE_ERROR_SPACE);
+      return false;
+    }
+
+#ifdef ATOMIC_FS_UPDATE
+    //address of the end of the space available for update
+    uintptr_t updateEndAddress = (uintptr_t)&_FS_start - 0x40200000;
+
+    updateStartAddress = (updateEndAddress > roundedSize)? (updateEndAddress - roundedSize) : 0;
+
+    if(updateStartAddress < currentSketchSize) {
+      _setError(UPDATE_ERROR_SPACE);
+      return false;
+    }
+#else
+    updateStartAddress = (uintptr_t)&_FS_start - 0x40200000;
+#endif
   }
   else {
     // unknown command
@@ -272,8 +292,19 @@ bool UpdaterClass::end(bool evenIfRemaining){
 
 #ifdef DEBUG_UPDATER
     DEBUG_UPDATER.printf_P(PSTR("Staged: address:0x%08X, size:0x%08zX\n"), _startAddress, _size);
+#endif
   }
   else if (_command == U_FS) {
+#ifdef ATOMIC_FS_UPDATE
+    eboot_command ebcmd;
+    ebcmd.action = ACTION_COPY_RAW;
+    ebcmd.args[0] = _startAddress;
+    ebcmd.args[1] = (uintptr_t)&_FS_start - 0x40200000;
+    ebcmd.args[2] = _size;
+    eboot_command_write(&ebcmd);
+#endif
+
+#ifdef DEBUG_UPDATER
     DEBUG_UPDATER.printf_P(PSTR("Filesystem: address:0x%08X, size:0x%08zX\n"), _startAddress, _size);
 #endif
   }
@@ -387,7 +418,7 @@ bool UpdaterClass::_verifyHeader(uint8_t data) {
         }
         return true;
     } else if(_command == U_FS) {
-        // no check of SPIFFS possible with first byte.
+        // no check of FS possible with first byte.
         return true;
     }
     return false;
@@ -421,7 +452,7 @@ bool UpdaterClass::_verifyEnd() {
 
         return true;
     } else if(_command == U_FS) {
-        // SPIFFS is already over written checks make no sense any more.
+        // FS is already over written checks make no sense any more.
         return true;
     }
     return false;
diff --git a/doc/ota_updates/readme.rst b/doc/ota_updates/readme.rst
index 67f4f9110..f5c76f0d2 100644
--- a/doc/ota_updates/readme.rst
+++ b/doc/ota_updates/readme.rst
@@ -646,6 +646,8 @@ Update process - memory view
 -  the new sketch is now copied "over" the old one.
 -  the new sketch is started.
 
+By default, filesystem updates are overriding the target flash directly. In order to use the same two step process set the `ATOMIC_FS_UPDATE` flag.
+
 .. figure:: update_memory_copy.png
    :alt: Memory layout for OTA updates
 
diff --git a/libraries/ESP8266httpUpdate/src/ESP8266httpUpdate.cpp b/libraries/ESP8266httpUpdate/src/ESP8266httpUpdate.cpp
index 0090390a7..8787562d6 100644
--- a/libraries/ESP8266httpUpdate/src/ESP8266httpUpdate.cpp
+++ b/libraries/ESP8266httpUpdate/src/ESP8266httpUpdate.cpp
@@ -389,7 +389,11 @@ HTTPUpdateResult ESP8266HTTPUpdate::handleUpdate(HTTPClient& http, const String&
                     DEBUG_HTTP_UPDATE("[httpUpdate] Update ok\n");
                     http.end();
 
+#ifdef ATOMIC_FS_UPDATE
+                    if(_rebootOnUpdate) {
+#else
                     if(_rebootOnUpdate && !spiffs) {
+#endif
                         ESP.restart();
                     }