Add progress callback to Update::writeStream(). (espressif#948)

everslick · me-no-dev · commit 2633fc3c7412 · 2018-01-18T01:02:58.000+02:00
diff --git a/libraries/Update/src/Update.h b/libraries/Update/src/Update.h
@@ -3,6 +3,7 @@
 
 #include <Arduino.h>
 #include <MD5Builder.h>
+#include <functional>
 #include "esp_partition.h"
 
 #define UPDATE_ERROR_OK                 (0)
@@ -27,7 +28,15 @@
 
 class UpdateClass {
   public:
+    typedef std::function<void(size_t, size_t)> THandlerFunction_Progress;
+
     UpdateClass();
+
+    /*
+      This callback will be called when Update is receiving data
+    */
+    UpdateClass& onProgress(THandlerFunction_Progress fn);
+
     /*
       Call this to check the space needed for the update
       Will return false if there is not enough space
@@ -153,6 +162,8 @@ class UpdateClass {
     bool _verifyHeader(uint8_t data);
     bool _verifyEnd();
 
+    THandlerFunction_Progress _progress_callback;
+
     uint8_t _error;
     uint8_t *_buffer;
     size_t _bufferLen;
diff --git a/libraries/Update/src/Updater.cpp b/libraries/Update/src/Updater.cpp
@@ -71,9 +71,15 @@ UpdateClass::UpdateClass()
 , _progress(0)
 , _command(U_FLASH)
 , _partition(NULL)
+, _progress_callback(NULL)
 {
 }
 
+UpdateClass& UpdateClass::onProgress(THandlerFunction_Progress fn) {
+    _progress_callback = fn;
+    return *this;
+}
+
 void UpdateClass::_reset() {
     if (_buffer)
         delete[] _buffer;
@@ -306,7 +312,9 @@ size_t UpdateClass::writeStream(Stream &data) {
         _reset();
         return 0;
     }
-
+    if (_progress_callback) {
+        _progress_callback(0, _size);
+    }
     while(remaining()) {
         toRead = data.readBytes(_buffer + _bufferLen,  (SPI_FLASH_SEC_SIZE - _bufferLen));
         if(toRead == 0) { //Timeout
@@ -321,6 +329,12 @@ size_t UpdateClass::writeStream(Stream &data) {
         if((_bufferLen == remaining() || _bufferLen == SPI_FLASH_SEC_SIZE) && !_writeBuffer())
             return written;
         written += toRead;
+        if(_progress_callback) {
+            _progress_callback(_progress, _size);
+        }
+    }
+    if(_progress_callback) {
+        _progress_callback(_size, _size);
     }
     return written;
 }
