From 2633fc3c7412d12a081d520f758f10d720988d5e Mon Sep 17 00:00:00 2001
From: Clemens Kirchgatterer <clemens@1541.org>
Date: Thu, 18 Jan 2018 00:02:58 +0100
Subject: [PATCH] Add progress callback to Update::writeStream(). (#948)

---
 libraries/Update/src/Update.h    | 11 +++++++++++
 libraries/Update/src/Updater.cpp | 16 +++++++++++++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/libraries/Update/src/Update.h b/libraries/Update/src/Update.h
index 6cdcff24..42fa591b 100644
--- 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
index 40d5f953..ad35a8d7 100644
--- 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;
 }