Fix ESP32 esp-idf OTA updates (#2424)

* WIP on separating out the OTA backends.

* Split off the individual OTA backends.

* Cleanup the three backends, split into .h and .cpp.

* After successfull flashing, activate the new boot partition.

* Fix linting issues.

* Minor cleanup

Co-authored-by: Maurice Makaay <mmakaay1@xs4all.net>
Co-authored-by: Otto winter <otto@otto-winter.com>

esphome/components/ota/ota_backend.h

@@ -0,0 +1,18 @@
+#pragma once
+#include "ota_component.h"
+
+namespace esphome {
+namespace ota {
+
+class OTABackend {
+ public:
+  virtual ~OTABackend() = default;
+  virtual OTAResponseTypes begin(size_t image_size) = 0;
+  virtual void set_update_md5(const char *md5) = 0;
+  virtual OTAResponseTypes write(uint8_t *data, size_t len) = 0;
+  virtual OTAResponseTypes end() = 0;
+  virtual void abort() = 0;
+};
+
+}  // namespace ota
+}  // namespace esphome

esphome/components/ota/ota_backend_arduino_esp32.cpp

@@ -0,0 +1,46 @@
+#include "esphome/core/defines.h"
+#ifdef USE_ESP32_FRAMEWORK_ARDUINO
+
+#include "ota_backend_arduino_esp32.h"
+#include "ota_component.h"
+#include "ota_backend.h"
+
+#include <Update.h>
+
+namespace esphome {
+namespace ota {
+
+OTAResponseTypes ArduinoESP32OTABackend::begin(size_t image_size) {
+  bool ret = Update.begin(image_size, U_FLASH);
+  if (ret) {
+    return OTA_RESPONSE_OK;
+  }
+
+  uint8_t error = Update.getError();
+  if (error == UPDATE_ERROR_SIZE)
+    return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
+  return OTA_RESPONSE_ERROR_UNKNOWN;
+}
+
+void ArduinoESP32OTABackend::set_update_md5(const char *md5) { Update.setMD5(md5); }
+
+OTAResponseTypes ArduinoESP32OTABackend::write(uint8_t *data, size_t len) {
+  size_t written = Update.write(data, len);
+  if (written != len) {
+    return OTA_RESPONSE_ERROR_WRITING_FLASH;
+  }
+  return OTA_RESPONSE_OK;
+}
+
+OTAResponseTypes ArduinoESP32OTABackend::end() {
+  if (!Update.end())
+    return OTA_RESPONSE_ERROR_UPDATE_END;
+  return OTA_RESPONSE_OK;
+}
+
+void ArduinoESP32OTABackend::abort() { Update.abort(); }
+
+}  // namespace ota
+}  // namespace esphome
+
+#endif  // USE_ESP32_FRAMEWORK_ARDUINO

esphome/components/ota/ota_backend_arduino_esp32.h

@@ -0,0 +1,22 @@
+#pragma once
+#include "esphome/core/defines.h"
+#ifdef USE_ESP32_FRAMEWORK_ARDUINO
+
+#include "ota_component.h"
+#include "ota_backend.h"
+
+namespace esphome {
+namespace ota {
+
+class ArduinoESP32OTABackend : public OTABackend {
+  OTAResponseTypes begin(size_t image_size) override;
+  void set_update_md5(const char *md5) override;
+  OTAResponseTypes write(uint8_t *data, size_t len) override;
+  OTAResponseTypes end() override;
+  void abort() override;
+};
+
+}  // namespace ota
+}  // namespace esphome
+
+#endif  // USE_ESP32_FRAMEWORK_ARDUINO

esphome/components/ota/ota_backend_arduino_esp8266.cpp

@@ -0,0 +1,58 @@
+#include "esphome/core/defines.h"
+#ifdef USE_ARDUINO
+#ifdef USE_ESP8266
+
+#include "ota_backend_arduino_esp8266.h"
+#include "ota_component.h"
+#include "ota_backend.h"
+#include "esphome/components/esp8266/preferences.h"
+
+#include <Updater.h>
+
+namespace esphome {
+namespace ota {
+
+OTAResponseTypes ArduinoESP8266OTABackend::begin(size_t image_size) {
+  bool ret = Update.begin(image_size, U_FLASH);
+  if (ret) {
+    esp8266::preferences_prevent_write(true);
+  }
+
+  uint8_t error = Update.getError();
+  if (error == UPDATE_ERROR_BOOTSTRAP)
+    return OTA_RESPONSE_ERROR_INVALID_BOOTSTRAPPING;
+  if (error == UPDATE_ERROR_NEW_FLASH_CONFIG)
+    return OTA_RESPONSE_ERROR_WRONG_NEW_FLASH_CONFIG;
+  if (error == UPDATE_ERROR_FLASH_CONFIG)
+    return OTA_RESPONSE_ERROR_WRONG_CURRENT_FLASH_CONFIG;
+  if (error == UPDATE_ERROR_SPACE)
+    return OTA_RESPONSE_ERROR_ESP8266_NOT_ENOUGH_SPACE;
+  return OTA_RESPONSE_ERROR_UNKNOWN;
+}
+
+void ArduinoESP8266OTABackend::set_update_md5(const char *md5) { Update.setMD5(md5); }
+
+OTAResponseTypes ArduinoESP8266OTABackend::write(uint8_t *data, size_t len) {
+  size_t written = Update.write(data, len);
+  if (written != len) {
+    return OTA_RESPONSE_ERROR_WRITING_FLASH;
+  }
+  return OTA_RESPONSE_OK;
+}
+
+OTAResponseTypes ArduinoESP8266OTABackend::end() {
+  if (!Update.end())
+    return OTA_RESPONSE_ERROR_UPDATE_END;
+  return OTA_RESPONSE_OK;
+}
+
+void ArduinoESP8266OTABackend::abort() {
+  Update.end();
+  esp8266::preferences_prevent_write(false);
+}
+
+}  // namespace ota
+}  // namespace esphome
+
+#endif
+#endif

esphome/components/ota/ota_backend_arduino_esp8266.h

@@ -0,0 +1,25 @@
+#pragma once
+#include "esphome/core/defines.h"
+#ifdef USE_ARDUINO
+#ifdef USE_ESP8266
+
+#include "ota_component.h"
+#include "ota_backend.h"
+
+namespace esphome {
+namespace ota {
+
+class ArduinoESP8266OTABackend : public OTABackend {
+ public:
+  OTAResponseTypes begin(size_t image_size) override;
+  void set_update_md5(const char *md5) override;
+  OTAResponseTypes write(uint8_t *data, size_t len) override;
+  OTAResponseTypes end() override;
+  void abort() override;
+};
+
+}  // namespace ota
+}  // namespace esphome
+
+#endif
+#endif

esphome/components/ota/ota_backend_esp_idf.cpp

@@ -0,0 +1,72 @@
+#include "esphome/core/defines.h"
+#ifdef USE_ESP_IDF
+
+#include "ota_backend_esp_idf.h"
+#include "ota_component.h"
+#include <esp_ota_ops.h>
+
+namespace esphome {
+namespace ota {
+
+OTAResponseTypes IDFOTABackend::begin(size_t image_size) {
+  this->partition_ = esp_ota_get_next_update_partition(nullptr);
+  if (this->partition_ == nullptr) {
+    return OTA_RESPONSE_ERROR_NO_UPDATE_PARTITION;
+  }
+  esp_err_t err = esp_ota_begin(this->partition_, image_size, &this->update_handle_);
+  if (err != ESP_OK) {
+    esp_ota_abort(this->update_handle_);
+    this->update_handle_ = 0;
+    if (err == ESP_ERR_INVALID_SIZE) {
+      return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
+    } else if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
+      return OTA_RESPONSE_ERROR_WRITING_FLASH;
+    }
+    return OTA_RESPONSE_ERROR_UNKNOWN;
+  }
+  return OTA_RESPONSE_OK;
+}
+
+void IDFOTABackend::set_update_md5(const char *md5) {
+  // pass
+}
+
+OTAResponseTypes IDFOTABackend::write(uint8_t *data, size_t len) {
+  esp_err_t err = esp_ota_write(this->update_handle_, data, len);
+  if (err != ESP_OK) {
+    if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
+      return OTA_RESPONSE_ERROR_MAGIC;
+    } else if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
+      return OTA_RESPONSE_ERROR_WRITING_FLASH;
+    }
+    return OTA_RESPONSE_ERROR_UNKNOWN;
+  }
+  return OTA_RESPONSE_OK;
+}
+
+OTAResponseTypes IDFOTABackend::end() {
+  esp_err_t err = esp_ota_end(this->update_handle_);
+  this->update_handle_ = 0;
+  if (err == ESP_OK) {
+    err = esp_ota_set_boot_partition(this->partition_);
+    if (err == ESP_OK) {
+      return OTA_RESPONSE_OK;
+    }
+  }
+  if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
+    return OTA_RESPONSE_ERROR_UPDATE_END;
+  }
+  if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
+    return OTA_RESPONSE_ERROR_WRITING_FLASH;
+  }
+  return OTA_RESPONSE_ERROR_UNKNOWN;
+}
+
+void IDFOTABackend::abort() {
+  esp_ota_abort(this->update_handle_);
+  this->update_handle_ = 0;
+}
+
+}  // namespace ota
+}  // namespace esphome
+#endif

esphome/components/ota/ota_backend_esp_idf.h

@@ -0,0 +1,27 @@
+#pragma once
+#include "esphome/core/defines.h"
+#ifdef USE_ESP_IDF
+
+#include "ota_component.h"
+#include "ota_backend.h"
+#include <esp_ota_ops.h>
+
+namespace esphome {
+namespace ota {
+
+class IDFOTABackend : public OTABackend {
+ public:
+  OTAResponseTypes begin(size_t image_size) override;
+  void set_update_md5(const char *md5) override;
+  OTAResponseTypes write(uint8_t *data, size_t len) override;
+  OTAResponseTypes end() override;
+  void abort() override;
+
+ private:
+  esp_ota_handle_t update_handle_{0};
+  const esp_partition_t *partition_;
+};
+
+}  // namespace ota
+}  // namespace esphome
+#endif

esphome/components/ota/ota_component.cpp

@@ -1,4 +1,8 @@
 #include "ota_component.h"
+#include "ota_backend.h"
+#include "ota_backend_arduino_esp32.h"
+#include "ota_backend_arduino_esp8266.h"
+#include "ota_backend_esp_idf.h"
 
 #include "esphome/core/log.h"
 #include "esphome/core/application.h"
@@ -9,23 +13,8 @@
 #include <cerrno>
 #include <cstdio>
 
-#ifdef USE_ARDUINO
 #ifdef USE_OTA_PASSWORD
 #include <MD5Builder.h>
-#endif  // USE_OTA_PASSWORD
-
-#ifdef USE_ESP32
-#include <Update.h>
-#endif  // USE_ESP32
-#endif  // USE_ARDUINO
-
-#ifdef USE_ESP8266
-#include <Updater.h>
-#include "esphome/components/esp8266/preferences.h"
-#endif  // USE_ESP8266
-
-#ifdef USE_ESP_IDF
-#include <esp_ota_ops.h>
 #endif
 
 namespace esphome {
@@ -35,125 +24,19 @@ static const char *const TAG = "ota";
 
 static const uint8_t OTA_VERSION_1_0 = 1;
 
-class OTABackend {
+std::unique_ptr<OTABackend> make_ota_backend() {
- public:
-  virtual ~OTABackend() = default;
-  virtual OTAResponseTypes begin(size_t image_size) = 0;
-  virtual void set_update_md5(const char *md5) = 0;
-  virtual OTAResponseTypes write(uint8_t *data, size_t len) = 0;
-  virtual OTAResponseTypes end() = 0;
-  virtual void abort() = 0;
-};
-
 #ifdef USE_ARDUINO
-class ArduinoOTABackend : public OTABackend {
- public:
-  OTAResponseTypes begin(size_t image_size) override {
-    bool ret = Update.begin(image_size, U_FLASH);
-    if (ret) {
 #ifdef USE_ESP8266
-      esp8266::preferences_prevent_write(true);
+  return make_unique<ArduinoESP8266OTABackend>();
-#endif
+#endif  // USE_ESP8266
-      return OTA_RESPONSE_OK;
-    }
-
-    uint8_t error = Update.getError();
-#ifdef USE_ESP8266
-    if (error == UPDATE_ERROR_BOOTSTRAP)
-      return OTA_RESPONSE_ERROR_INVALID_BOOTSTRAPPING;
-    if (error == UPDATE_ERROR_NEW_FLASH_CONFIG)
-      return OTA_RESPONSE_ERROR_WRONG_NEW_FLASH_CONFIG;
-    if (error == UPDATE_ERROR_FLASH_CONFIG)
-      return OTA_RESPONSE_ERROR_WRONG_CURRENT_FLASH_CONFIG;
-    if (error == UPDATE_ERROR_SPACE)
-      return OTA_RESPONSE_ERROR_ESP8266_NOT_ENOUGH_SPACE;
-#endif
 #ifdef USE_ESP32
-    if (error == UPDATE_ERROR_SIZE)
+  return make_unique<ArduinoESP32OTABackend>();
-      return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
+#endif  // USE_ESP32
-#endif
-    return OTA_RESPONSE_ERROR_UNKNOWN;
-  }
-  void set_update_md5(const char *md5) override { Update.setMD5(md5); }
-  OTAResponseTypes write(uint8_t *data, size_t len) override {
-    size_t written = Update.write(data, len);
-    if (written != len) {
-      return OTA_RESPONSE_ERROR_WRITING_FLASH;
-    }
-    return OTA_RESPONSE_OK;
-  }
-  OTAResponseTypes end() override {
-    if (!Update.end())
-      return OTA_RESPONSE_ERROR_UPDATE_END;
-    return OTA_RESPONSE_OK;
-  }
-  void abort() override {
-#ifdef USE_ESP32
-    Update.abort();
-#endif
-
-#ifdef USE_ESP8266
-    Update.end();
-    esp8266::preferences_prevent_write(false);
-#endif
-  }
-};
-std::unique_ptr<OTABackend> make_ota_backend() { return make_unique<ArduinoOTABackend>(); }
 #endif  // USE_ARDUINO
-
 #ifdef USE_ESP_IDF
-class IDFOTABackend : public OTABackend {
+  return make_unique<IDFOTABackend>();
- public:
-  esp_ota_handle_t update_handle = 0;
-
-  OTAResponseTypes begin(size_t image_size) override {
-    const esp_partition_t *update_partition = esp_ota_get_next_update_partition(nullptr);
-    if (update_partition == nullptr) {
-      return OTA_RESPONSE_ERROR_NO_UPDATE_PARTITION;
-    }
-    esp_err_t err = esp_ota_begin(update_partition, image_size, &update_handle);
-    if (err != ESP_OK) {
-      esp_ota_abort(update_handle);
-      update_handle = 0;
-      if (err == ESP_ERR_INVALID_SIZE) {
-        return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
-      } else if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
-        return OTA_RESPONSE_ERROR_WRITING_FLASH;
-      }
-      return OTA_RESPONSE_ERROR_UNKNOWN;
-    }
-    return OTA_RESPONSE_OK;
-  }
-  void set_update_md5(const char *md5) override {
-    // pass
-  }
-  OTAResponseTypes write(uint8_t *data, size_t len) override {
-    esp_err_t err = esp_ota_write(update_handle, data, len);
-    if (err != ESP_OK) {
-      if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
-        return OTA_RESPONSE_ERROR_MAGIC;
-      } else if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
-        return OTA_RESPONSE_ERROR_WRITING_FLASH;
-      }
-      return OTA_RESPONSE_ERROR_UNKNOWN;
-    }
-    return OTA_RESPONSE_OK;
-  }
-  OTAResponseTypes end() override {
-    esp_err_t err = esp_ota_end(update_handle);
-    update_handle = 0;
-    if (err != ESP_OK) {
-      if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
-        return OTA_RESPONSE_ERROR_UPDATE_END;
-      }
-      return OTA_RESPONSE_ERROR_UNKNOWN;
-    }
-    return OTA_RESPONSE_OK;
-  }
-  void abort() override { esp_ota_abort(update_handle); }
-};
-std::unique_ptr<OTABackend> make_ota_backend() { return make_unique<IDFOTABackend>(); }
 #endif  // USE_ESP_IDF
+}
 
 void OTAComponent::setup() {
   server_ = socket::socket(AF_INET, SOCK_STREAM, 0);