Updater lifetime callbacks

cores/esp8266/Updater.cpp

@@ -41,21 +41,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
   }
@@ -170,7 +171,15 @@ 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) {
+#ifdef DEBUG_UPDATER
+    DEBUG_UPDATER.println(F("[begin] Unable to allocate a temporary buffer."));
+#endif
+    _reset(false);
+    return false;
+  }
+
   _command = command;
 
 #ifdef DEBUG_UPDATER
@@ -182,6 +191,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;
 }
 
@@ -555,45 +569,59 @@ 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);
+String UpdaterClass::getErrorString() const {
+  String out;
+
   if(_error == UPDATE_ERROR_OK){
-    out.println(F("No Error"));
+    out = F("No Error");
   } else if(_error == UPDATE_ERROR_WRITE){
-    out.println(F("Flash Write Failed"));
+    out = F("Flash Write Failed");
   } else if(_error == UPDATE_ERROR_ERASE){
-    out.println(F("Flash Erase Failed"));
+    out = F("Flash Erase Failed");
   } else if(_error == UPDATE_ERROR_READ){
-    out.println(F("Flash Read Failed"));
+    out = F("Flash Read Failed");
   } else if(_error == UPDATE_ERROR_SPACE){
-    out.println(F("Not Enough Space"));
+    out = F("Not Enough Space");
   } else if(_error == UPDATE_ERROR_SIZE){
-    out.println(F("Bad Size Given"));
+    out = F("Bad Size Given");
   } else if(_error == UPDATE_ERROR_STREAM){
-    out.println(F("Stream Read Timeout"));
+    out = F("Stream Read Timeout");
   } else if(_error == UPDATE_ERROR_NO_DATA){
-    out.println(F("No data supplied"));
+    out = 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());
+    out += F("MD5 verification failed: ");
+    out += F("expected: ") + _target_md5;
+    out += F(", calculated: ") + _md5.toString();
   } else if(_error == UPDATE_ERROR_SIGN){
-    out.println(F("Signature verification failed"));
+    out = 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());
+    out += F("Flash config wrong: ");
+    out += F("real: ") + String(ESP.getFlashChipRealSize(), 10);
+    out += F(", SDK: ") + String(ESP.getFlashChipSize(), 10);
   } else if(_error == UPDATE_ERROR_NEW_FLASH_CONFIG){
-    out.printf_P(PSTR("new Flash config wrong real: %d\n"), ESP.getFlashChipRealSize());
+    out += F("new Flash config wrong, real size: ") + String(ESP.getFlashChipRealSize(), 10);
   } else if(_error == UPDATE_ERROR_MAGIC_BYTE){
-    out.println(F("Magic byte is wrong, not 0xE9"));
+    out = F("Magic byte is not 0xE9");
   } else if (_error == UPDATE_ERROR_BOOTSTRAP){
-    out.println(F("Invalid bootstrapping state, reset ESP8266 before updating"));
+    out = F("Invalid bootstrapping state, reset ESP8266 before updating");
   } else {
-    out.println(F("UNKNOWN"));
+    out = F("UNKNOWN");
   }
+
+  return out;
+}
+
+void UpdaterClass::printError(Print &out){
+  out.printf_P(PSTR("ERROR[%u]: %s\n"), _error, getErrorString().c_str());
 }
 
 UpdaterClass Update;

cores/esp8266/Updater.h

@@ -51,7 +51,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 +99,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 +127,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 +209,7 @@ class UpdaterClass {
     }
 
   private:
-    void _reset();
+    void _reset(bool callback = true);
     bool _writeBuffer();
 
     bool _verifyHeader(uint8_t data);
@@ -202,8 +236,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;