Simplify pass: one error channel in the inflate loop, whole-window flushes where acks go out on receipt, shared log helpers, one backend alias, leaner tests

esphome/components/esphome/ota/ota_esphome.cpp

@@ -297,8 +297,8 @@ void ESPHomeOTAComponent::handle_handshake_() {
 
       this->transition_ota_state_(OTAState::FEATURE_ACK);
 
-      const bool supports_compression = (this->ota_features_ & CLIENT_FEATURE_SUPPORTS_COMPRESSION) != 0 &&
-                                        ota::OTABackendPtr::element_type::supports_compression();
+      const bool supports_compression =
+          (this->ota_features_ & CLIENT_FEATURE_SUPPORTS_COMPRESSION) != 0 && ota::OTABackend::supports_compression();
 
       // Compose the feature-ack response. When the client negotiates the extended protocol we emit
       // a 2-byte response (marker + server feature flags); otherwise we emit the single-byte
@@ -781,12 +781,11 @@ ssize_t ESPHomeOTAComponent::receive_data_(uint8_t *buf, DataTransfer &xfer) {
     if (read > 0)
       break;
     if (read == 0) {
-      ESP_LOGW(TAG, "Remote closed");
+      this->log_remote_closed_(LOG_STR("data"));
       return -1;
     }
-    const int err = errno;
-    if (!this->would_block_(err)) {
-      ESP_LOGW(TAG, "Read err %d", err);
+    if (!this->would_block_(errno)) {
+      this->log_socket_error_(LOG_STR("data"));
       return -1;
     }
     // read() already waited up to SO_RCVTIMEO for data, just feed WDT
@@ -838,6 +837,7 @@ ota::OTAResponseTypes ESPHomeOTAComponent::inflate_flush_(InflateSession &sessio
   session.written += pending;
   // A compressible region yields many windows per socket read
   App.feed_wdt();
+  this->ack_written_(*session.xfer);
   return ota::OTA_RESPONSE_OK;
 }
 
@@ -850,21 +850,25 @@ ota::OTAResponseTypes ESPHomeOTAComponent::inflate_data_(uint8_t *in, size_t ima
   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. Everything
-  // received so far is decoded by then, so it is written first and, on the
-  // platforms that ack after the write, acked.
+  // Pulls the next compressed chunk when the decoder runs dry. Where the ack
+  // must follow the write, everything decoded so far is written and acked
+  // first, or the client would wait for an ack while the decoder waits for it.
   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->ack_written_(*s->xfer);
+    if (ACK_AFTER_WRITE) {
+      s->error = s->self->inflate_flush_(*s);
+      if (s->error != ota::OTA_RESPONSE_OK)
+        return -1;
+    }
     // The stream wants more than announced; the size check below reports it
     if (s->xfer->total >= s->xfer->ota_size)
       return -1;
     ssize_t read = s->self->receive_data_(s->in, *s->xfer);
-    if (read <= 0)
+    if (read <= 0) {
+      // Already logged by receive_data_
+      s->error = ota::OTA_RESPONSE_ERROR_UNKNOWN;
       return -1;
+    }
     d->source = s->in + 1;
     d->source_limit = s->in + read;
     return s->in[0];
@@ -883,17 +887,14 @@ ota::OTAResponseTypes ESPHomeOTAComponent::inflate_data_(uint8_t *in, size_t ima
     if (res < 0 || session.eof)
       break;
     session.error = this->inflate_flush_(session);
-    if (session.error != ota::OTA_RESPONSE_OK)
-      break;
-    this->ack_written_(xfer);
-  } while (res != OTA_INFLATE_DONE);
+  } while (res != OTA_INFLATE_DONE && session.error == ota::OTA_RESPONSE_OK);
 
+  // Transport and flash failures are logged where they happen
+  if (session.error != ota::OTA_RESPONSE_OK)
+    return session.error;
   if (res != OTA_INFLATE_DONE || session.written != image_size || xfer.total != xfer.ota_size) {
-    // A transport failure is already logged; a flash error is reported by write_flash_
-    if (session.error == ota::OTA_RESPONSE_OK && (!session.eof || xfer.total == xfer.ota_size)) {
-      ESP_LOGW(TAG, "Inflate err %d", res);
-    }
-    return session.error != ota::OTA_RESPONSE_OK ? session.error : ota::OTA_RESPONSE_ERROR_UNKNOWN;
+    ESP_LOGW(TAG, "Inflate err %d", res);
+    return ota::OTA_RESPONSE_ERROR_UNKNOWN;
   }
   ESP_LOGD(TAG, "Inflated %zu bytes from %zu", session.written, xfer.total);
   return ota::OTA_RESPONSE_OK;

esphome/components/esphome/ota/ota_esphome.h

@@ -152,7 +152,6 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
     if (ACK_AFTER_WRITE)
       this->send_chunk_acks_(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);
@@ -227,7 +226,7 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
     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(),
+  static_assert(!ota::OTABackend::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

esphome/components/esphome/ota/ota_esphome_inflate.c

@@ -61,7 +61,7 @@
   }
 
 /* --------------------------------------------------- *
- * -- uninitialized global data (static structures) -- *
+ * -- constant tables (upstream builds them at runtime) -- *
  * --------------------------------------------------- */
 
 static const unsigned char LENGTH_BITS[30] = {0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2,
@@ -127,7 +127,7 @@ static unsigned char uzlib_get_byte(TINF_DATA *d) {
   /* Otherwise if there's callback and we haven't seen EOF yet, try to
      read next byte using it. (Note: the callback can also update ->source
      and ->source_limit). */
-  if (d->source_read_cb && !d->eof) {
+  if (!d->eof) {
     int val = d->source_read_cb(d);
     if (val >= 0) {
       return (unsigned char) val;

esphome/components/ota/ota_backend_factory.h

@@ -33,6 +33,7 @@ std::unique_ptr<StubOTABackend> make_ota_backend();
 
 namespace esphome::ota {
 using OTABackendPtr = decltype(make_ota_backend());
-static_assert(OTABackendContract<OTABackendPtr::element_type>,
+using OTABackend = OTABackendPtr::element_type;
+static_assert(OTABackendContract<OTABackend>,
               "The platform's OTA backend is missing part of the backend surface (ota_backend.h)");
 }  // namespace esphome::ota

esphome/espota2.py

@@ -654,14 +654,12 @@ def perform_ota(
                 f"retry {flag_name}."
             )
 
-    deflate = False
-    if extended_proto and features & SERVER_FEATURE_SUPPORTS_DEFLATE:
-        # The device inflates while receiving through a small ring window; the
-        # offer is binding, so it wins over gzip should a device set both bits
+    deflate = bool(extended_proto and features & SERVER_FEATURE_SUPPORTS_DEFLATE)
+    if deflate:
+        # The device inflates while receiving through a small ring window
         upload_contents = zlib.compress(
             file_contents, COMPRESS_LEVEL, wbits=-DEFLATE_WINDOW_BITS
         )
-        deflate = True
         _LOGGER.info("Compressed to %s bytes (deflate)", len(upload_contents))
     elif features & SERVER_FEATURE_SUPPORTS_COMPRESSION:
         # The device stores the gzip file and inflates it when it reboots
@@ -733,9 +731,7 @@ def perform_ota(
     if deflate:
         # The device sizes the partition by the inflated image; its own frame,
         # as an encrypted session carries one field per frame
-        send_check(
-            sock, len(file_contents).to_bytes(SIZE_FIELD_BYTES, "big"), "image size"
-        )
+        send_check(sock, file_size.to_bytes(SIZE_FIELD_BYTES, "big"), "image size")
     receive_exactly(sock, 1, "update prepare result", RESPONSE_UPDATE_PREPARE_OK)
     prepare_duration = time.perf_counter() - prepare_start
     _LOGGER.info("Preparing for upload took %.2f seconds", prepare_duration)

tests/components/esphome/test_ota_inflate.cpp

@@ -209,16 +209,20 @@ static const uint8_t STORED[] = {
 static constexpr size_t WINDOW = 4096;
 static constexpr size_t PLAIN_SIZE = 16000;
 
+// Pseudo random bytes reproducible from Python for the vectors above
+static uint8_t lcg_next(uint32_t &x) {
+  x = (x * 1103515245u + 12345u) & 0x7fffffffu;
+  return (x >> 16) & 0xff;
+}
+
 static std::vector<uint8_t> build_plain() {
   std::vector<uint8_t> plain;
   const char *text = "esphome ota deflate ";
   for (int i = 0; i < 300; i++)
     plain.insert(plain.end(), text, text + strlen(text));
   uint32_t x = 1;
-  for (int i = 0; i < 3000; i++) {
-    x = (x * 1103515245u + 12345u) & 0x7fffffffu;
-    plain.push_back((x >> 16) & 0xff);
-  }
+  for (int i = 0; i < 3000; i++)
+    plain.push_back(lcg_next(x));
   plain.insert(plain.end(), 5000, 0);
   for (int i = 0; i < 100; i++)
     plain.insert(plain.end(), text, text + strlen(text));
@@ -233,7 +237,6 @@ struct Session : OtaInflateState {
   size_t in_pos;
   size_t chunk;
   std::vector<uint8_t> out;
-  size_t out_limit;
   uint8_t window[WINDOW];
 };
 
@@ -249,13 +252,12 @@ static int read_cb(OtaInflateState *d) {
 }
 
 // Inflates the whole input; returns the decoder result and fills s.out
-static int inflate_all(Session &s, const uint8_t *in, size_t in_len, size_t chunk, size_t out_limit) {
+static int inflate_all(Session &s, const uint8_t *in, size_t in_len, size_t chunk) {
   s.in = in;
   s.in_len = in_len;
   s.in_pos = 0;
   s.chunk = chunk;
   s.out.clear();
-  s.out_limit = out_limit;
   memset(s.window, 0, sizeof(s.window));
   ota_inflate_init(&s, s.window, WINDOW);
   s.source_read_cb = read_cb;
@@ -267,7 +269,7 @@ static int inflate_all(Session &s, const uint8_t *in, size_t in_len, size_t chun
     if (res < 0 || s.eof)
       return res < 0 ? res : OTA_INFLATE_DATA_ERROR;
     s.out.insert(s.out.end(), s.window, s.dest);
-    if (s.out.size() > s.out_limit)
+    if (s.out.size() > PLAIN_SIZE)
       return OTA_INFLATE_DATA_ERROR;
   } while (res != OTA_INFLATE_DONE);
   return res;
@@ -275,20 +277,20 @@ static int inflate_all(Session &s, const uint8_t *in, size_t in_len, size_t chun
 
 TEST(OtaInflate, RoundTripThroughWindow) {
   auto s = std::make_unique<Session>();
-  ASSERT_EQ(inflate_all(*s, DEFLATED, sizeof(DEFLATED), 1040, PLAIN_SIZE), OTA_INFLATE_DONE);
+  ASSERT_EQ(inflate_all(*s, DEFLATED, sizeof(DEFLATED), 1040), OTA_INFLATE_DONE);
   EXPECT_EQ(s->out, build_plain());
   EXPECT_EQ(s->in_pos, sizeof(DEFLATED));
 }
 
 TEST(OtaInflate, SmallReadChunks) {
   auto s = std::make_unique<Session>();
-  ASSERT_EQ(inflate_all(*s, DEFLATED, sizeof(DEFLATED), 7, PLAIN_SIZE), OTA_INFLATE_DONE);
+  ASSERT_EQ(inflate_all(*s, DEFLATED, sizeof(DEFLATED), 7), OTA_INFLATE_DONE);
   EXPECT_EQ(s->out, build_plain());
 }
 
 TEST(OtaInflate, StoredBlock) {
   auto s = std::make_unique<Session>();
-  ASSERT_EQ(inflate_all(*s, STORED, sizeof(STORED), 64, PLAIN_SIZE), OTA_INFLATE_DONE);
+  ASSERT_EQ(inflate_all(*s, STORED, sizeof(STORED), 64), OTA_INFLATE_DONE);
   auto plain = build_plain();
   plain.resize(300);
   EXPECT_EQ(s->out, plain);
@@ -297,37 +299,33 @@ TEST(OtaInflate, StoredBlock) {
 TEST(OtaInflate, TruncatedStreamFails) {
   auto s = std::make_unique<Session>();
   for (size_t cut : {size_t{1}, size_t{100}, size_t{1000}, sizeof(DEFLATED) - 1}) {
-    EXPECT_LT(inflate_all(*s, DEFLATED, cut, 1040, PLAIN_SIZE), 0) << "cut at " << cut;
+    EXPECT_LT(inflate_all(*s, DEFLATED, cut, 1040), 0) << "cut at " << cut;
     EXPECT_LE(s->out.size(), PLAIN_SIZE);
   }
 }
 
 TEST(OtaInflate, TruncatedStoredBlockFails) {
   auto s = std::make_unique<Session>();
-  EXPECT_LT(inflate_all(*s, STORED, sizeof(STORED) - 50, 64, PLAIN_SIZE), 0);
+  EXPECT_LT(inflate_all(*s, STORED, sizeof(STORED) - 50, 64), 0);
 }
 
 TEST(OtaInflate, CorruptStreamsNeverEscapeTheWindow) {
-  // Every byte of the stream flipped in turn, plus pseudo random garbage: the
-  // decoder must fail or finish without ever reading or writing out of bounds
-  // (the sanitizers check that) and without producing more than announced.
+  // Every third byte of the stream flipped in turn, plus pseudo random garbage:
+  // the sanitizers check that the decoder never reads or writes out of bounds
+  // whatever it returns.
   auto s = std::make_unique<Session>();
   std::vector<uint8_t> bad(DEFLATED, DEFLATED + sizeof(DEFLATED));
   for (size_t i = 0; i < bad.size(); i += 3) {
     bad[i] ^= 0x5a;
-    int res = inflate_all(*s, bad.data(), bad.size(), 1040, PLAIN_SIZE);
-    EXPECT_TRUE(res < 0 || res == OTA_INFLATE_DONE);
+    inflate_all(*s, bad.data(), bad.size(), 1040);
     bad[i] ^= 0x5a;
   }
   uint32_t x = 99;
   std::vector<uint8_t> garbage(2000);
   for (int round = 0; round < 50; round++) {
-    for (auto &b : garbage) {
-      x = (x * 1103515245u + 12345u) & 0x7fffffffu;
-      b = (x >> 16) & 0xff;
-    }
-    int res = inflate_all(*s, garbage.data(), garbage.size(), 1040, PLAIN_SIZE);
-    EXPECT_TRUE(res < 0 || res == OTA_INFLATE_DONE);
+    for (auto &b : garbage)
+      b = lcg_next(x);
+    inflate_all(*s, garbage.data(), garbage.size(), 1040);
   }
 }
 

tests/unit_tests/test_espota2.py

@@ -1528,14 +1528,20 @@ _UPLOAD_TAIL = [
 
 
 @pytest.mark.usefixtures("mock_time")
-def test_perform_ota_with_deflate(mock_socket: Mock) -> None:
+@pytest.mark.parametrize(
+    "server_features",
+    [
+        espota2.SERVER_FEATURE_SUPPORTS_DEFLATE,
+        # A deflate offer is binding, so it wins should a device set both bits
+        espota2.SERVER_FEATURE_SUPPORTS_DEFLATE
+        | espota2.SERVER_FEATURE_SUPPORTS_COMPRESSION,
+    ],
+)
+def test_perform_ota_with_deflate(mock_socket: Mock, server_features: int) -> 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 = (
-        _no_auth_handshake(
-            espota2.OTA_VERSION_2_0, espota2.SERVER_FEATURE_SUPPORTS_DEFLATE
-        )
-        + _UPLOAD_TAIL
+        _no_auth_handshake(espota2.OTA_VERSION_2_0, server_features) + _UPLOAD_TAIL
     )
 
     espota2.perform_ota(mock_socket, None, io.BytesIO(original_content), "test.bin")
@@ -1543,29 +1549,9 @@ def test_perform_ota_with_deflate(mock_socket: Mock) -> None:
     sent = [c[0][0] for c in mock_socket.sendall.call_args_list]
     # magic, features, ota type, size, image size, md5, data, end ack
     sent_size = struct.unpack(">I", sent[3])[0]
-    assert sent[4] == len(original_content).to_bytes(4, "big")
+    assert sent[4] == len(original_content).to_bytes(espota2.SIZE_FIELD_BYTES, "big")
     payload = sent[6]
     assert len(payload) == sent_size < len(original_content)
     # The device decodes through a window of 1 << DEFLATE_WINDOW_BITS bytes
     assert zlib.decompress(payload, -espota2.DEFLATE_WINDOW_BITS) == original_content
     assert sent[5] == hashlib.md5(original_content).hexdigest().encode()
-
-
-@pytest.mark.usefixtures("mock_time")
-def test_perform_ota_deflate_wins_over_gzip(mock_socket: Mock) -> None:
-    """A deflate offer is binding on the device, so it wins should a device set both bits."""
-    original_content = b"firmware" * 100
-    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")
-
-    sent = [c[0][0] for c in mock_socket.sendall.call_args_list]
-    assert sent[4] == len(original_content).to_bytes(4, "big")
-    assert zlib.decompress(sent[6], -espota2.DEFLATE_WINDOW_BITS) == original_content