OTA 2 which confirm each written chunk (#6066)

Co-authored-by: Jesse Hills <3060199+jesserockz@users.noreply.github.com>

esphome/components/ota/__init__.py

@@ -12,6 +12,7 @@ from esphome.const import (
     CONF_TRIGGER_ID,
     CONF_OTA,
     KEY_PAST_SAFE_MODE,
+    CONF_VERSION,
 )
 from esphome.core import CORE, coroutine_with_priority
 
@@ -41,6 +42,7 @@ CONFIG_SCHEMA = cv.Schema(
     {
         cv.GenerateID(): cv.declare_id(OTAComponent),
         cv.Optional(CONF_SAFE_MODE, default=True): cv.boolean,
+        cv.Optional(CONF_VERSION, default=2): cv.one_of(1, 2, int=True),
         cv.SplitDefault(
             CONF_PORT,
             esp8266=8266,
@@ -93,6 +95,7 @@ async def to_code(config):
     if CONF_PASSWORD in config:
         cg.add(var.set_auth_password(config[CONF_PASSWORD]))
         cg.add_define("USE_OTA_PASSWORD")
+    cg.add_define("USE_OTA_VERSION", config[CONF_VERSION])
 
     await cg.register_component(var, config)
 

esphome/components/ota/ota_component.cpp

@@ -20,8 +20,7 @@ namespace esphome {
 namespace ota {
 
 static const char *const TAG = "ota";
-
-static const uint8_t OTA_VERSION_1_0 = 1;
+static constexpr u_int16_t OTA_BLOCK_SIZE = 8192;
 
 OTAComponent *global_ota_component = nullptr;  // NOLINT(cppcoreguidelines-avoid-non-const-global-variables)
 
@@ -101,6 +100,7 @@ void OTAComponent::dump_config() {
     ESP_LOGCONFIG(TAG, "  Using Password.");
   }
 #endif
+  ESP_LOGCONFIG(TAG, "  OTA version: %d.", USE_OTA_VERSION);
   if (this->has_safe_mode_ && this->safe_mode_rtc_value_ > 1 &&
       this->safe_mode_rtc_value_ != esphome::ota::OTAComponent::ENTER_SAFE_MODE_MAGIC) {
     ESP_LOGW(TAG, "Last Boot was an unhandled reset, will proceed to safe mode in %" PRIu32 " restarts",
@@ -132,6 +132,9 @@ void OTAComponent::handle_() {
   uint8_t ota_features;
   std::unique_ptr<OTABackend> backend;
   (void) ota_features;
+#if USE_OTA_VERSION == 2
+  size_t size_acknowledged = 0;
+#endif
 
   if (client_ == nullptr) {
     struct sockaddr_storage source_addr;
@@ -168,7 +171,7 @@ void OTAComponent::handle_() {
 
   // Send OK and version - 2 bytes
   buf[0] = OTA_RESPONSE_OK;
-  buf[1] = OTA_VERSION_1_0;
+  buf[1] = USE_OTA_VERSION;
   this->writeall_(buf, 2);
 
   backend = make_ota_backend();
@@ -312,6 +315,13 @@ void OTAComponent::handle_() {
       goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
     }
     total += read;
+#if USE_OTA_VERSION == 2
+    while (size_acknowledged + OTA_BLOCK_SIZE <= total || (total == ota_size && size_acknowledged < ota_size)) {
+      buf[0] = OTA_RESPONSE_CHUNK_OK;
+      this->writeall_(buf, 1);
+      size_acknowledged += OTA_BLOCK_SIZE;
+    }
+#endif
 
     uint32_t now = millis();
     if (now - last_progress > 1000) {

esphome/components/ota/ota_component.h

@@ -10,31 +10,32 @@ namespace esphome {
 namespace ota {
 
 enum OTAResponseTypes {
-  OTA_RESPONSE_OK = 0,
-  OTA_RESPONSE_REQUEST_AUTH = 1,
+  OTA_RESPONSE_OK = 0x00,
+  OTA_RESPONSE_REQUEST_AUTH = 0x01,
 
-  OTA_RESPONSE_HEADER_OK = 64,
-  OTA_RESPONSE_AUTH_OK = 65,
-  OTA_RESPONSE_UPDATE_PREPARE_OK = 66,
-  OTA_RESPONSE_BIN_MD5_OK = 67,
-  OTA_RESPONSE_RECEIVE_OK = 68,
-  OTA_RESPONSE_UPDATE_END_OK = 69,
-  OTA_RESPONSE_SUPPORTS_COMPRESSION = 70,
+  OTA_RESPONSE_HEADER_OK = 0x40,
+  OTA_RESPONSE_AUTH_OK = 0x41,
+  OTA_RESPONSE_UPDATE_PREPARE_OK = 0x42,
+  OTA_RESPONSE_BIN_MD5_OK = 0x43,
+  OTA_RESPONSE_RECEIVE_OK = 0x44,
+  OTA_RESPONSE_UPDATE_END_OK = 0x45,
+  OTA_RESPONSE_SUPPORTS_COMPRESSION = 0x46,
+  OTA_RESPONSE_CHUNK_OK = 0x47,
 
-  OTA_RESPONSE_ERROR_MAGIC = 128,
-  OTA_RESPONSE_ERROR_UPDATE_PREPARE = 129,
-  OTA_RESPONSE_ERROR_AUTH_INVALID = 130,
-  OTA_RESPONSE_ERROR_WRITING_FLASH = 131,
-  OTA_RESPONSE_ERROR_UPDATE_END = 132,
-  OTA_RESPONSE_ERROR_INVALID_BOOTSTRAPPING = 133,
-  OTA_RESPONSE_ERROR_WRONG_CURRENT_FLASH_CONFIG = 134,
-  OTA_RESPONSE_ERROR_WRONG_NEW_FLASH_CONFIG = 135,
-  OTA_RESPONSE_ERROR_ESP8266_NOT_ENOUGH_SPACE = 136,
-  OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE = 137,
-  OTA_RESPONSE_ERROR_NO_UPDATE_PARTITION = 138,
-  OTA_RESPONSE_ERROR_MD5_MISMATCH = 139,
-  OTA_RESPONSE_ERROR_RP2040_NOT_ENOUGH_SPACE = 140,
-  OTA_RESPONSE_ERROR_UNKNOWN = 255,
+  OTA_RESPONSE_ERROR_MAGIC = 0x80,
+  OTA_RESPONSE_ERROR_UPDATE_PREPARE = 0x81,
+  OTA_RESPONSE_ERROR_AUTH_INVALID = 0x82,
+  OTA_RESPONSE_ERROR_WRITING_FLASH = 0x83,
+  OTA_RESPONSE_ERROR_UPDATE_END = 0x84,
+  OTA_RESPONSE_ERROR_INVALID_BOOTSTRAPPING = 0x85,
+  OTA_RESPONSE_ERROR_WRONG_CURRENT_FLASH_CONFIG = 0x86,
+  OTA_RESPONSE_ERROR_WRONG_NEW_FLASH_CONFIG = 0x87,
+  OTA_RESPONSE_ERROR_ESP8266_NOT_ENOUGH_SPACE = 0x88,
+  OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE = 0x89,
+  OTA_RESPONSE_ERROR_NO_UPDATE_PARTITION = 0x8A,
+  OTA_RESPONSE_ERROR_MD5_MISMATCH = 0x8B,
+  OTA_RESPONSE_ERROR_RP2040_NOT_ENOUGH_SPACE = 0x8C,
+  OTA_RESPONSE_ERROR_UNKNOWN = 0xFF,
 };
 
 enum OTAState { OTA_COMPLETED = 0, OTA_STARTED, OTA_IN_PROGRESS, OTA_ERROR };

esphome/core/defines.h

@@ -37,6 +37,7 @@
 #define USE_OTA
 #define USE_OTA_PASSWORD
 #define USE_OTA_STATE_CALLBACK
+#define USE_OTA_VERSION 1
 #define USE_OUTPUT
 #define USE_POWER_SUPPLY
 #define USE_QR_CODE

esphome/espota2.py

@@ -12,32 +12,34 @@ import time
 from esphome.core import EsphomeError
 from esphome.helpers import is_ip_address, resolve_ip_address
 
-RESPONSE_OK = 0
-RESPONSE_REQUEST_AUTH = 1
+RESPONSE_OK = 0x00
+RESPONSE_REQUEST_AUTH = 0x01
 
-RESPONSE_HEADER_OK = 64
-RESPONSE_AUTH_OK = 65
-RESPONSE_UPDATE_PREPARE_OK = 66
-RESPONSE_BIN_MD5_OK = 67
-RESPONSE_RECEIVE_OK = 68
-RESPONSE_UPDATE_END_OK = 69
-RESPONSE_SUPPORTS_COMPRESSION = 70
+RESPONSE_HEADER_OK = 0x40
+RESPONSE_AUTH_OK = 0x41
+RESPONSE_UPDATE_PREPARE_OK = 0x42
+RESPONSE_BIN_MD5_OK = 0x43
+RESPONSE_RECEIVE_OK = 0x44
+RESPONSE_UPDATE_END_OK = 0x45
+RESPONSE_SUPPORTS_COMPRESSION = 0x46
+RESPONSE_CHUNK_OK = 0x47
 
-RESPONSE_ERROR_MAGIC = 128
-RESPONSE_ERROR_UPDATE_PREPARE = 129
-RESPONSE_ERROR_AUTH_INVALID = 130
-RESPONSE_ERROR_WRITING_FLASH = 131
-RESPONSE_ERROR_UPDATE_END = 132
-RESPONSE_ERROR_INVALID_BOOTSTRAPPING = 133
-RESPONSE_ERROR_WRONG_CURRENT_FLASH_CONFIG = 134
-RESPONSE_ERROR_WRONG_NEW_FLASH_CONFIG = 135
-RESPONSE_ERROR_ESP8266_NOT_ENOUGH_SPACE = 136
-RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE = 137
-RESPONSE_ERROR_NO_UPDATE_PARTITION = 138
-RESPONSE_ERROR_MD5_MISMATCH = 139
-RESPONSE_ERROR_UNKNOWN = 255
+RESPONSE_ERROR_MAGIC = 0x80
+RESPONSE_ERROR_UPDATE_PREPARE = 0x81
+RESPONSE_ERROR_AUTH_INVALID = 0x82
+RESPONSE_ERROR_WRITING_FLASH = 0x83
+RESPONSE_ERROR_UPDATE_END = 0x84
+RESPONSE_ERROR_INVALID_BOOTSTRAPPING = 0x85
+RESPONSE_ERROR_WRONG_CURRENT_FLASH_CONFIG = 0x86
+RESPONSE_ERROR_WRONG_NEW_FLASH_CONFIG = 0x87
+RESPONSE_ERROR_ESP8266_NOT_ENOUGH_SPACE = 0x88
+RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE = 0x89
+RESPONSE_ERROR_NO_UPDATE_PARTITION = 0x8A
+RESPONSE_ERROR_MD5_MISMATCH = 0x8B
+RESPONSE_ERROR_UNKNOWN = 0xFF
 
 OTA_VERSION_1_0 = 1
+OTA_VERSION_2_0 = 2
 
 MAGIC_BYTES = [0x6C, 0x26, 0xF7, 0x5C, 0x45]
 
@@ -203,7 +205,8 @@ def perform_ota(
     send_check(sock, MAGIC_BYTES, "magic bytes")
 
     _, version = receive_exactly(sock, 2, "version", RESPONSE_OK)
-    if version != OTA_VERSION_1_0:
+    _LOGGER.debug("Device support OTA version: %s", version)
+    if version not in (OTA_VERSION_1_0, OTA_VERSION_2_0):
         raise OTAError(f"Unsupported OTA version {version}")
 
     # Features
@@ -279,6 +282,8 @@ def perform_ota(
 
         try:
             sock.sendall(chunk)
+            if version >= OTA_VERSION_2_0:
+                receive_exactly(sock, 1, "chunk OK", RESPONSE_CHUNK_OK)
         except OSError as err:
             sys.stderr.write("\n")
             raise OTAError(f"Error sending data: {err}") from err

tests/test3.1.yaml

@@ -49,6 +49,7 @@ spi:
     number: GPIO14
 
 ota:
+  version: 2
 
 logger: