Updater lifetime callbacks (#8653)

follow-up of #8598 
similar to ArduinoOTA API, execute certain callback in the Updater context.

Author: mcspr

cores/esp8266/Updater.cpp

@@ -44,21 +44,22 @@ UpdaterClass::~UpdaterClass()
 #endif
 }
 
-UpdaterClass& UpdaterClass::onProgress(THandlerFunction_Progress fn) {
-    _progress_callback = fn;
-    return *this;
-}
-
-void UpdaterClass::_reset() {
-  if (_buffer)
+void UpdaterClass::_reset(bool callback) {
+  if (_buffer) {
     delete[] _buffer;
-  _buffer = 0;
+  }
+
+  _buffer = nullptr;
   _bufferLen = 0;
   _startAddress = 0;
   _currentAddress = 0;
   _size = 0;
   _command = U_FLASH;
 
+  if (callback && _end_callback) {
+    _end_callback();
+  }
+
   if(_ledPin != -1) {
     digitalWrite(_ledPin, !_ledOn); // off
   }
@@ -173,7 +174,13 @@ bool UpdaterClass::begin(size_t size, int command, int ledPin, uint8_t ledOn) {
   } else {
     _bufferSize = 256;
   }
-  _buffer = new uint8_t[_bufferSize];
+  _buffer = new (std::nothrow) uint8_t[_bufferSize];
+  if (!_buffer) {
+    _setError(UPDATE_ERROR_OOM);
+    _reset(false);
+    return false;
+  }
+
   _command = command;
 
 #ifdef DEBUG_UPDATER
@@ -185,6 +192,11 @@ bool UpdaterClass::begin(size_t size, int command, int ledPin, uint8_t ledOn) {
   if (!_verify) {
     _md5.begin();
   }
+
+  if (_start_callback) {
+    _start_callback();
+  }
+
   return true;
 }
 
@@ -293,7 +305,7 @@ bool UpdaterClass::end(bool evenIfRemaining){
       const uint32_t sigAddr = _startAddress + binSize;
       sig.reset(new (std::nothrow) uint8_t[sigLen]);
       if (!sig) {
-        _setError(UPDATE_ERROR_SIGN);
+        _setError(UPDATE_ERROR_OOM);
         _reset();
         return false;
       }
@@ -574,45 +586,79 @@ size_t UpdaterClass::writeStream(Stream &data, uint16_t streamTimeout) {
 
 void UpdaterClass::_setError(int error){
   _error = error;
+  if (_error_callback) {
+    _error_callback(error);
+  }
 #ifdef DEBUG_UPDATER
   printError(DEBUG_UPDATER);
 #endif
   _reset(); // Any error condition invalidates the entire update, so clear partial status
 }
 
-void UpdaterClass::printError(Print &out){
-  out.printf_P(PSTR("ERROR[%u]: "), _error);
-  if(_error == UPDATE_ERROR_OK){
-    out.println(F("No Error"));
-  } else if(_error == UPDATE_ERROR_WRITE){
-    out.println(F("Flash Write Failed"));
-  } else if(_error == UPDATE_ERROR_ERASE){
-    out.println(F("Flash Erase Failed"));
-  } else if(_error == UPDATE_ERROR_READ){
-    out.println(F("Flash Read Failed"));
-  } else if(_error == UPDATE_ERROR_SPACE){
-    out.println(F("Not Enough Space"));
-  } else if(_error == UPDATE_ERROR_SIZE){
-    out.println(F("Bad Size Given"));
-  } else if(_error == UPDATE_ERROR_STREAM){
-    out.println(F("Stream Read Timeout"));
-  } else if(_error == UPDATE_ERROR_NO_DATA){
-    out.println(F("No data supplied"));
-  } else if(_error == UPDATE_ERROR_MD5){
-    out.printf_P(PSTR("MD5 Failed: expected:%s, calculated:%s\n"), _target_md5.c_str(), _md5.toString().c_str());
-  } else if(_error == UPDATE_ERROR_SIGN){
-    out.println(F("Signature verification failed"));
-  } else if(_error == UPDATE_ERROR_FLASH_CONFIG){
-    out.printf_P(PSTR("Flash config wrong real: %d IDE: %d\n"), ESP.getFlashChipRealSize(), ESP.getFlashChipSize());
-  } else if(_error == UPDATE_ERROR_NEW_FLASH_CONFIG){
-    out.printf_P(PSTR("new Flash config wrong real: %d\n"), ESP.getFlashChipRealSize());
-  } else if(_error == UPDATE_ERROR_MAGIC_BYTE){
-    out.println(F("Magic byte is wrong, not 0xE9"));
-  } else if (_error == UPDATE_ERROR_BOOTSTRAP){
-    out.println(F("Invalid bootstrapping state, reset ESP8266 before updating"));
-  } else {
-    out.println(F("UNKNOWN"));
+String UpdaterClass::getErrorString() const {
+  String out;
+
+  switch (_error) {
+  case UPDATE_ERROR_OK:
+    out = F("No Error");
+    break;
+  case UPDATE_ERROR_WRITE:
+    out = F("Flash Write Failed");
+    break;
+  case UPDATE_ERROR_ERASE:
+    out = F("Flash Erase Failed");
+    break;
+  case UPDATE_ERROR_READ:
+    out = F("Flash Read Failed");
+    break;
+  case UPDATE_ERROR_SPACE:
+    out = F("Not Enough Space");
+    break;
+  case UPDATE_ERROR_SIZE:
+    out = F("Bad Size Given");
+    break;
+  case UPDATE_ERROR_STREAM:
+    out = F("Stream Read Timeout");
+    break;
+  case UPDATE_ERROR_MD5:
+    out += F("MD5 verification failed: ");
+    out += F("expected: ") + _target_md5;
+    out += F(", calculated: ") + _md5.toString();
+    break;
+  case UPDATE_ERROR_FLASH_CONFIG:
+    out += F("Flash config wrong: ");
+    out += F("real: ") + String(ESP.getFlashChipRealSize(), 10);
+    out += F(", SDK: ") + String(ESP.getFlashChipSize(), 10);
+    break;
+  case UPDATE_ERROR_NEW_FLASH_CONFIG:
+    out += F("new Flash config wrong, real size: ");
+    out += String(ESP.getFlashChipRealSize(), 10);
+    break;
+  case UPDATE_ERROR_MAGIC_BYTE:
+    out = F("Magic byte is not 0xE9");
+    break;
+  case UPDATE_ERROR_BOOTSTRAP:
+    out = F("Invalid bootstrapping state, reset ESP8266 before updating");
+    break;
+  case UPDATE_ERROR_SIGN:
+    out = F("Signature verification failed");
+    break;
+  case UPDATE_ERROR_NO_DATA:
+    out = F("No data supplied");
+    break;
+  case UPDATE_ERROR_OOM:
+    out = F("Out of memory");
+    break;
+  default:
+    out = F("UNKNOWN");
+    break;
   }
+
+  return out;
+}
+
+void UpdaterClass::printError(Print &out){
+  out.printf_P(PSTR("ERROR[%hhu]: %s\n"), _error, getErrorString().c_str());
 }
 
 UpdaterClass Update;

cores/esp8266/Updater.h

@@ -20,6 +20,7 @@
 #define UPDATE_ERROR_BOOTSTRAP          (11)
 #define UPDATE_ERROR_SIGN               (12)
 #define UPDATE_ERROR_NO_DATA            (13)
+#define UPDATE_ERROR_OOM                (14)
 
 #define U_FLASH   0
 #define U_FS      100
@@ -51,7 +52,9 @@ class UpdaterVerifyClass {
 
 class UpdaterClass {
   public:
-    typedef std::function<void(size_t, size_t)> THandlerFunction_Progress;
+    using THandlerFunction_Progress = std::function<void(size_t, size_t)>;
+    using THandlerFunction_Error = std::function<void(uint8_t)>;
+    using THandlerFunction = std::function<void()>;
 
     UpdaterClass();
     ~UpdaterClass();
@@ -97,6 +100,11 @@ class UpdaterClass {
     */
     bool end(bool evenIfRemaining = false);
 
+    /*
+      Gets the last error description as string
+    */
+    String getErrorString() const;
+
     /*
       Prints the last error to an output stream
     */
@@ -120,7 +128,34 @@ class UpdaterClass {
     /*
       This callback will be called when Updater is receiving data
     */
-    UpdaterClass& onProgress(THandlerFunction_Progress fn);
+    UpdaterClass& onProgress(THandlerFunction_Progress fn) {
+      _progress_callback = std::move(fn);
+      return *this;
+    }
+
+    /*
+      This callback will be called when Updater ends
+    */
+    UpdaterClass& onError(THandlerFunction_Error fn) {
+      _error_callback = std::move(fn);
+      return *this;
+    }
+
+    /*
+      This callback will be called when Updater begins
+    */
+    UpdaterClass& onStart(THandlerFunction fn) {
+      _start_callback = std::move(fn);
+      return *this;
+    }
+
+    /*
+      This callback will be called when Updater ends
+    */
+    UpdaterClass& onEnd(THandlerFunction fn) {
+      _end_callback = std::move(fn);
+      return *this;
+    }
 
     //Helpers
     uint8_t getError(){ return _error; }
@@ -175,7 +210,7 @@ class UpdaterClass {
     }
 
   private:
-    void _reset();
+    void _reset(bool callback = true);
     bool _writeBuffer();
 
     bool _verifyHeader(uint8_t data);
@@ -202,8 +237,12 @@ class UpdaterClass {
     // Optional signed binary verification
     UpdaterHashClass *_hash = nullptr;
     UpdaterVerifyClass *_verify = nullptr;
-    // Optional progress callback function
+
+    // Optional lifetime callback functions
     THandlerFunction_Progress _progress_callback = nullptr;
+    THandlerFunction_Error _error_callback = nullptr;
+    THandlerFunction _start_callback = nullptr;
+    THandlerFunction _end_callback = nullptr;
 };
 
 extern UpdaterClass Update;

doc/ota_updates/readme.rst

@@ -668,9 +668,29 @@ Updater class
 
 Updater is in the Core and deals with writing the firmware to the flash, checking its integrity and telling the bootloader (eboot) to load the new firmware on the next boot.
 
-**Note:** The bootloader command will be stored into the first 128 bytes of user RTC memory, then it will be retrieved by eboot on boot. That means that user data present there will be lost `(per discussion in #5330) <https://github.com/esp8266/Arduino/pull/5330#issuecomment-437803456>`__.
+The following `Updater <https://github.com/esp8266/Arduino/tree/master/cores/esp8266/Updater.h` methods could be used to be notified about OTA progress:
 
-**Note:** For uncompressed firmware images, the Updater will change the flash mode bits if they differ from the flash mode the device is currently running at. This ensures that the flash mode is not changed to an incompatible mode when the device is in a remote or hard to access area. Compressed images are not modified, thus changing the flash mode in this instance could result in damage to the ESP8266 and/or flash memory chip or your device no longer be accessible via OTA, and requiring re-flashing via a serial connection `(per discussion in #7307) <https://github.com/esp8266/Arduino/issues/7307#issuecomment-631523053>`__.
+.. code:: cpp
+
+    using THandlerFunction_Progress = std::function<void(size_t, size_t)>;
+    void onProgress(THandlerFunction_Progress); // current and total number of bytes
+
+    using THandlerFunction_Error = std::function<void(uint8_t)>;
+    void onStart(THandlerFunction_Error); // error code
+
+    using THandlerFunction = std::function<void()>;
+    void onEnd(THandlerFunction);
+    void onError(THandlerFunction);
+
+Using RTC memory
+~~~~~~~~~~~~~~~~
+
+The bootloader command will be stored into the first 128 bytes of user RTC memory, then it will be retrieved by eboot on boot. That means that user data present there will be lost `(per discussion in #5330) <https://github.com/esp8266/Arduino/pull/5330#issuecomment-437803456>`__.
+
+Flash mode and size
+~~~~~~~~~~~~~~~~~~~
+
+For uncompressed firmware images, the Updater will change the flash mode bits if they differ from the flash mode the device is currently running at. This ensures that the flash mode is not changed to an incompatible mode when the device is in a remote or hard to access area. Compressed images are not modified, thus changing the flash mode in this instance could result in damage to the ESP8266 and/or flash memory chip or your device no longer be accessible via OTA, and requiring re-flashing via a serial connection `(per discussion in #7307) <https://github.com/esp8266/Arduino/issues/7307#issuecomment-631523053>`__.
 
 Update process - memory view
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~