From db55c1d43f4f6d309bfde2b8b86a0bc182a348e1 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Tue, 8 Sep 2026 10:27:20 +0200 Subject: [PATCH] Trim the inflate glue and tie the deflate gate to the backend Measured on ESP32: the decoder is 1472 B at -Os and cannot shrink without dropping dynamic Huffman; the glue loses its extra log sites and strings, the flash-write log is shared, and the single-call helpers inline on the platforms without the decoder so the ESP8266 and RP2040 images do not grow. supports_compression() is constexpr so the deflate build asserts that its backend cannot store gzip. --- esphome/__main__.py | 7 +-- .../components/esphome/ota/ota_esphome.cpp | 57 +++++++++---------- esphome/components/esphome/ota/ota_esphome.h | 16 ++++-- .../ota/ota_backend_arduino_libretiny.h | 2 +- .../components/ota/ota_backend_arduino_rp2.h | 2 +- esphome/components/ota/ota_backend_esp8266.h | 2 +- esphome/components/ota/ota_backend_esp_idf.h | 2 +- esphome/components/ota/ota_backend_factory.h | 2 +- esphome/components/ota/ota_backend_host.h | 2 +- esphome/espota2.py | 8 +-- tests/unit_tests/test_espota2.py | 49 +++++++++------- 11 files changed, 78 insertions(+), 71 deletions(-) diff --git a/esphome/__main__.py b/esphome/__main__.py index 30e97f55eb..d4c3d787c9 100644 --- a/esphome/__main__.py +++ b/esphome/__main__.py @@ -1289,10 +1289,9 @@ def _choose_ota_platform(config: ConfigType, requested: str | None) -> str: The native API uses challenge-response auth with MD5/SHA256 hashing of a server-issued nonce, so the password is never sent over the wire; the ``web_server`` path uses HTTP Basic auth which transmits credentials in - cleartext over the LAN. (The native path also supports gzip compression - on ESP8266, where flash space is tight; on ESP32/RP2040/LibreTiny the - backend reports ``supports_compression() == false`` and the firmware is - sent uncompressed regardless of which platform is used.) Falls back to + cleartext over the LAN. (The native path also compresses the upload: + gzip on ESP8266 and RP2040, which inflate it at reboot, and a deflate + stream on ESP32/LibreTiny, which inflate it as it arrives.) Falls back to ``web_server`` only when that is the only available platform. """ # Use a dict (insertion-ordered) instead of a list so error messages and diff --git a/esphome/components/esphome/ota/ota_esphome.cpp b/esphome/components/esphome/ota/ota_esphome.cpp index 4a2df014d7..d226726ef0 100644 --- a/esphome/components/esphome/ota/ota_esphome.cpp +++ b/esphome/components/esphome/ota/ota_esphome.cpp @@ -24,7 +24,6 @@ #include #include -#include #include #include @@ -319,15 +318,14 @@ void ESPHomeOTAComponent::handle_handshake_() { this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_NOISE; #endif #ifdef USE_OTA_DEFLATE - // The backend cannot store gzip here (USE_OTA_DEFLATE is not set on the - // one that can), so inflate on the fly when the client offers it and the - // session memory (a few KB) is in hand; otherwise stay uncompressed - if ((this->ota_features_ & CLIENT_FEATURE_SUPPORTS_DEFLATE) != 0 && !supports_compression) { + // Offer to inflate on the fly once the session memory (a few KB) is in + // hand; otherwise the upload stays uncompressed + if ((this->ota_features_ & CLIENT_FEATURE_SUPPORTS_DEFLATE) != 0) { this->inflate_.reset(new (std::nothrow) InflateSession()); if (this->inflate_ != nullptr) { this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_DEFLATE; } else { - ESP_LOGW(TAG, "No memory to inflate, upload will be uncompressed"); + ESP_LOGW(TAG, "No memory to inflate"); } } #endif @@ -536,11 +534,9 @@ void ESPHomeOTAComponent::handle_data_() { error_code = ota::OTA_RESPONSE_ERROR_UNKNOWN; goto error; // NOLINT(cppcoreguidelines-avoid-goto) } - error_code = this->backend_->write(buf, read); - if (error_code != ota::OTA_RESPONSE_OK) { - ESP_LOGW(TAG, "Flash write err %d", error_code); + error_code = this->write_flash_(buf, read); + if (error_code != ota::OTA_RESPONSE_OK) goto error; // NOLINT(cppcoreguidelines-avoid-goto) - } } } @@ -747,6 +743,14 @@ bool ESPHomeOTAComponent::read_size_(uint8_t *buf, size_t &size, const LogString return true; } +ota::OTAResponseTypes ESPHomeOTAComponent::write_flash_(uint8_t *data, size_t len) { + ota::OTAResponseTypes result = this->backend_->write(data, len); + if (result != ota::OTA_RESPONSE_OK) { + ESP_LOGW(TAG, "Flash write err %d", result); + } + return result; +} + ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) { const size_t remaining = xfer.ota_size - xfer.total; const size_t requested = std::min(remaining, OTA_BUFFER_SIZE); @@ -814,16 +818,14 @@ ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) { // the backend, and its bytes remain available as the back-reference history for // the next windowful. ota::OTAResponseTypes ESPHomeOTAComponent::inflate_data_(uint8_t *in, size_t image_size, DataTransfer &xfer) { - static_assert(offsetof(InflateSession, state) == 0, "the read callback recovers the session from &state"); InflateSession &session = *this->inflate_; - OtaInflateState &state = session.state; session.self = this; session.xfer = &xfer; session.in = in; - ota_inflate_init(&state, session.window, OTA_INFLATE_WINDOW_SIZE); + ota_inflate_init(&session, session.window, OTA_INFLATE_WINDOW_SIZE); // Pulls the next compressed chunk when the decoder runs dry - state.source_read_cb = [](OtaInflateState *d) -> int { - auto *s = reinterpret_cast(d); + session.source_read_cb = [](OtaInflateState *d) -> int { + auto *s = static_cast(d); ssize_t read = s->self->receive_data_(s->in, *s->xfer); if (read <= 0) return -1; @@ -835,31 +837,28 @@ ota::OTAResponseTypes ESPHomeOTAComponent::inflate_data_(uint8_t *in, size_t ima size_t written = 0; int res; do { - state.dest = session.window; - state.dest_limit = session.window + OTA_INFLATE_WINDOW_SIZE; - res = ota_inflate(&state); + session.dest = session.window; + session.dest_limit = session.window + OTA_INFLATE_WINDOW_SIZE; + res = ota_inflate(&session); if (res < 0) { // eof means the read callback failed, which is already logged - if (!state.eof) { + if (!session.eof) { ESP_LOGW(TAG, "Inflate err %d", res); } return ota::OTA_RESPONSE_ERROR_UNKNOWN; } - const size_t produced = state.dest - session.window; - if (produced > image_size - written) { - ESP_LOGW(TAG, "Image exceeds announced size"); - return ota::OTA_RESPONSE_ERROR_UNKNOWN; - } - ota::OTAResponseTypes write_result = this->backend_->write(session.window, produced); - if (write_result != ota::OTA_RESPONSE_OK) { - ESP_LOGW(TAG, "Flash write err %d", write_result); + const size_t produced = session.dest - session.window; + // More output than announced: stop before the write and report it below + if (produced > image_size - written) + break; + ota::OTAResponseTypes write_result = this->write_flash_(session.window, produced); + if (write_result != ota::OTA_RESPONSE_OK) return write_result; - } written += produced; } while (res != OTA_INFLATE_DONE); if (written != image_size || xfer.total != xfer.ota_size) { - ESP_LOGW(TAG, "Inflated %zu of %zu bytes from %zu of %zu", written, image_size, xfer.total, xfer.ota_size); + ESP_LOGW(TAG, "Inflate size mismatch"); return ota::OTA_RESPONSE_ERROR_UNKNOWN; } ESP_LOGD(TAG, "Inflated %zu bytes from %zu", written, xfer.total); diff --git a/esphome/components/esphome/ota/ota_esphome.h b/esphome/components/esphome/ota/ota_esphome.h index 17bf2778d7..3a2b720e97 100644 --- a/esphome/components/esphome/ota/ota_esphome.h +++ b/esphome/components/esphome/ota/ota_esphome.h @@ -135,9 +135,11 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { // Receives up to OTA_BUFFER_SIZE bytes of upload data into buf, waiting up to // the data timeout; updates xfer and sends chunk acks. Returns bytes read, -1 // on failure (logged). - ssize_t receive_data_(uint8_t *buf, DataTransfer &xfer); - // Reads a 4 byte MSB first size field; buf must hold OTA_BUFFER_SIZE bytes - bool read_size_(uint8_t *buf, size_t &size, const LogString *desc); + inline ssize_t receive_data_(uint8_t *buf, DataTransfer &xfer); + // Reads a 4 byte MSB first size field into size + inline bool read_size_(uint8_t *buf, size_t &size, const LogString *desc); + // Writes to the backend and logs a failure + inline ota::OTAResponseTypes write_flash_(uint8_t *data, size_t len); bool try_read_(size_t to_read, const LogString *desc); bool try_write_(size_t to_write, const LogString *desc); @@ -196,14 +198,16 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { // ring window must be at least that. It also serves as the inflate output // buffer, so it is flushed to the backend one windowful at a time. static constexpr size_t OTA_INFLATE_WINDOW_SIZE = 4096; - // Heap-allocated only while a deflate-compressed upload is negotiated. - struct InflateSession { - OtaInflateState state; // first member: the read callback casts back from it + // Heap-allocated only while a deflate-compressed upload is negotiated; the + // decoder state is the base so its read callback can recover the session + struct InflateSession : OtaInflateState { ESPHomeOTAComponent *self; DataTransfer *xfer; uint8_t *in; // caller's buffer for the compressed input, valid during inflate_data_ uint8_t window[OTA_INFLATE_WINDOW_SIZE]; }; + static_assert(!ota::OTABackendPtr::element_type::supports_compression(), + "USE_OTA_DEFLATE is for backends that cannot store a gzip image"); ota::OTAResponseTypes inflate_data_(uint8_t *in, size_t image_size, DataTransfer &xfer); std::unique_ptr inflate_; #endif diff --git a/esphome/components/ota/ota_backend_arduino_libretiny.h b/esphome/components/ota/ota_backend_arduino_libretiny.h index c2716a44d1..c322ed21f2 100644 --- a/esphome/components/ota/ota_backend_arduino_libretiny.h +++ b/esphome/components/ota/ota_backend_arduino_libretiny.h @@ -13,7 +13,7 @@ class ArduinoLibreTinyOTABackend final { OTAResponseTypes write(uint8_t *data, size_t len); OTAResponseTypes end(); void abort(); - bool supports_compression() { return false; } + static constexpr bool supports_compression() { return false; } private: bool md5_set_{false}; diff --git a/esphome/components/ota/ota_backend_arduino_rp2.h b/esphome/components/ota/ota_backend_arduino_rp2.h index 75f12f44b0..47a23e2d0a 100644 --- a/esphome/components/ota/ota_backend_arduino_rp2.h +++ b/esphome/components/ota/ota_backend_arduino_rp2.h @@ -17,7 +17,7 @@ class ArduinoRP2OTABackend final { void abort(); // A gzip image is staged on LittleFS as is; the core's OTA stub inflates it // into the app region at reboot, the same way the ESP8266 bootloader does - bool supports_compression() { return true; } + static constexpr bool supports_compression() { return true; } private: bool md5_set_{false}; diff --git a/esphome/components/ota/ota_backend_esp8266.h b/esphome/components/ota/ota_backend_esp8266.h index 21b5c12c2d..1f1ec37eee 100644 --- a/esphome/components/ota/ota_backend_esp8266.h +++ b/esphome/components/ota/ota_backend_esp8266.h @@ -20,7 +20,7 @@ class ESP8266OTABackend final { OTAResponseTypes end(); void abort(); // Compression supported in all ESP8266 Arduino versions ESPHome supports (>= 2.7.0) - bool supports_compression() { return true; } + static constexpr bool supports_compression() { return true; } protected: /// Erase flash sector if current address is at sector boundary diff --git a/esphome/components/ota/ota_backend_esp_idf.h b/esphome/components/ota/ota_backend_esp_idf.h index c991f896e8..4f4093a594 100644 --- a/esphome/components/ota/ota_backend_esp_idf.h +++ b/esphome/components/ota/ota_backend_esp_idf.h @@ -33,7 +33,7 @@ class IDFOTABackend final { OTAResponseTypes write(uint8_t *data, size_t len); OTAResponseTypes end(); void abort(); - bool supports_compression() { return false; } + static constexpr bool supports_compression() { return false; } protected: #ifdef USE_OTA_PARTITIONS diff --git a/esphome/components/ota/ota_backend_factory.h b/esphome/components/ota/ota_backend_factory.h index 82d001ed9e..29da4f9d7d 100644 --- a/esphome/components/ota/ota_backend_factory.h +++ b/esphome/components/ota/ota_backend_factory.h @@ -25,7 +25,7 @@ struct StubOTABackend { OTAResponseTypes write(uint8_t *data, size_t len) { return OTA_RESPONSE_ERROR_UNKNOWN; } OTAResponseTypes end() { return OTA_RESPONSE_ERROR_UNKNOWN; } void abort() {} - bool supports_compression() { return false; } + static constexpr bool supports_compression() { return false; } }; std::unique_ptr make_ota_backend(); } // namespace esphome::ota diff --git a/esphome/components/ota/ota_backend_host.h b/esphome/components/ota/ota_backend_host.h index 51ffdaeda3..e53868f102 100644 --- a/esphome/components/ota/ota_backend_host.h +++ b/esphome/components/ota/ota_backend_host.h @@ -19,7 +19,7 @@ class HostOTABackend final { OTAResponseTypes write(uint8_t *data, size_t len); OTAResponseTypes end(); void abort(); - bool supports_compression() { return false; } + static constexpr bool supports_compression() { return false; } protected: md5::MD5Digest md5_{}; diff --git a/esphome/espota2.py b/esphome/espota2.py index 0220815a75..379c4bd796 100644 --- a/esphome/espota2.py +++ b/esphome/espota2.py @@ -658,10 +658,9 @@ def perform_ota( _LOGGER.info("Compressed to %s bytes", len(upload_contents)) elif extended_proto and features & SERVER_FEATURE_SUPPORTS_DEFLATE: # The device inflates while receiving through a small ring window - compressor = zlib.compressobj( - COMPRESS_LEVEL, zlib.DEFLATED, -DEFLATE_WINDOW_BITS + upload_contents = zlib.compress( + file_contents, COMPRESS_LEVEL, wbits=-DEFLATE_WINDOW_BITS ) - upload_contents = compressor.compress(file_contents) + compressor.flush() deflate = True _LOGGER.info("Compressed to %s bytes (deflate)", len(upload_contents)) else: @@ -722,12 +721,11 @@ def perform_ota( send_check(sock, ota_type, "ota type") upload_size = len(upload_contents) - upload_size_encoded = upload_size.to_bytes(SIZE_FIELD_BYTES, "big") # The device erases flash between receiving the size and acking the # prepare, so this window shows the erase cost (near zero when the # device erases lazily during the upload) prepare_start = time.perf_counter() - send_check(sock, upload_size_encoded, "binary size") + send_check(sock, upload_size.to_bytes(SIZE_FIELD_BYTES, "big"), "binary size") if deflate: # The device sizes the partition by the inflated image; its own frame, # as an encrypted session carries one field per frame diff --git a/tests/unit_tests/test_espota2.py b/tests/unit_tests/test_espota2.py index 54112e4d2b..956b4cc2ab 100644 --- a/tests/unit_tests/test_espota2.py +++ b/tests/unit_tests/test_espota2.py @@ -600,12 +600,19 @@ def test_perform_ota_upload_error(mock_socket: Mock, mock_file: io.BytesIO) -> N espota2.perform_ota(mock_socket, None, mock_file, "test.bin") -def _no_auth_handshake(version: int) -> list[bytes]: - """Recv responses for a handshake without auth, up to the MD5 check.""" +def _no_auth_handshake(version: int, server_features: int | None = None) -> list[bytes]: + """Recv responses for a handshake without auth, up to the MD5 check. + + With server_features the device answers with the extended feature flags. + """ + if server_features is None: + features = [bytes([espota2.RESPONSE_HEADER_OK])] + else: + features = [bytes([espota2.RESPONSE_FEATURE_FLAGS]), bytes([server_features])] return [ bytes([espota2.RESPONSE_OK]), # First byte of version response bytes([version]), # Version number - bytes([espota2.RESPONSE_HEADER_OK]), # Features response + *features, bytes([espota2.RESPONSE_AUTH_OK]), # No auth required bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]), # Binary size OK bytes([espota2.RESPONSE_BIN_MD5_OK]), # MD5 checksum OK @@ -1512,27 +1519,23 @@ def test_check_error_passes_non_error_when_expect_is_none() -> None: espota2.check_error([espota2.RESPONSE_FEATURE_FLAGS], None) -def _deflate_handshake(server_features: int) -> list[bytes]: - return [ - bytes([espota2.RESPONSE_OK]), - bytes([espota2.OTA_VERSION_2_0]), - bytes([espota2.RESPONSE_FEATURE_FLAGS]), - bytes([server_features]), - bytes([espota2.RESPONSE_AUTH_OK]), - bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]), - bytes([espota2.RESPONSE_BIN_MD5_OK]), - bytes([espota2.RESPONSE_CHUNK_OK]), - bytes([espota2.RESPONSE_RECEIVE_OK]), - bytes([espota2.RESPONSE_UPDATE_END_OK]), - ] +# Device replies after the MD5 check for a one-chunk upload +_UPLOAD_TAIL = [ + bytes([espota2.RESPONSE_CHUNK_OK]), + bytes([espota2.RESPONSE_RECEIVE_OK]), + bytes([espota2.RESPONSE_UPDATE_END_OK]), +] @pytest.mark.usefixtures("mock_time") def test_perform_ota_with_deflate(mock_socket: Mock) -> None: """A device that inflates on the fly gets a raw deflate stream, both sizes and the image MD5.""" original_content = b"firmware" * 100 - mock_socket.recv.side_effect = _deflate_handshake( - espota2.SERVER_FEATURE_SUPPORTS_DEFLATE + mock_socket.recv.side_effect = ( + _no_auth_handshake( + espota2.OTA_VERSION_2_0, espota2.SERVER_FEATURE_SUPPORTS_DEFLATE + ) + + _UPLOAD_TAIL ) espota2.perform_ota(mock_socket, None, io.BytesIO(original_content), "test.bin") @@ -1552,9 +1555,13 @@ def test_perform_ota_with_deflate(mock_socket: Mock) -> None: def test_perform_ota_gzip_wins_over_deflate(mock_socket: Mock) -> None: """A device that can store gzip keeps getting gzip even when it also offers deflate.""" original_content = b"firmware" * 100 - mock_socket.recv.side_effect = _deflate_handshake( - espota2.SERVER_FEATURE_SUPPORTS_COMPRESSION - | espota2.SERVER_FEATURE_SUPPORTS_DEFLATE + mock_socket.recv.side_effect = ( + _no_auth_handshake( + espota2.OTA_VERSION_2_0, + espota2.SERVER_FEATURE_SUPPORTS_COMPRESSION + | espota2.SERVER_FEATURE_SUPPORTS_DEFLATE, + ) + + _UPLOAD_TAIL ) espota2.perform_ota(mock_socket, None, io.BytesIO(original_content), "test.bin")