From 884e417a4943ae64897fc06db47915f24c00b59a Mon Sep 17 00:00:00 2001
From: Jeroen88 <info@el-mago.nl>
Date: Thu, 6 Dec 2018 20:39:52 +0100
Subject: [PATCH] Fix HTTPUpdate flash size check and add SPIFFS size check
 (#2161)

* Fix error in PR #2048: if ::available() is called before ::connect() _rxBuffer is not initialised

* Fixed flash size check and added SPIFFS size check

* Rewriting ESP.getFreeSketchSpace(), moving code from HTTPUpdate.cpp
---
 cores/esp32/Esp.cpp                     |  7 +++++-
 libraries/HTTPUpdate/src/HTTPUpdate.cpp | 29 +++++++++++++++++++------
 libraries/HTTPUpdate/src/HTTPUpdate.h   |  1 +
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/cores/esp32/Esp.cpp b/cores/esp32/Esp.cpp
index 7f57973ea..9e5dbc252 100644
--- a/cores/esp32/Esp.cpp
+++ b/cores/esp32/Esp.cpp
@@ -169,7 +169,12 @@ uint32_t EspClass::getSketchSize () {
 }
 
 uint32_t EspClass::getFreeSketchSpace () {
-    return sketchSize(SKETCH_SIZE_FREE);
+    const esp_partition_t* _partition = esp_ota_get_next_update_partition(NULL);
+    if(!_partition){
+        return 0;
+    }
+
+    return _partition->size;
 }
 
 uint8_t EspClass::getChipRevision(void)
diff --git a/libraries/HTTPUpdate/src/HTTPUpdate.cpp b/libraries/HTTPUpdate/src/HTTPUpdate.cpp
index af263936f..dd4f86236 100644
--- a/libraries/HTTPUpdate/src/HTTPUpdate.cpp
+++ b/libraries/HTTPUpdate/src/HTTPUpdate.cpp
@@ -127,6 +127,8 @@ String HTTPUpdate::getLastErrorString(void)
         return "Verify Bin Header Failed";
     case HTTP_UE_BIN_FOR_WRONG_FLASH:
         return "New Binary Does Not Fit Flash Size";
+    case HTTP_UE_NO_PARTITION:
+        return "Partition Could Not be Found";
     }
 
     return String();
@@ -238,14 +240,25 @@ HTTPUpdateResult HTTPUpdate::handleUpdate(HTTPClient& http, const String& curren
         if(len > 0) {
             bool startUpdate = true;
             if(spiffs) {
-// To do                size_t spiffsSize = ((size_t) &_SPIFFS_end - (size_t) &_SPIFFS_start);
-// To do                if(len > (int) spiffsSize) {
-// To do                    log_e("spiffsSize to low (%d) needed: %d\n", spiffsSize, len);
-// To do                    startUpdate = false;
-// To do                }
+                const esp_partition_t* _partition = esp_partition_find_first(ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_DATA_SPIFFS, NULL);
+                if(!_partition){
+                    _lastError = HTTP_UE_NO_PARTITION;
+                    return HTTP_UPDATE_FAILED;
+                }
+
+                if(len > _partition->size) {
+                    log_e("spiffsSize to low (%d) needed: %d\n", _partition->size, len);
+                    startUpdate = false;
+                }
             } else {
-                if(len > (int) ESP.getFreeSketchSpace()) {
-                    log_e("FreeSketchSpace to low (%d) needed: %d\n", ESP.getFreeSketchSpace(), len);
+                int sketchFreeSpace = ESP.getFreeSketchSpace();
+                if(!sketchFreeSpace){
+                    _lastError = HTTP_UE_NO_PARTITION;
+                    return HTTP_UPDATE_FAILED;
+                }
+
+                if(len > sketchFreeSpace) {
+                    log_e("FreeSketchSpace to low (%d) needed: %d\n", sketchFreeSpace, len);
                     startUpdate = false;
                 }
             }
@@ -375,6 +388,8 @@ bool HTTPUpdate::runUpdate(Stream& in, uint32_t size, String md5, int command)
         }
     }
 
+// To do: the SHA256 could be checked if the server sends it
+
     if(Update.writeStream(in) != size) {
         _lastError = Update.getError();
         Update.printError(error);
diff --git a/libraries/HTTPUpdate/src/HTTPUpdate.h b/libraries/HTTPUpdate/src/HTTPUpdate.h
index 33b778401..f126cba06 100644
--- a/libraries/HTTPUpdate/src/HTTPUpdate.h
+++ b/libraries/HTTPUpdate/src/HTTPUpdate.h
@@ -42,6 +42,7 @@
 #define HTTP_UE_SERVER_FAULTY_MD5           (-105)
 #define HTTP_UE_BIN_VERIFY_HEADER_FAILED    (-106)
 #define HTTP_UE_BIN_FOR_WRONG_FLASH         (-107)
+#define HTTP_UE_NO_PARTITION                (-108)
 
 enum HTTPUpdateResult {
     HTTP_UPDATE_FAILED,
-- 
GitLab