Ack chunks after the flash write on the inflate path and harden the decoder

Chunk acks again mean the block is in flash on both paths: the read callback
writes and acks everything decoded so far before it waits for more input.
The decoder rejects a negative distance symbol, initialises its whole state,
and the flush feeds the watchdog once per window. The static_assert on the
backend skips the static analysis build, which sets every define at once.

esphome/components/esphome/ota/ota_esphome.cpp

@@ -537,6 +537,7 @@ void ESPHomeOTAComponent::handle_data_() {
       error_code = this->write_flash_(buf, read);
       if (error_code != ota::OTA_RESPONSE_OK)
         goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
+      this->send_chunk_acks_(xfer);
     }
   }
 
@@ -793,13 +794,6 @@ ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) {
   const uint32_t now = millis();
   xfer.last_data_ms = now;
   xfer.total += read;
-#if USE_OTA_VERSION == 2
-  while (xfer.acknowledged + OTA_BLOCK_SIZE <= xfer.total ||
-         (xfer.total == xfer.ota_size && xfer.acknowledged < xfer.ota_size)) {
-    this->data_write_byte_(ota::OTA_RESPONSE_CHUNK_OK);
-    xfer.acknowledged += OTA_BLOCK_SIZE;
-  }
-#endif
   if (now - xfer.last_progress > OTA_PROGRESS_INTERVAL_MS) {
     xfer.last_progress = now;
     float percentage = (xfer.total * 100.0f) / xfer.ota_size;
@@ -813,19 +807,61 @@ ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) {
   return read;
 }
 
+void ESPHomeOTAComponent::send_chunk_acks_(DataTransfer &xfer) {
+#if USE_OTA_VERSION == 2
+  while (xfer.acknowledged + OTA_BLOCK_SIZE <= xfer.total ||
+         (xfer.total == xfer.ota_size && xfer.acknowledged < xfer.ota_size)) {
+    this->data_write_byte_(ota::OTA_RESPONSE_CHUNK_OK);
+    xfer.acknowledged += OTA_BLOCK_SIZE;
+  }
+#endif
+}
+
 #ifdef USE_OTA_DEFLATE
 // The window doubles as the output buffer: the decoder fills it, we flush it to
 // the backend, and its bytes remain available as the back-reference history for
 // the next windowful.
+ota::OTAResponseTypes ESPHomeOTAComponent::inflate_flush_(InflateSession &session) {
+  const size_t produced = session.dest - session.window;
+  const size_t pending = produced - session.flushed;
+  if (pending == 0)
+    return ota::OTA_RESPONSE_OK;
+  if (pending > session.image_size - session.written) {
+    ESP_LOGW(TAG, "Inflate size mismatch");
+    return ota::OTA_RESPONSE_ERROR_UNKNOWN;
+  }
+  ota::OTAResponseTypes result = this->write_flash_(session.window + session.flushed, pending);
+  if (result != ota::OTA_RESPONSE_OK)
+    return result;
+  session.flushed = produced;
+  session.written += pending;
+  // A compressible region yields many windows per socket read
+  App.feed_wdt();
+  return ota::OTA_RESPONSE_OK;
+}
+
 ota::OTAResponseTypes ESPHomeOTAComponent::inflate_data_(uint8_t *in, size_t image_size, DataTransfer &xfer) {
   InflateSession &session = *this->inflate_;
   session.self = this;
   session.xfer = &xfer;
   session.in = in;
+  session.image_size = image_size;
+  session.written = 0;
+  session.error = ota::OTA_RESPONSE_OK;
   ota_inflate_init(&session, session.window, OTA_INFLATE_WINDOW_SIZE);
-  // Pulls the next compressed chunk when the decoder runs dry
+  // Pulls the next compressed chunk when the decoder runs dry. Everything
+  // received so far is decoded by then, so it is written and acked first,
+  // keeping a chunk ack meaning "in flash" as on the uncompressed path.
   session.source_read_cb = [](OtaInflateState *d) -> int {
     auto *s = static_cast<InflateSession *>(d);
+    s->error = s->self->inflate_flush_(*s);
+    if (s->error != ota::OTA_RESPONSE_OK)
+      return -1;
+    s->self->send_chunk_acks_(*s->xfer);
+    if (s->xfer->total >= s->xfer->ota_size) {
+      ESP_LOGW(TAG, "Inflate size mismatch");
+      return -1;
+    }
     ssize_t read = s->self->receive_data_(s->in, *s->xfer);
     if (read <= 0)
       return -1;
@@ -834,34 +870,32 @@ ota::OTAResponseTypes ESPHomeOTAComponent::inflate_data_(uint8_t *in, size_t ima
     return s->in[0];
   };
 
-  size_t written = 0;
   int res;
   do {
+    // The ring index wrapped to 0 exactly when the window filled, so the
+    // window and the output cursor stay in lockstep
     session.dest = session.window;
     session.dest_limit = session.window + OTA_INFLATE_WINDOW_SIZE;
+    session.flushed = 0;
     res = ota_inflate(&session);
     if (res < 0) {
       // eof means the read callback failed, which is already logged
       if (!session.eof) {
         ESP_LOGW(TAG, "Inflate err %d", res);
       }
-      return ota::OTA_RESPONSE_ERROR_UNKNOWN;
+      return session.error != ota::OTA_RESPONSE_OK ? session.error : ota::OTA_RESPONSE_ERROR_UNKNOWN;
     }
-    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;
+    ota::OTAResponseTypes flush_result = this->inflate_flush_(session);
+    if (flush_result != ota::OTA_RESPONSE_OK)
+      return flush_result;
+    this->send_chunk_acks_(xfer);
   } while (res != OTA_INFLATE_DONE);
 
-  if (written != image_size || xfer.total != xfer.ota_size) {
+  if (session.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);
+  ESP_LOGD(TAG, "Inflated %zu bytes from %zu", session.written, xfer.total);
   return ota::OTA_RESPONSE_OK;
 }
 #endif  // USE_OTA_DEFLATE

esphome/components/esphome/ota/ota_esphome.h

@@ -133,9 +133,10 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
     uint32_t last_progress{0};
   };
   // 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).
+  // the data timeout, and updates xfer. Returns bytes read, -1 on failure (logged).
   inline ssize_t receive_data_(uint8_t *buf, DataTransfer &xfer);
+  // Acks every received block once it is in flash, so an ack means written
+  inline void send_chunk_acks_(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
@@ -204,10 +205,18 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
     ESPHomeOTAComponent *self;
     DataTransfer *xfer;
     uint8_t *in;  // caller's buffer for the compressed input, valid during inflate_data_
+    size_t image_size;
+    size_t written;               // inflated bytes in flash
+    size_t flushed;               // bytes of the current window already in flash
+    ota::OTAResponseTypes error;  // first failure inside the read callback
     uint8_t window[OTA_INFLATE_WINDOW_SIZE];
   };
+#ifndef CLANG_TIDY  // static analysis sets every define at once
   static_assert(!ota::OTABackendPtr::element_type::supports_compression(),
                 "USE_OTA_DEFLATE is for backends that cannot store a gzip image");
+#endif
+  // Writes the decoded bytes not yet in flash; dest stays put so the ring history is intact
+  ota::OTAResponseTypes inflate_flush_(InflateSession &session);
   ota::OTAResponseTypes inflate_data_(uint8_t *in, size_t image_size, DataTransfer &xfer);
   std::unique_ptr<InflateSession> inflate_;
 #endif

esphome/components/esphome/ota/ota_esphome_inflate.c

@@ -41,6 +41,8 @@
 
 #include "ota_esphome_inflate.h"
 
+#include <stddef.h>
+
 #define TINF_OK OTA_INFLATE_OK
 #define TINF_DONE OTA_INFLATE_DONE
 #define TINF_DATA_ERROR OTA_INFLATE_DATA_ERROR
@@ -353,7 +355,7 @@ static int tinf_inflate_block_data(TINF_DATA *d, TINF_TREE *lt, TINF_TREE *dt) {
     d->curlen = tinf_read_bits(d, LENGTH_BITS[sym], LENGTH_BASE[sym]);
 
     dist = tinf_decode_symbol(d, dt);
-    if (dist >= 30) {
+    if (dist < 0 || dist >= 30) {
       return TINF_DATA_ERROR;
     }
 
@@ -425,8 +427,12 @@ static int tinf_inflate_uncompressed_block(TINF_DATA *d) {
 
 /* initialize decompression structure */
 void ota_inflate_init(TINF_DATA *d, unsigned char *dict, unsigned int dict_len) {
+  d->source = NULL;
+  d->source_limit = NULL;
+  d->tag = 0;
   d->eof = 0;
   d->bitcount = 0;
+  d->lz_off = 0;
   d->bfinal = 0;
   d->btype = -1;
   d->dict_size = dict_len;

esphome/components/ota/ota_backend_arduino_rp2.h

@@ -16,7 +16,9 @@ class ArduinoRP2OTABackend final {
   OTAResponseTypes end();
   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
+  // into the app region at reboot, the same way the ESP8266 bootloader does.
+  // begin() then sees the gzip size, so only the staging space is checked up
+  // front; the inflated size is not known until the stub reads the trailer.
   static constexpr bool supports_compression() { return true; }
 
  private:

tests/integration/test_host_ota.py

@@ -186,7 +186,7 @@ async def test_host_ota_self_update(
         if "OTA staged at" in line:
             staged.set()
         # The host backend has no gzip support, so the upload negotiates deflate
-        if "Inflated " in line:
+        if " bytes from " in line:
             inflated.set()
         dev.on_log(line)
 

tests/unit_tests/test_espota2.py

@@ -1567,7 +1567,8 @@ def test_perform_ota_gzip_wins_over_deflate(mock_socket: Mock) -> None:
     espota2.perform_ota(mock_socket, None, io.BytesIO(original_content), "test.bin")
 
     sent = [c[0][0] for c in mock_socket.sendall.call_args_list]
-    compressed = gzip.compress(original_content, compresslevel=9)
-    assert sent[3] == len(compressed).to_bytes(4, "big")
-    assert sent[4] == hashlib.md5(compressed).hexdigest().encode()
-    assert sent[5] == compressed
+    # gzip output embeds the time of compression, so compare what it holds
+    payload = sent[5]
+    assert gzip.decompress(payload) == original_content
+    assert sent[3] == len(payload).to_bytes(4, "big")
+    assert sent[4] == hashlib.md5(payload).hexdigest().encode()