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.
2026-09-11 07:17:33 +00:00 · 2026-09-08 10:27:20 +02:00
parent 27a8768986
commit db55c1d43f
11 changed files with 78 additions and 71 deletions
@@ -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
@@ -24,7 +24,6 @@

 #include <cerrno>
 #include <cstdio>
-#include <cstddef>
 #include <new>
 #include <sys/time.h>

@@ -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<InflateSession *>(d);
+  session.source_read_cb = [](OtaInflateState *d) -> int {
+    auto *s = static_cast<InflateSession *>(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);
@@ -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<InflateSession> inflate_;
 #endif
@@ -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};
@@ -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};
@@ -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
@@ -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
@@ -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<StubOTABackend> make_ota_backend();
 }  // namespace esphome::ota
@@ -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_{};
@@ -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
@@ -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")