Fold the duplicated noise gate and inflate bookkeeping

One predicate for the noise offer and one place that allocates the
session; the window is already zero from value initialization, so drop
the second clearing pass. The RP2 backend reads the framework version
from the define ESPHome already emits rather than a vendor header, the
inflate counter moves into the test device helper, and the corrupt
stream sweep takes a coarser stride for the same coverage.

esphome/components/esphome/ota/ota_esphome.cpp

@@ -22,6 +22,7 @@
 #include "esphome/core/lwip_fast_select.h"
 #endif
 
+#include <algorithm>
 #include <cerrno>
 #include <cstdio>
 #include <cstring>
@@ -197,6 +198,13 @@ static constexpr uint8_t SERVER_FEATURE_SUPPORTS_NOISE = 0x04;
 // client must then send the image size frame and a deflate stream.
 static constexpr uint8_t SERVER_FEATURE_SUPPORTS_DEFLATE = 0x08;
 
+#ifdef USE_OTA_ENCRYPTION
+inline bool ESPHomeOTAComponent::noise_offered_() const {
+  return (this->handshake_buf_[1] & SERVER_FEATURE_SUPPORTS_NOISE) != 0 &&
+         (this->ota_features_ & CLIENT_NOISE_FEATURES) == CLIENT_NOISE_FEATURES;
+}
+#endif
+
 inline bool ESPHomeOTAComponent::extended_proto_() const {
 #ifdef USE_OTA_ENCRYPTION_REQUIRED
   // FEATURE_READ already refused every client without the extended protocol
@@ -324,17 +332,16 @@ void ESPHomeOTAComponent::handle_handshake_() {
 #endif
 #ifdef USE_OTA_ENCRYPTION
         // Reserve the noise session before the optional inflate buffer, so the
-        // required allocation is not starved by the compression window. Gated
-        // on the same condition that starts the session in FEATURE_ACK.
-        if ((this->handshake_buf_[1] & SERVER_FEATURE_SUPPORTS_NOISE) != 0 &&
-            (this->ota_features_ & CLIENT_NOISE_FEATURES) == CLIENT_NOISE_FEATURES) {
-          // NOLINTNEXTLINE(clang-analyzer-cplusplus.NewDeleteLeaks)
-          this->noise_ = std::unique_ptr<NoiseSession>(new (std::nothrow) NoiseSession);
+        // required allocation is not starved by the compression window
+        if (this->noise_offered_()) {
+          this->noise_reserve_session_();
         }
 #endif
 #ifdef USE_OTA_DEFLATE
         // Offered only once the session memory is in hand; else uncompressed
         if ((this->ota_features_ & CLIENT_FEATURE_SUPPORTS_DEFLATE) != 0) {
+          // Value initialized: a corrupt stream that back references the
+          // window before it is filled then copies zeros, never stale memory
           this->inflate_.reset(new (std::nothrow) InflateSession());
           if (this->inflate_ != nullptr) {
             this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_DEFLATE;
@@ -360,8 +367,7 @@ void ESPHomeOTAComponent::handle_handshake_() {
 #ifdef USE_OTA_ENCRYPTION
       // Latch the offer actually sent: a key activating between the two
       // states must not start a session the client never expects
-      if ((this->handshake_buf_[1] & SERVER_FEATURE_SUPPORTS_NOISE) != 0 &&
-          (this->ota_features_ & CLIENT_NOISE_FEATURES) == CLIENT_NOISE_FEATURES) {
+      if (this->noise_offered_()) {
         // handshake_buf_ still holds the feature ack composed above; a
         // would-block re-entry lands here without rebuilding it
         if (!this->noise_start_session_(this->handshake_buf_[1])) {
@@ -862,9 +868,6 @@ ota::OTAResponseTypes ESPHomeOTAComponent::inflate_data_(uint8_t *in, size_t ima
   session.image_size = image_size;
   session.written = 0;
   session.error = ota::OTA_RESPONSE_OK;
-  // A corrupt stream may back reference the window before it is filled; zero it
-  // so such a read copies zeros, never stale memory
-  std::memset(session.window, 0, sizeof(session.window));
   ota_inflate_init(&session, session.window, OTA_INFLATE_WINDOW_SIZE);
   // Where the ack must follow the write, flush and ack before waiting for
   // input, or the client waits for an ack while the decoder waits for data

esphome/components/esphome/ota/ota_esphome.h

@@ -91,6 +91,9 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
   };
   // The api server's live context when the api has encryption, else our own
   const noise::NoiseContext &noise_context_() const;
+  // True once the feature ack offers noise and the client asked for it
+  bool noise_offered_() const;
+  void noise_reserve_session_();
   bool noise_start_session_(uint8_t server_feature_flags);
   bool handle_noise_handshake_();
   bool noise_try_read_frame_();

esphome/components/esphome/ota/ota_esphome_noise.cpp

@@ -33,7 +33,13 @@ ESPHomeOTAComponent::NoiseSession::~NoiseSession() {
   }
 }
 
-/** Allocate the session and start the responder handshake.
+void ESPHomeOTAComponent::noise_reserve_session_() {
+  // Default-init: the frame buffer is written before it is read
+  // NOLINTNEXTLINE(clang-analyzer-cplusplus.NewDeleteLeaks)
+  this->noise_ = std::unique_ptr<NoiseSession>(new (std::nothrow) NoiseSession);
+}
+
+/** Start the responder handshake, on the session reserved at offer time.
  *
  * The prologue binds the whole plaintext preamble, so any tampering with the
  * negotiation (a stripped feature flag, a changed version) breaks the first
@@ -42,13 +48,7 @@ ESPHomeOTAComponent::NoiseSession::~NoiseSession() {
  */
 bool ESPHomeOTAComponent::noise_start_session_(uint8_t server_feature_flags) {
   // A provisioned key cleared between the offer and here is not guarded: the
-  // session runs on the zero key load_psk fills in and fails the client's MAC.
-  // Reuse the session reserved at offer time, else allocate now. Default-init:
-  // the frame buffer is written before it is read
-  if (this->noise_ == nullptr) {
-    // NOLINTNEXTLINE(clang-analyzer-cplusplus.NewDeleteLeaks)
-    this->noise_ = std::unique_ptr<NoiseSession>(new (std::nothrow) NoiseSession);
-  }
+  // session runs on the zero key load_psk fills in and fails the client's MAC
   static constexpr size_t PROLOGUE_ACK_LEN = 2;  // OTA_RESPONSE_OK + version
   static constexpr size_t PROLOGUE_CLIENT_FEATURES_LEN = 1;
   static constexpr size_t PROLOGUE_FEATURE_ACK_LEN = 2;  // OTA_RESPONSE_FEATURE_FLAGS + server flags

esphome/components/ota/ota_backend_arduino_rp2.h

@@ -6,8 +6,6 @@
 #include "esphome/core/defines.h"
 #include "esphome/core/macros.h"
 
-#include <RP2040Version.h>
-
 namespace esphome::ota {
 
 class ArduinoRP2OTABackend final {
@@ -20,9 +18,7 @@ class ArduinoRP2OTABackend final {
   // The core's OTA stub inflates a staged gzip image at reboot, on every chip
   // from 4.0.3 (ESPHome pins 6.0.0). begin() only sees the gzip size; the
   // inflated size is known when the stub reads the trailer.
-  static constexpr bool supports_compression() {
-    return VERSION_CODE(ARDUINO_PICO_MAJOR, ARDUINO_PICO_MINOR, ARDUINO_PICO_REVISION) >= VERSION_CODE(4, 0, 3);
-  }
+  static constexpr bool supports_compression() { return USE_ARDUINO_VERSION_CODE >= VERSION_CODE(4, 0, 3); }
 
  private:
   bool md5_set_{false};

tests/components/esphome/test_ota_inflate.cpp

@@ -306,7 +306,8 @@ TEST(OtaInflate, CorruptStreamsNeverEscapeTheWindow) {
   // Flipped bytes and garbage; the sanitizers check the decoder stays in bounds
   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) {
+  // A coarse, non-aligned stride: neighbouring offsets hit the same paths
+  for (size_t i = 0; i < bad.size(); i += 29) {
     bad[i] ^= 0x5a;
     inflate_all(*s, bad.data(), bad.size(), 1040);
     bad[i] ^= 0x5a;

tests/integration/test_host_ota.py

@@ -122,6 +122,7 @@ class _Device:
     binary_path: Path
     proc: asyncio.subprocess.Process | None = None
     reboots: int = 0
+    inflates: int = 0
 
     def __post_init__(self) -> None:
         self._rebooted = asyncio.Event()
@@ -130,6 +131,8 @@ class _Device:
         if "Rebooting safely" in line:
             self.reboots += 1
             self._rebooted.set()
+        if "Inflated " in line and " bytes from " in line:
+            self.inflates += 1
 
     async def wait_reboot(self, count: int, timeout: float = 10.0) -> None:
         async with asyncio.timeout(timeout):
@@ -180,14 +183,10 @@ async def test_host_ota_self_update(
         )
     )
     staged = asyncio.Event()
-    inflated = asyncio.Event()
 
     def on_log(line: str) -> None:
         if "OTA staged at" in line:
             staged.set()
-        # The host backend cannot store gzip, so the upload negotiates deflate
-        if "Inflated " in line and " bytes from " in line:
-            inflated.set()
         dev.on_log(line)
 
     async with run_binary(dev.binary_path, line_callback=on_log) as (proc, _lines):
@@ -199,7 +198,6 @@ async def test_host_ota_self_update(
 
         await dev.ota(None, None, "espota2 reported failure")
         assert staged.is_set()
-        assert inflated.is_set(), "upload was not deflate compressed"
 
         async with wait_and_connect_api_client(port=dev.api_port) as client:
             info_after = await client.device_info()
@@ -224,13 +222,15 @@ async def test_host_ota_deflate(
             yaml_config, write_yaml_config, compile_esphome, reserved_tcp_port
         )
     )
-    inflated: list[str] = []
     errors: list[str] = []
 
     def on_log(line: str) -> None:
-        if "Inflated " in line and " bytes from " in line:
-            inflated.append(line)
-        if "Inflate err" in line or "End update err" in line:
+        # A corrupt stream is caught by the decoder, by the size check or by
+        # the MD5 at the end, depending on where the damage lands
+        if any(
+            text in line
+            for text in ("Inflate err", "Inflate overrun", "End update err")
+        ):
             errors.append(line)
         dev.on_log(line)
 
@@ -247,12 +247,12 @@ async def test_host_ota_deflate(
 
         # Default: the host backend cannot store gzip, so the CLI sends deflate
         await dev.ota(None, None, "deflate upload failed")
-        assert len(inflated) == 1, "device did not inflate the upload"
+        assert dev.inflates == 1, "device did not inflate the upload"
 
         # A client that does not offer deflate is served uncompressed
         monkeypatch.setattr(espota2, "CLIENT_FEATURE_SUPPORTS_DEFLATE", 0)
         await dev.ota(None, None, "uncompressed upload failed")
-        assert len(inflated) == 1, "device inflated without a client offer"
+        assert dev.inflates == 1, "device inflated without a client offer"
         monkeypatch.undo()
 
         # A corrupt stream fails the upload and leaves the device running
@@ -262,9 +262,8 @@ async def test_host_ota_deflate(
         assert errors, "device did not report the corrupt stream"
 
         # and it still takes a good upload afterwards
-        inflated_before = len(inflated)
         await dev.ota(None, None, "upload after a rejected stream failed")
-        assert len(inflated) == inflated_before + 1
+        assert dev.inflates == 2
 
 
 @pytest.mark.asyncio