From 5be52f71f9c036d78fa1fe6799650899a6180794 Mon Sep 17 00:00:00 2001 From: Otto Winter Date: Fri, 22 Oct 2021 13:02:55 +0200 Subject: [PATCH] Add OTA upload compression for ESP8266 (#2601) --- esphome/components/ota/ota_backend.h | 1 + .../ota/ota_backend_arduino_esp32.h | 1 + .../ota/ota_backend_arduino_esp8266.h | 1 + esphome/components/ota/ota_backend_esp_idf.h | 1 + esphome/components/ota/ota_component.cpp | 9 +++- esphome/components/ota/ota_component.h | 1 + esphome/espota2.py | 44 ++++++++++++------- 7 files changed, 42 insertions(+), 16 deletions(-) diff --git a/esphome/components/ota/ota_backend.h b/esphome/components/ota/ota_backend.h index c253e009c6..5c5b61a278 100644 --- a/esphome/components/ota/ota_backend.h +++ b/esphome/components/ota/ota_backend.h @@ -12,6 +12,7 @@ class OTABackend { virtual OTAResponseTypes write(uint8_t *data, size_t len) = 0; virtual OTAResponseTypes end() = 0; virtual void abort() = 0; + virtual bool supports_compression() = 0; }; } // namespace ota diff --git a/esphome/components/ota/ota_backend_arduino_esp32.h b/esphome/components/ota/ota_backend_arduino_esp32.h index 6b712502fb..f86a70d678 100644 --- a/esphome/components/ota/ota_backend_arduino_esp32.h +++ b/esphome/components/ota/ota_backend_arduino_esp32.h @@ -15,6 +15,7 @@ class ArduinoESP32OTABackend : public OTABackend { OTAResponseTypes write(uint8_t *data, size_t len) override; OTAResponseTypes end() override; void abort() override; + bool supports_compression() override { return false; } }; } // namespace ota diff --git a/esphome/components/ota/ota_backend_arduino_esp8266.h b/esphome/components/ota/ota_backend_arduino_esp8266.h index d1195af911..cf29a90fc1 100644 --- a/esphome/components/ota/ota_backend_arduino_esp8266.h +++ b/esphome/components/ota/ota_backend_arduino_esp8266.h @@ -16,6 +16,7 @@ class ArduinoESP8266OTABackend : public OTABackend { OTAResponseTypes write(uint8_t *data, size_t len) override; OTAResponseTypes end() override; void abort() override; + bool supports_compression() override { return true; } }; } // namespace ota diff --git a/esphome/components/ota/ota_backend_esp_idf.h b/esphome/components/ota/ota_backend_esp_idf.h index 49c6e124fa..af09d0d693 100644 --- a/esphome/components/ota/ota_backend_esp_idf.h +++ b/esphome/components/ota/ota_backend_esp_idf.h @@ -17,6 +17,7 @@ class IDFOTABackend : public OTABackend { OTAResponseTypes write(uint8_t *data, size_t len) override; OTAResponseTypes end() override; void abort() override; + bool supports_compression() override { return false; } private: esp_ota_handle_t update_handle_{0}; diff --git a/esphome/components/ota/ota_component.cpp b/esphome/components/ota/ota_component.cpp index 6d51087882..e49c108320 100644 --- a/esphome/components/ota/ota_component.cpp +++ b/esphome/components/ota/ota_component.cpp @@ -104,6 +104,8 @@ void OTAComponent::loop() { } } +static const uint8_t FEATURE_SUPPORTS_COMPRESSION = 0x01; + void OTAComponent::handle_() { OTAResponseTypes error_code = OTA_RESPONSE_ERROR_UNKNOWN; bool update_started = false; @@ -154,6 +156,8 @@ void OTAComponent::handle_() { buf[1] = OTA_VERSION_1_0; this->writeall_(buf, 2); + backend = make_ota_backend(); + // Read features - 1 byte if (!this->readall_(buf, 1)) { ESP_LOGW(TAG, "Reading features failed!"); @@ -164,6 +168,10 @@ void OTAComponent::handle_() { // Acknowledge header - 1 byte buf[0] = OTA_RESPONSE_HEADER_OK; + if ((ota_features & FEATURE_SUPPORTS_COMPRESSION) != 0 && backend->supports_compression()) { + buf[0] = OTA_RESPONSE_SUPPORTS_COMPRESSION; + } + this->writeall_(buf, 1); #ifdef USE_OTA_PASSWORD @@ -241,7 +249,6 @@ void OTAComponent::handle_() { } ESP_LOGV(TAG, "OTA size is %u bytes", ota_size); - backend = make_ota_backend(); error_code = backend->begin(ota_size); if (error_code != OTA_RESPONSE_OK) goto error; diff --git a/esphome/components/ota/ota_component.h b/esphome/components/ota/ota_component.h index e08e187df6..5647d52eeb 100644 --- a/esphome/components/ota/ota_component.h +++ b/esphome/components/ota/ota_component.h @@ -19,6 +19,7 @@ enum OTAResponseTypes { OTA_RESPONSE_BIN_MD5_OK = 67, OTA_RESPONSE_RECEIVE_OK = 68, OTA_RESPONSE_UPDATE_END_OK = 69, + OTA_RESPONSE_SUPPORTS_COMPRESSION = 70, OTA_RESPONSE_ERROR_MAGIC = 128, OTA_RESPONSE_ERROR_UPDATE_PREPARE = 129, diff --git a/esphome/espota2.py b/esphome/espota2.py index f8a2fab94c..8f299395dd 100644 --- a/esphome/espota2.py +++ b/esphome/espota2.py @@ -4,6 +4,7 @@ import random import socket import sys import time +import gzip from esphome.core import EsphomeError from esphome.helpers import is_ip_address, resolve_ip_address @@ -17,6 +18,7 @@ RESPONSE_UPDATE_PREPARE_OK = 66 RESPONSE_BIN_MD5_OK = 67 RESPONSE_RECEIVE_OK = 68 RESPONSE_UPDATE_END_OK = 69 +RESPONSE_SUPPORTS_COMPRESSION = 70 RESPONSE_ERROR_MAGIC = 128 RESPONSE_ERROR_UPDATE_PREPARE = 129 @@ -34,6 +36,8 @@ OTA_VERSION_1_0 = 1 MAGIC_BYTES = [0x6C, 0x26, 0xF7, 0x5C, 0x45] +FEATURE_SUPPORTS_COMPRESSION = 0x01 + _LOGGER = logging.getLogger(__name__) @@ -170,11 +174,9 @@ def send_check(sock, data, msg): def perform_ota(sock, password, file_handle, filename): - file_md5 = hashlib.md5(file_handle.read()).hexdigest() - file_size = file_handle.tell() + file_contents = file_handle.read() + file_size = len(file_contents) _LOGGER.info("Uploading %s (%s bytes)", filename, file_size) - file_handle.seek(0) - _LOGGER.debug("MD5 of binary is %s", file_md5) # Enable nodelay, we need it for phase 1 sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) @@ -185,8 +187,16 @@ def perform_ota(sock, password, file_handle, filename): raise OTAError(f"Unsupported OTA version {version}") # Features - send_check(sock, 0x00, "features") - receive_exactly(sock, 1, "features", RESPONSE_HEADER_OK) + send_check(sock, FEATURE_SUPPORTS_COMPRESSION, "features") + features = receive_exactly( + sock, 1, "features", [RESPONSE_HEADER_OK, RESPONSE_SUPPORTS_COMPRESSION] + )[0] + + if features == RESPONSE_SUPPORTS_COMPRESSION: + upload_contents = gzip.compress(file_contents, compresslevel=9) + _LOGGER.info("Compressed to %s bytes", len(upload_contents)) + else: + upload_contents = file_contents (auth,) = receive_exactly( sock, 1, "auth", [RESPONSE_REQUEST_AUTH, RESPONSE_AUTH_OK] @@ -213,16 +223,20 @@ def perform_ota(sock, password, file_handle, filename): send_check(sock, result, "auth result") receive_exactly(sock, 1, "auth result", RESPONSE_AUTH_OK) - file_size_encoded = [ - (file_size >> 24) & 0xFF, - (file_size >> 16) & 0xFF, - (file_size >> 8) & 0xFF, - (file_size >> 0) & 0xFF, + upload_size = len(upload_contents) + upload_size_encoded = [ + (upload_size >> 24) & 0xFF, + (upload_size >> 16) & 0xFF, + (upload_size >> 8) & 0xFF, + (upload_size >> 0) & 0xFF, ] - send_check(sock, file_size_encoded, "binary size") + send_check(sock, upload_size_encoded, "binary size") receive_exactly(sock, 1, "binary size", RESPONSE_UPDATE_PREPARE_OK) - send_check(sock, file_md5, "file checksum") + upload_md5 = hashlib.md5(upload_contents).hexdigest() + _LOGGER.debug("MD5 of upload is %s", upload_md5) + + send_check(sock, upload_md5, "file checksum") receive_exactly(sock, 1, "file checksum", RESPONSE_BIN_MD5_OK) # Disable nodelay for transfer @@ -236,7 +250,7 @@ def perform_ota(sock, password, file_handle, filename): offset = 0 progress = ProgressBar() while True: - chunk = file_handle.read(1024) + chunk = upload_contents[offset : offset + 1024] if not chunk: break offset += len(chunk) @@ -247,7 +261,7 @@ def perform_ota(sock, password, file_handle, filename): sys.stderr.write("\n") raise OTAError(f"Error sending data: {err}") from err - progress.update(offset / float(file_size)) + progress.update(offset / upload_size) progress.done() # Enable nodelay for last checks