Merge branch 'ota-extended-protocol' into integration

esphome/components/esphome/ota/ota_esphome.cpp

@@ -212,6 +212,7 @@ void ESPHomeOTAComponent::handle_handshake_() {
       // legacy response.
       this->extended_proto_ = (this->ota_features_ & CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL) != 0;
       if (this->extended_proto_) {
+        static_assert(HANDSHAKE_BUF_SIZE >= 2, "handshake_buf_ must hold the 2-byte extended-protocol feature ack");
         this->handshake_buf_[0] = ota::OTA_RESPONSE_FEATURE_FLAGS;
         this->handshake_buf_[1] = (supports_compression ? SERVER_FEATURE_SUPPORTS_COMPRESSION : 0);
       } else {
@@ -346,6 +347,11 @@ void ESPHomeOTAComponent::handle_data_() {
              (static_cast<size_t>(buf[2]) << 8) | buf[3];
   ESP_LOGV(TAG, "Size is %u bytes", ota_size);
 
+  if (ota_type != ota::OTA_TYPE_UPDATE_APP) {
+    error_code = ota::OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
+    goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
+  }
+
   // Now that we've passed authentication and are actually
   // starting the update, set the warning status and notify
   // listeners. This ensures that port scanners do not
@@ -356,10 +362,6 @@ void ESPHomeOTAComponent::handle_data_() {
   this->notify_state_(ota::OTA_STARTED, 0.0f, 0);
 #endif
 
-  if (ota_type != ota::OTA_TYPE_UPDATE_APP) {
-    error_code = ota::OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
-    goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
-  }
   // This will block for a few seconds as it locks flash
   error_code = this->backend_->begin(ota_size);
   if (error_code != ota::OTA_RESPONSE_OK)

esphome/components/esphome/ota/ota_esphome.h

@@ -91,15 +91,15 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
   std::string password_;
   std::unique_ptr<uint8_t[]> auth_buf_;
 #endif  // USE_OTA_PASSWORD
-  bool extended_proto_{false};
 
   socket::ListenSocket *server_{nullptr};
   std::unique_ptr<socket::Socket> client_;
   ota::OTABackendPtr backend_;
 
   uint32_t client_connect_time_{0};
+  static constexpr size_t HANDSHAKE_BUF_SIZE = 5;
   uint16_t port_;
-  uint8_t handshake_buf_[5];
+  uint8_t handshake_buf_[HANDSHAKE_BUF_SIZE];
   OTAState ota_state_{OTAState::IDLE};
   uint8_t handshake_buf_pos_{0};
   uint8_t ota_features_{0};
@@ -107,6 +107,7 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
   uint8_t auth_buf_pos_{0};
   uint8_t auth_type_{0};  // Store auth type to know which hasher to use
 #endif                    // USE_OTA_PASSWORD
+  bool extended_proto_{false};
 };
 
 }  // namespace esphome

esphome/espota2.py

@@ -285,7 +285,19 @@ def perform_ota(
         features = 0
 
     if ota_type != OTA_TYPE_UPDATE_APP:
-        raise OTAError(f"Unsupported OTA type: 0x{ota_type:02X}")
+        # Any non-app OTA type requires the extended protocol and the
+        # partition-access server feature. Reject up front so the user gets
+        # a clear capability error instead of a post-auth 0x8E from the device.
+        if not extended_proto:
+            raise OTAError(
+                f"Device does not support extended OTA protocol; "
+                f"OTA type 0x{ota_type:02X} requires it"
+            )
+        if not (features & SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS):
+            raise OTAError(
+                f"Device does not support partition access; "
+                f"OTA type 0x{ota_type:02X} cannot be used"
+            )
 
     if features & SERVER_FEATURE_SUPPORTS_COMPRESSION:
         upload_contents = gzip.compress(file_contents, compresslevel=9)

tests/unit_tests/test_espota2.py

@@ -832,20 +832,91 @@ def test_perform_ota_extended_protocol_app(
 
 
 @pytest.mark.usefixtures("mock_time")
-def test_perform_ota_extended_protocol_unsupported_type(
+def test_perform_ota_device_rejects_with_unsupported_ota_type(
     mock_socket: Mock, mock_file: io.BytesIO
 ) -> None:
-    """Test OTA fails when OTA type is unsupported by the client."""
-    # Setup socket responses for recv calls
+    """End-to-end: device returns 0x8E after the size byte; perform_ota must
+    surface the human-readable 'unsupported OTA type' error from the lookup
+    table in check_error()."""
     recv_responses = [
         bytes([espota2.RESPONSE_OK]),  # First byte of version response
         bytes([espota2.OTA_VERSION_2_0]),  # Version number
-        bytes([espota2.RESPONSE_HEADER_OK]),  # Features response
+        bytes([espota2.RESPONSE_FEATURE_FLAGS]),  # Extended protocol marker
+        bytes(
+            [
+                espota2.SERVER_FEATURE_SUPPORTS_COMPRESSION
+                | espota2.SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS
+            ]
+        ),  # Feature flags
+        bytes([espota2.RESPONSE_AUTH_OK]),  # No auth required
+        bytes([espota2.RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE]),  # Reject at size step
     ]
 
     mock_socket.recv.side_effect = recv_responses
 
-    with pytest.raises(espota2.OTAError, match="Unsupported OTA type"):
+    with pytest.raises(
+        espota2.OTAError,
+        match="The requested OTA type is not supported by the device",
+    ):
+        espota2.perform_ota(
+            mock_socket,
+            "testpass",
+            mock_file,
+            "test.bin",
+            espota2.OTA_TYPE_UPDATE_APP,
+        )
+
+    # Verify the client did send the OTA type byte before the size step
+    assert mock_socket.sendall.call_args_list[2] == call(
+        bytes([espota2.OTA_TYPE_UPDATE_APP])
+    )
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_non_app_type_requires_extended_protocol(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Non-app OTA type must fail when device only supports the legacy protocol."""
+    recv_responses = [
+        bytes([espota2.RESPONSE_OK]),  # First byte of version response
+        bytes([espota2.OTA_VERSION_2_0]),  # Version number
+        bytes([espota2.RESPONSE_HEADER_OK]),  # Legacy single-byte feature ack
+    ]
+
+    mock_socket.recv.side_effect = recv_responses
+
+    with pytest.raises(
+        espota2.OTAError, match="Device does not support extended OTA protocol"
+    ):
+        espota2.perform_ota(
+            mock_socket,
+            "testpass",
+            mock_file,
+            "test.bin",
+            255,
+        )
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_non_app_type_requires_partition_access(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Non-app OTA type must fail when device advertises extended protocol but
+    not the partition-access feature."""
+    recv_responses = [
+        bytes([espota2.RESPONSE_OK]),  # First byte of version response
+        bytes([espota2.OTA_VERSION_2_0]),  # Version number
+        bytes([espota2.RESPONSE_FEATURE_FLAGS]),  # Extended protocol marker
+        bytes(
+            [espota2.SERVER_FEATURE_SUPPORTS_COMPRESSION]
+        ),  # Compression only, no partition access
+    ]
+
+    mock_socket.recv.side_effect = recv_responses
+
+    with pytest.raises(
+        espota2.OTAError, match="Device does not support partition access"
+    ):
         espota2.perform_ota(
             mock_socket,
             "testpass",