OTA firmware MD5 check + password support for esp-idf (#2507)

Co-authored-by: Maurice Makaay <account-github@makaay.nl>

CODEOWNERS

@@ -89,6 +89,7 @@ esphome/components/mcp23x17_base/* @jesserockz
 esphome/components/mcp23xxx_base/* @jesserockz
 esphome/components/mcp2515/* @danielschramm @mvturnho
 esphome/components/mcp9808/* @k7hpn
+esphome/components/md5/* @esphome/core
 esphome/components/mdns/* @esphome/core
 esphome/components/midea/* @dudanov
 esphome/components/mitsubishi/* @RubyBailey

esphome/components/md5/__init__.py

@@ -0,0 +1 @@
+CODEOWNERS = ["@esphome/core"]

esphome/components/md5/md5.cpp

@@ -0,0 +1,51 @@
+#include <cstdio>
+#include <cstring>
+#include "md5.h"
+#include "esphome/core/helpers.h"
+
+namespace esphome {
+namespace md5 {
+
+void MD5Digest::init() {
+  memset(this->digest_, 0, 16);
+  MD5Init(&this->ctx_);
+}
+
+void MD5Digest::add(const uint8_t *data, size_t len) { MD5Update(&this->ctx_, data, len); }
+
+void MD5Digest::calculate() { MD5Final(this->digest_, &this->ctx_); }
+
+void MD5Digest::get_bytes(uint8_t *output) { memcpy(output, this->digest_, 16); }
+
+void MD5Digest::get_hex(char *output) {
+  for (size_t i = 0; i < 16; i++) {
+    sprintf(output + i * 2, "%02x", this->digest_[i]);
+  }
+}
+
+bool MD5Digest::equals_bytes(const char *expected) {
+  for (size_t i = 0; i < 16; i++) {
+    if (expected[i] != this->digest_[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool MD5Digest::equals_hex(const char *expected) {
+  for (size_t i = 0; i < 16; i++) {
+    auto high = parse_hex(expected[i * 2]);
+    auto low = parse_hex(expected[i * 2 + 1]);
+    if (!high.has_value() || !low.has_value()) {
+      return false;
+    }
+    auto value = (*high << 4) | *low;
+    if (value != this->digest_[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+}  // namespace md5
+}  // namespace esphome

esphome/components/md5/md5.h

@@ -0,0 +1,58 @@
+#pragma once
+
+#include "esphome/core/defines.h"
+
+#ifdef USE_ESP_IDF
+#include "esp32/rom/md5_hash.h"
+#define MD5_CTX_TYPE MD5Context
+#endif
+
+#if defined(USE_ARDUINO) && defined(USE_ESP32)
+#include "rom/md5_hash.h"
+#define MD5_CTX_TYPE MD5Context
+#endif
+
+#if defined(USE_ARDUINO) && defined(USE_ESP8266)
+#include <md5.h>
+#define MD5_CTX_TYPE md5_context_t
+#endif
+
+namespace esphome {
+namespace md5 {
+
+class MD5Digest {
+ public:
+  MD5Digest() = default;
+  ~MD5Digest() = default;
+
+  /// Initialize a new MD5 digest computation.
+  void init();
+
+  /// Add bytes of data for the digest.
+  void add(const uint8_t *data, size_t len);
+  void add(const char *data, size_t len) { this->add((const uint8_t *) data, len); }
+
+  /// Compute the digest, based on the provided data.
+  void calculate();
+
+  /// Retrieve the MD5 digest as bytes.
+  /// The output must be able to hold 16 bytes or more.
+  void get_bytes(uint8_t *output);
+
+  /// Retrieve the MD5 digest as hex characters.
+  /// The output must be able to hold 32 bytes or more.
+  void get_hex(char *output);
+
+  /// Compare the digest against a provided byte-encoded digest (16 bytes).
+  bool equals_bytes(const char *expected);
+
+  /// Compare the digest against a provided hex-encoded digest (32 bytes).
+  bool equals_hex(const char *expected);
+
+ protected:
+  MD5_CTX_TYPE ctx_{};
+  uint8_t digest_[16];
+};
+
+}  // namespace md5
+}  // namespace esphome

esphome/components/ota/__init__.py

@@ -15,7 +15,7 @@ from esphome.core import CORE, coroutine_with_priority
 
 CODEOWNERS = ["@esphome/core"]
 DEPENDENCIES = ["network"]
-AUTO_LOAD = ["socket"]
+AUTO_LOAD = ["socket", "md5"]
 
 CONF_ON_STATE_CHANGE = "on_state_change"
 CONF_ON_BEGIN = "on_begin"
@@ -35,20 +35,12 @@ OTAEndTrigger = ota_ns.class_("OTAEndTrigger", automation.Trigger.template())
 OTAErrorTrigger = ota_ns.class_("OTAErrorTrigger", automation.Trigger.template())
 
 
-def validate_password_support(value):
-    if CORE.using_arduino:
-        return value
-    if CORE.using_esp_idf:
-        raise cv.Invalid("Password support is not implemented yet for ESP-IDF")
-    raise NotImplementedError
-
-
 CONFIG_SCHEMA = cv.Schema(
     {
         cv.GenerateID(): cv.declare_id(OTAComponent),
         cv.Optional(CONF_SAFE_MODE, default=True): cv.boolean,
         cv.SplitDefault(CONF_PORT, esp8266=8266, esp32=3232): cv.port,
-        cv.Optional(CONF_PASSWORD): cv.All(cv.string, validate_password_support),
+        cv.Optional(CONF_PASSWORD): cv.string,
         cv.Optional(
             CONF_REBOOT_TIMEOUT, default="5min"
         ): cv.positive_time_period_milliseconds,

esphome/components/ota/ota_backend_arduino_esp32.h

@@ -9,6 +9,7 @@ namespace esphome {
 namespace ota {
 
 class ArduinoESP32OTABackend : 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;

esphome/components/ota/ota_backend_esp_idf.cpp

@@ -4,6 +4,7 @@
 #include "ota_backend_esp_idf.h"
 #include "ota_component.h"
 #include <esp_ota_ops.h>
+#include "esphome/components/md5/md5.h"
 
 namespace esphome {
 namespace ota {
@@ -24,15 +25,15 @@ OTAResponseTypes IDFOTABackend::begin(size_t image_size) {
     }
     return OTA_RESPONSE_ERROR_UNKNOWN;
   }
+  this->md5_.init();
   return OTA_RESPONSE_OK;
 }
 
-void IDFOTABackend::set_update_md5(const char *md5) {
-  // pass
-}
+void IDFOTABackend::set_update_md5(const char *expected_md5) { memcpy(this->expected_bin_md5_, expected_md5, 32); }
 
 OTAResponseTypes IDFOTABackend::write(uint8_t *data, size_t len) {
   esp_err_t err = esp_ota_write(this->update_handle_, data, len);
+  this->md5_.add(data, len);
   if (err != ESP_OK) {
     if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
       return OTA_RESPONSE_ERROR_MAGIC;
@@ -45,6 +46,11 @@ OTAResponseTypes IDFOTABackend::write(uint8_t *data, size_t len) {
 }
 
 OTAResponseTypes IDFOTABackend::end() {
+  this->md5_.calculate();
+  if (!this->md5_.equals_hex(this->expected_bin_md5_)) {
+    this->abort();
+    return OTA_RESPONSE_ERROR_UPDATE_END;
+  }
   esp_err_t err = esp_ota_end(this->update_handle_);
   this->update_handle_ = 0;
   if (err == ESP_OK) {

esphome/components/ota/ota_backend_esp_idf.h

@@ -5,6 +5,7 @@
 #include "ota_component.h"
 #include "ota_backend.h"
 #include <esp_ota_ops.h>
+#include "esphome/components/md5/md5.h"
 
 namespace esphome {
 namespace ota {
@@ -20,6 +21,8 @@ class IDFOTABackend : public OTABackend {
  private:
   esp_ota_handle_t update_handle_{0};
   const esp_partition_t *partition_;
+  md5::MD5Digest md5_{};
+  char expected_bin_md5_[32];
 };
 
 }  // namespace ota

esphome/components/ota/ota_component.cpp

@@ -8,15 +8,12 @@
 #include "esphome/core/application.h"
 #include "esphome/core/hal.h"
 #include "esphome/core/util.h"
+#include "esphome/components/md5/md5.h"
 #include "esphome/components/network/util.h"
 
 #include <cerrno>
 #include <cstdio>
 
-#ifdef USE_OTA_PASSWORD
-#include <MD5Builder.h>
-#endif
-
 namespace esphome {
 namespace ota {
 
@@ -173,12 +170,12 @@ void OTAComponent::handle_() {
   if (!this->password_.empty()) {
     buf[0] = OTA_RESPONSE_REQUEST_AUTH;
     this->writeall_(buf, 1);
-    MD5Builder md5_builder{};
-    md5_builder.begin();
+    md5::MD5Digest md5{};
+    md5.init();
     sprintf(sbuf, "%08X", random_uint32());
-    md5_builder.add(sbuf);
-    md5_builder.calculate();
-    md5_builder.getChars(sbuf);
+    md5.add(sbuf, 8);
+    md5.calculate();
+    md5.get_hex(sbuf);
     ESP_LOGV(TAG, "Auth: Nonce is %s", sbuf);
 
     // Send nonce, 32 bytes hex MD5
@@ -188,10 +185,10 @@ void OTAComponent::handle_() {
     }
 
     // prepare challenge
-    md5_builder.begin();
-    md5_builder.add(this->password_.c_str());
+    md5.init();
+    md5.add(this->password_.c_str(), this->password_.length());
     // add nonce
-    md5_builder.add(sbuf);
+    md5.add(sbuf, 32);
 
     // Receive cnonce, 32 bytes hex MD5
     if (!this->readall_(buf, 32)) {
@@ -201,11 +198,11 @@ void OTAComponent::handle_() {
     sbuf[32] = '\0';
     ESP_LOGV(TAG, "Auth: CNonce is %s", sbuf);
     // add cnonce
-    md5_builder.add(sbuf);
+    md5.add(sbuf, 32);
 
     // calculate result
-    md5_builder.calculate();
-    md5_builder.getChars(sbuf);
+    md5.calculate();
+    md5.get_hex(sbuf);
     ESP_LOGV(TAG, "Auth: Result is %s", sbuf);
 
     // Receive result, 32 bytes hex MD5

esphome/core/defines.h

@@ -26,6 +26,7 @@
 #define USE_LOGGER
 #define USE_MDNS
 #define USE_NUMBER
+#define USE_OTA_PASSWORD
 #define USE_OTA_STATE_CALLBACK
 #define USE_POWER_SUPPLY
 #define USE_PROMETHEUS