From cadbad88506ac49f710439580975de2ed7855dc9 Mon Sep 17 00:00:00 2001
From: thewavelength <thewavelength@users.noreply.github.com>
Date: Mon, 2 Nov 2020 17:49:24 +0100
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. (#4442)

The UpdateClass in the Updater component has the ability to update data to
a SPIFFS partition. It selects the first available partition using the
ESP-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 the
new 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 884c9d34..fe85580f 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 db0ead63..cd0ba05f 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 6d025d5c..15c7624c 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 bc061127..e18a5364 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;