[ota] Add extended OTA protocol (#16164)

Co-authored-by: J. Nick Koston <nick@home-assistant.io>

tests/unit_tests/test_espota2.py

@@ -185,6 +185,14 @@ def test_receive_exactly_socket_error(mock_socket: Mock) -> None:
             "Error: The OTA partition on the ESP couldn't be found",
         ),
         (espota2.RESPONSE_ERROR_MD5_MISMATCH, "Error: Application MD5 code mismatch"),
+        (
+            espota2.RESPONSE_ERROR_SIGNATURE_INVALID,
+            "Error: Firmware signature verification failed",
+        ),
+        (
+            espota2.RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE,
+            "Error: The requested OTA type is not supported by the device",
+        ),
         (espota2.RESPONSE_ERROR_UNKNOWN, "Unknown error from ESP"),
     ],
 )
@@ -270,12 +278,13 @@ def test_perform_ota_successful_md5_auth(
     # Verify magic bytes were sent
     assert mock_socket.sendall.call_args_list[0] == call(bytes(espota2.MAGIC_BYTES))
 
-    # Verify features were sent (compression + SHA256 support)
+    # Verify features were sent (compression + SHA256 support + extended protocol)
     assert mock_socket.sendall.call_args_list[1] == call(
         bytes(
             [
-                espota2.FEATURE_SUPPORTS_COMPRESSION
-                | espota2.FEATURE_SUPPORTS_SHA256_AUTH
+                espota2.CLIENT_FEATURE_SUPPORTS_COMPRESSION
+                | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
+                | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
             ]
         )
     )
@@ -640,12 +649,13 @@ def test_perform_ota_successful_sha256_auth(
     # Verify magic bytes were sent
     assert mock_socket.sendall.call_args_list[0] == call(bytes(espota2.MAGIC_BYTES))
 
-    # Verify features were sent (compression + SHA256 support)
+    # Verify features were sent (compression + SHA256 support + extended protocol)
     assert mock_socket.sendall.call_args_list[1] == call(
         bytes(
             [
-                espota2.FEATURE_SUPPORTS_COMPRESSION
-                | espota2.FEATURE_SUPPORTS_SHA256_AUTH
+                espota2.CLIENT_FEATURE_SUPPORTS_COMPRESSION
+                | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
+                | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
             ]
         )
     )
@@ -699,8 +709,9 @@ def test_perform_ota_sha256_fallback_to_md5(
     assert mock_socket.sendall.call_args_list[1] == call(
         bytes(
             [
-                espota2.FEATURE_SUPPORTS_COMPRESSION
-                | espota2.FEATURE_SUPPORTS_SHA256_AUTH
+                espota2.CLIENT_FEATURE_SUPPORTS_COMPRESSION
+                | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
+                | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
             ]
         )
     )
@@ -765,3 +776,220 @@ def test_perform_ota_version_differences(
 
     # For v2.0, verify more recv calls due to chunk acknowledgments
     assert mock_socket.recv.call_count == 9  # v2.0 has 9 recv calls (includes chunk OK)
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_extended_protocol_app(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Test OTA extended protocol app update."""
+    recv_responses = [
+        bytes([espota2.RESPONSE_OK]),  # First byte of version response
+        bytes([espota2.OTA_VERSION_2_0]),  # Version number
+        bytes([espota2.RESPONSE_FEATURE_FLAGS]),  # Device supports extended protocol
+        bytes(
+            [
+                espota2.SERVER_FEATURE_SUPPORTS_COMPRESSION
+                | espota2.SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS
+            ]
+        ),  # Device feature flags
+        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
+        bytes([espota2.RESPONSE_CHUNK_OK]),  # Chunk OK
+        bytes([espota2.RESPONSE_RECEIVE_OK]),  # Receive OK
+        bytes([espota2.RESPONSE_UPDATE_END_OK]),  # Update end OK
+    ]
+
+    mock_socket.recv.side_effect = recv_responses
+
+    espota2.perform_ota(
+        mock_socket,
+        "testpass",
+        mock_file,
+        "test.bin",
+        espota2.OTA_TYPE_UPDATE_APP,
+    )
+
+    # Verify magic bytes were sent
+    assert mock_socket.sendall.call_args_list[0] == call(bytes(espota2.MAGIC_BYTES))
+
+    # Verify features were sent (compression + SHA256 support + extended protocol)
+    assert mock_socket.sendall.call_args_list[1] == call(
+        bytes(
+            [
+                espota2.CLIENT_FEATURE_SUPPORTS_COMPRESSION
+                | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
+                | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
+            ]
+        )
+    )
+
+    # Verify ota type was sent
+    assert mock_socket.sendall.call_args_list[2] == call(
+        bytes([espota2.OTA_TYPE_UPDATE_APP])
+    )
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_device_rejects_with_unsupported_ota_type(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """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_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="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_unsupported_type_rejected_early(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """ota_type values not in _SUPPORTED_OTA_TYPES are rejected before any I/O."""
+    with pytest.raises(espota2.OTAError, match="Unsupported OTA type 0xFF"):
+        espota2.perform_ota(
+            mock_socket,
+            "testpass",
+            mock_file,
+            "test.bin",
+            0xFF,
+        )
+    # No bytes should have been transmitted to the device.
+    mock_socket.sendall.assert_not_called()
+
+
+@pytest.mark.parametrize("bad_type", [-1, 256, 0x10000, "app", None, 1.5])
+def test_perform_ota_rejects_out_of_range_type(
+    mock_socket: Mock, mock_file: io.BytesIO, bad_type: object
+) -> None:
+    """Out-of-range or non-int ota_type must raise OTAError, not ValueError."""
+    with pytest.raises(espota2.OTAError, match="Invalid ota_type"):
+        espota2.perform_ota(
+            mock_socket,
+            "testpass",
+            mock_file,
+            "test.bin",
+            bad_type,  # type: ignore[arg-type]
+        )
+    mock_socket.sendall.assert_not_called()
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_non_app_type_requires_extended_protocol(
+    mock_socket: Mock, mock_file: io.BytesIO, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Non-app OTA type must fail when device only supports the legacy protocol."""
+    monkeypatch.setattr(
+        espota2,
+        "_SUPPORTED_OTA_TYPES",
+        frozenset({espota2.OTA_TYPE_UPDATE_APP, 0xFF}),
+    )
+    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",
+            0xFF,
+        )
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_non_app_type_requires_partition_access(
+    mock_socket: Mock, mock_file: io.BytesIO, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Non-app OTA type must fail when device advertises extended protocol but
+    not the partition-access feature."""
+    monkeypatch.setattr(
+        espota2,
+        "_SUPPORTED_OTA_TYPES",
+        frozenset({espota2.OTA_TYPE_UPDATE_APP, 0xFF}),
+    )
+    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",
+            mock_file,
+            "test.bin",
+            0xFF,
+        )
+
+
+def test_check_error_detects_errors_when_expect_is_none() -> None:
+    """check_error must surface device error bytes even when expect is None.
+
+    Regression test: previously, receive_exactly(..., expect=None) calls (used
+    during feature negotiation and nonce reads) silently passed error bytes
+    through, turning clean device errors into confusing later failures.
+    """
+    with pytest.raises(espota2.OTAError, match="Error: Authentication invalid"):
+        espota2.check_error([espota2.RESPONSE_ERROR_AUTH_INVALID], None)
+
+
+def test_check_error_detects_empty_when_expect_is_none() -> None:
+    """Empty data with expect=None must still raise (connection closed)."""
+    with pytest.raises(
+        espota2.OTAError, match="Device closed connection without responding"
+    ):
+        espota2.check_error([], None)
+
+
+def test_check_error_passes_non_error_when_expect_is_none() -> None:
+    """Non-error bytes with expect=None must pass through silently."""
+    espota2.check_error([espota2.RESPONSE_OK], None)
+    espota2.check_error([espota2.RESPONSE_HEADER_OK], None)
+    espota2.check_error([espota2.RESPONSE_FEATURE_FLAGS], None)

tests/unit_tests/test_main.py

@@ -83,6 +83,7 @@ from esphome.const import (
     PLATFORM_RP2040,
 )
 from esphome.core import CORE, EsphomeError
+from esphome.espota2 import OTA_TYPE_UPDATE_APP
 from esphome.util import BootselResult
 from esphome.zeroconf import _await_discovery, discover_mdns_devices
 
@@ -1593,7 +1594,7 @@ def test_upload_program_ota_success(
         tmp_path / ".esphome" / "build" / "test" / ".pioenvs" / "test" / "firmware.bin"
     )
     mock_run_ota.assert_called_once_with(
-        ["192.168.1.100"], 3232, "secret", expected_firmware
+        ["192.168.1.100"], 3232, "secret", expected_firmware, OTA_TYPE_UPDATE_APP
     )
 
 
@@ -1624,7 +1625,7 @@ def test_upload_program_ota_with_file_arg(
     assert exit_code == 0
     assert host == "192.168.1.100"
     mock_run_ota.assert_called_once_with(
-        ["192.168.1.100"], 3232, None, Path("custom.bin")
+        ["192.168.1.100"], 3232, None, Path("custom.bin"), OTA_TYPE_UPDATE_APP
     )
 
 
@@ -1682,7 +1683,7 @@ def test_upload_program_ota_with_mqtt_resolution(
         tmp_path / ".esphome" / "build" / "test" / ".pioenvs" / "test" / "firmware.bin"
     )
     mock_run_ota.assert_called_once_with(
-        ["192.168.1.100"], 3232, None, expected_firmware
+        ["192.168.1.100"], 3232, None, expected_firmware, OTA_TYPE_UPDATE_APP
     )
 
 
@@ -1730,7 +1731,7 @@ def test_upload_program_ota_with_mqtt_empty_broker(
         tmp_path / ".esphome" / "build" / "test" / ".pioenvs" / "test" / "firmware.bin"
     )
     mock_run_ota.assert_called_once_with(
-        ["192.168.1.50"], 3232, None, expected_firmware
+        ["192.168.1.50"], 3232, None, expected_firmware, OTA_TYPE_UPDATE_APP
     )
     # Verify warning was logged
     assert "MQTT IP discovery failed" in caplog.text
@@ -3207,7 +3208,11 @@ def test_upload_program_ota_static_ip_with_mqttip(
         tmp_path / ".esphome" / "build" / "test" / ".pioenvs" / "test" / "firmware.bin"
     )
     mock_run_ota.assert_called_once_with(
-        ["192.168.1.100", "192.168.2.50"], 3232, None, expected_firmware
+        ["192.168.1.100", "192.168.2.50"],
+        3232,
+        None,
+        expected_firmware,
+        OTA_TYPE_UPDATE_APP,
     )
 
 
@@ -3250,7 +3255,11 @@ def test_upload_program_ota_multiple_mqttip_resolves_once(
         tmp_path / ".esphome" / "build" / "test" / ".pioenvs" / "test" / "firmware.bin"
     )
     mock_run_ota.assert_called_once_with(
-        ["192.168.2.50", "192.168.2.51", "192.168.1.100"], 3232, None, expected_firmware
+        ["192.168.2.50", "192.168.2.51", "192.168.1.100"],
+        3232,
+        None,
+        expected_firmware,
+        OTA_TYPE_UPDATE_APP,
     )
 
 
@@ -3415,7 +3424,7 @@ def test_upload_program_ota_mqtt_timeout_fallback(
         tmp_path / ".esphome" / "build" / "test" / ".pioenvs" / "test" / "firmware.bin"
     )
     mock_run_ota.assert_called_once_with(
-        ["192.168.1.100"], 3232, None, expected_firmware
+        ["192.168.1.100"], 3232, None, expected_firmware, OTA_TYPE_UPDATE_APP
     )