From 8f42e15ee9df9b402cfb123a5f1117620099e95e Mon Sep 17 00:00:00 2001
From: thewavelength <github@cb92.de>
Date: Sat, 24 Oct 2020 23:16:17 +0200
Subject: [PATCH] Add partition label argument to Update and ArduinoOTA
 classThe UpdateClass in the Updater component has the ability to update data
 toa SPIFFS partition. It selects the first available partition using
 theESP-IDF esp_partition_find_first() function.That behaviour is problematic
 if one has multiple SPIFFS partitions.This change allows a user to pass the
 label argument (defaults to NULL)to UpdateClass::begin() so a specific SPIFFS
 partition can be updated.Additionally, ArduinoOTA can set this partition
 label using thenew method ArduinoOTAClass::setPartitionLabel which is
 optional.This change does not break compatibility.

---
 libraries/ArduinoOTA/src/ArduinoOTA.cpp | 16 ++++++++++++++--
 libraries/ArduinoOTA/src/ArduinoOTA.h   |  5 +++++
 libraries/Update/src/Update.h           |  2 +-
 libraries/Update/src/Updater.cpp        |  4 ++--
 4 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/libraries/ArduinoOTA/src/ArduinoOTA.cpp b/libraries/ArduinoOTA/src/ArduinoOTA.cpp
index 884c9d341c1..fe85580f566 100644
--- a/libraries/ArduinoOTA/src/ArduinoOTA.cpp
+++ b/libraries/ArduinoOTA/src/ArduinoOTA.cpp
@@ -88,6 +88,17 @@ ArduinoOTAClass& ArduinoOTAClass::setPasswordHash(const char * password) {
     return *this;
 }
 
+ArduinoOTAClass& ArduinoOTAClass::setPartitionLabel(const char * partition_label) {
+    if (!_initialized && !_partition_label.length() && partition_label) {
+        _partition_label = partition_label;
+    }
+    return *this;
+}
+
+String ArduinoOTAClass::getPartitionLabel() {
+    return _partition_label;
+}
+
 ArduinoOTAClass& ArduinoOTAClass::setRebootOnSuccess(bool reboot){
     _rebootOnSuccess = reboot;
     return *this;
@@ -235,7 +246,8 @@ void ArduinoOTAClass::_onRx(){
 }
 
 void ArduinoOTAClass::_runUpdate() {
-    if (!Update.begin(_size, _cmd)) {
+    const char *partition_label = _partition_label.length() ? _partition_label.c_str() : NULL;
+    if (!Update.begin(_size, _cmd, -1, LOW, partition_label)) {
 
         log_e("Begin ERROR: %s", Update.errorString());
 
@@ -358,7 +370,7 @@ void ArduinoOTAClass::end() {
 
 void ArduinoOTAClass::handle() {
     if (!_initialized) {
-        return; 
+        return;
     }
     if (_state == OTA_RUNUPDATE) {
         _runUpdate();
diff --git a/libraries/ArduinoOTA/src/ArduinoOTA.h b/libraries/ArduinoOTA/src/ArduinoOTA.h
index db0ead631db..cd0ba05f547 100644
--- a/libraries/ArduinoOTA/src/ArduinoOTA.h
+++ b/libraries/ArduinoOTA/src/ArduinoOTA.h
@@ -44,6 +44,10 @@ class ArduinoOTAClass
     //Sets the password as above but in the form MD5(password). Default NULL
     ArduinoOTAClass& setPasswordHash(const char *password);
 
+    //Sets the partition label to write to when updating SPIFFS. Default NULL
+    ArduinoOTAClass &setPartitionLabel(const char *partition_label);
+    String getPartitionLabel();
+
     //Sets if the device should be rebooted after successful update. Default true
     ArduinoOTAClass& setRebootOnSuccess(bool reboot);
 
@@ -80,6 +84,7 @@ class ArduinoOTAClass
     int _port;
     String _password;
     String _hostname;
+    String _partition_label;
     String _nonce;
     WiFiUDP _udp_ota;
     bool _initialized;
diff --git a/libraries/Update/src/Update.h b/libraries/Update/src/Update.h
index 6d025d5c4cd..15c7624c46d 100644
--- a/libraries/Update/src/Update.h
+++ b/libraries/Update/src/Update.h
@@ -43,7 +43,7 @@ class UpdateClass {
       Call this to check the space needed for the update
       Will return false if there is not enough space
     */
-    bool begin(size_t size=UPDATE_SIZE_UNKNOWN, int command = U_FLASH, int ledPin = -1, uint8_t ledOn = LOW);
+    bool begin(size_t size=UPDATE_SIZE_UNKNOWN, int command = U_FLASH, int ledPin = -1, uint8_t ledOn = LOW, const char *label = NULL);
 
     /*
       Writes a buffer to the flash and increments the address
diff --git a/libraries/Update/src/Updater.cpp b/libraries/Update/src/Updater.cpp
index bc061127510..e18a53644ec 100644
--- a/libraries/Update/src/Updater.cpp
+++ b/libraries/Update/src/Updater.cpp
@@ -104,7 +104,7 @@ bool UpdateClass::rollBack(){
     return _partitionIsBootable(partition) && !esp_ota_set_boot_partition(partition);
 }
 
-bool UpdateClass::begin(size_t size, int command, int ledPin, uint8_t ledOn) {
+bool UpdateClass::begin(size_t size, int command, int ledPin, uint8_t ledOn, const char *label) {
     if(_size > 0){
         log_w("already running");
         return false;
@@ -132,7 +132,7 @@ bool UpdateClass::begin(size_t size, int command, int ledPin, uint8_t ledOn) {
         log_d("OTA Partition: %s", _partition->label);
     }
     else if (command == U_SPIFFS) {
-        _partition = esp_partition_find_first(ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_DATA_SPIFFS, NULL);
+        _partition = esp_partition_find_first(ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_DATA_SPIFFS, label);
         if(!_partition){
             _error = UPDATE_ERROR_NO_PARTITION;
             return false;