Merge branch 'dev' into libretiny-ota-trim-uf2

uf2ota rejects a selector shorter than the three bytes of the scheme map;
reading a short zero-filled one as no device scheme would drop its image
and hide that.

esphome/__main__.py

@@ -1305,10 +1305,7 @@ def upload_program(
                 f"{option_string} is only supported with the esphome OTA platform; "
                 "the web_server OTA path can only update the firmware image."
             )
-        binary = CORE.firmware_bin
-        if getattr(args, "file", None) is not None:
-            binary = Path(args.file)
-        return _upload_via_web_server(config, network_devices, binary)
+        return _upload_via_web_server(config, network_devices, _app_ota_image(args))
 
     return _upload_via_native_api(config, network_devices, args)
 
@@ -1364,6 +1361,16 @@ def _esphome_ota_conf(config: ConfigType) -> ConfigType | None:
     return next(iter(esphome_ota_items(config)), None)
 
 
+def _app_ota_image(args: ArgsProtocol) -> Path:
+    """The file an app update sends: ``--file`` or the build output, as the
+    target platform's ota_image hook picks it."""
+    binary = CORE.firmware_bin
+    if (file := getattr(args, "file", None)) is not None:
+        binary = Path(file)
+    ota_image = platform_hooks.get_platform_hook(CORE.target_platform, "ota_image")
+    return ota_image(binary) if ota_image is not None else binary
+
+
 def _upload_via_native_api(
     config: ConfigType, network_devices: list[str], args: ArgsProtocol
 ) -> tuple[int, str | None]:
@@ -1413,27 +1420,27 @@ def _upload_via_native_api(
                 f"flash a build with the option enabled, and then retry {option_string}."
             )
 
-    binary = CORE.firmware_bin
     ota_type = espota2.OTA_TYPE_UPDATE_APP
+    file = getattr(args, "file", None)
     if getattr(args, "partition_table", False):
         # Fail fast if the resolved ESPHome OTA config does not enable allow_partition_access.
         # The device-side handshake also rejects this with "Device only supports app updates",
         # but checking here surfaces the misconfiguration before opening a network connection.
         check_partition_access("--partition-table")
-        binary = CORE.partition_table_bin
+        binary = Path(file) if file is not None else CORE.partition_table_bin
         ota_type = espota2.OTA_TYPE_UPDATE_PARTITION_TABLE
     elif getattr(args, "bootloader", False):
         check_partition_access("--bootloader")
         if (
-            getattr(args, "file", None) is None
+            file is None
             and (native := native_backend())
             and (hint := getattr(native, "missing_image_hint", lambda: None)())
         ):
             raise EsphomeError(hint)
-        binary = CORE.bootloader_bin
+        binary = Path(file) if file is not None else CORE.bootloader_bin
         ota_type = espota2.OTA_TYPE_UPDATE_BOOTLOADER
-    if getattr(args, "file", None) is not None:
-        binary = Path(args.file)
+    else:
+        binary = _app_ota_image(args)
 
     if ota_type == espota2.OTA_TYPE_UPDATE_PARTITION_TABLE:
         _validate_partition_table_binary(binary)

esphome/components/bk72xx/__init__.py

@@ -78,3 +78,7 @@ async def pin_to_code(config: ConfigType) -> MockObj:
 # Called by writer.py; delegates to the shared libretiny implementation.
 def copy_files() -> None:
     libretiny.copy_files()
+
+
+# Called by the OTA upload; the shared libretiny implementation.
+ota_image = libretiny.device_ota_image

esphome/components/libretiny/__init__.py

@@ -26,7 +26,7 @@ from esphome.const import (
 )
 from esphome.core import CORE
 from esphome.core.config import BOARD_MAX_LENGTH
-from esphome.helpers import copy_file_if_changed
+from esphome.helpers import copy_file_if_changed, write_file
 from esphome.platformio.toolchain import copy_ccache_script
 from esphome.storage_json import StorageJSON
 
@@ -656,6 +656,26 @@ async def component_to_code(config):
     await cg.register_component(var, config)
 
 
+# Called by the OTA upload through the platforms' ota_image hook
+def device_ota_image(firmware: Path) -> Path:
+    """The file to send for an app update: the UF2 without the images the
+    device skips, written next to the build output, or ``firmware`` itself
+    when there is nothing to drop."""
+    from .uf2 import device_ota_uf2
+
+    data = firmware.read_bytes()
+    if (trimmed := device_ota_uf2(data)) is None:
+        return firmware
+    image = CORE.relative_pioenvs_path(CORE.name, "firmware.ota.uf2")
+    write_file(image, trimmed)
+    _LOGGER.info(
+        "Sending the images this device writes: %s of %s bytes",
+        len(trimmed),
+        len(data),
+    )
+    return image
+
+
 # Called by writer.py
 def copy_files() -> None:
     script_dir = Path(__file__).parent

esphome/components/libretiny/generate_components.py

@@ -91,6 +91,10 @@ async def pin_to_code(config: ConfigType) -> MockObj:
 # Called by writer.py; delegates to the shared libretiny implementation.
 def copy_files() -> None:
     libretiny.copy_files()
+
+
+# Called by the OTA upload; the shared libretiny implementation.
+ota_image = libretiny.device_ota_image
 '''
 
 BASE_CODE_BOARDS = '''

esphome/components/libretiny/uf2.py

@@ -0,0 +1,98 @@
+"""Trim a LibreTiny UF2 package to the images a device writes over OTA.
+
+A LibreTiny UF2 holds one image per OTA scheme: the device schemes (single,
+dual slot 1, dual slot 2) and the flasher schemes used over UART. The device
+skips every block of an image whose scheme it does not run (uf2ota), so an
+image no device scheme writes is only sent to be thrown away. On single OTA
+chips (BK72xx, LN882H) that is the whole flasher image, more than half the
+package. The legacy OTA format names the same device partitions, so devices
+that read it write the same images.
+"""
+
+import struct
+
+UF2_BLOCK_SIZE = 512
+UF2_MAGIC_HEAD = struct.pack("<2I", 0x0A324655, 0x9E5D5157)
+UF2_MAGIC_TAIL = struct.pack("<I", 0x0AB16F30)
+UF2_FLAG_HAS_TAGS = 0x8000
+# Header words after the magic: flags, target address, payload length, block
+# number, block count
+_HEADER_WORDS = struct.Struct("<5I")
+_SEQ_OFFSET = 20
+_DATA_START = 32
+_DATA_END = UF2_BLOCK_SIZE - len(UF2_MAGIC_TAIL)
+
+# One partition index nibble per scheme, high nibble first: device single,
+# dual 1 and dual 2 are the first three, the flasher schemes the next three
+UF2_TAG_OTA_PART_INFO = 0xC0EE0C
+# Tags uf2ota acts on beyond the image they sit in; one in an image that
+# would be dropped leaves the package untrimmed
+_KEEP_TAGS = frozenset(
+    {
+        0x5D57D0,  # OTA_FORMAT_1
+        0x6C8492,  # OTA_FORMAT_2
+        0x6EC68A,  # OTA_PART_LIST
+        0xB948DE,  # BINPATCH
+        0x8288ED,  # FAL_PTABLE
+    }
+)
+
+
+def _tags(block: bytes, length: int) -> dict[int, bytes] | None:
+    """The extension tags after a block's payload, or None if malformed."""
+    tags: dict[int, bytes] = {}
+    pos = _DATA_START + length
+    while pos + 4 <= _DATA_END and (size := block[pos]):
+        if size < 4 or pos + size > _DATA_END:
+            return None
+        tags[int.from_bytes(block[pos + 1 : pos + 4], "little")] = block[
+            pos + 4 : pos + size
+        ]
+        pos += (size + 3) & ~3
+    return tags
+
+
+def device_ota_uf2(data: bytes) -> bytes | None:
+    """``data`` without the images no device scheme writes.
+
+    None when there is nothing to drop, or when ``data`` is not a well
+    formed UF2, so the caller sends the file as it is.
+    """
+    if len(data) % UF2_BLOCK_SIZE:
+        return None
+    count = len(data) // UF2_BLOCK_SIZE
+    kept: list[bytes] = []
+    keep = True
+    for seq in range(count):
+        block = data[seq * UF2_BLOCK_SIZE : (seq + 1) * UF2_BLOCK_SIZE]
+        flags, _, length, block_seq, block_count = _HEADER_WORDS.unpack_from(block, 8)
+        # A missing or repeated block must stay visible to the device, which
+        # takes the update size from the block count
+        if (
+            block[:8] != UF2_MAGIC_HEAD
+            or block[_DATA_END:] != UF2_MAGIC_TAIL
+            or _DATA_START + length > _DATA_END
+            or (block_seq, block_count) != (seq, count)
+        ):
+            return None
+        tags = _tags(block, length) if flags & UF2_FLAG_HAS_TAGS else {}
+        if tags is None:
+            return None
+        if (part_info := tags.get(UF2_TAG_OTA_PART_INFO)) is not None:
+            # uf2ota rejects a selector without the whole scheme map
+            if len(part_info) < 3:
+                return None
+            # Any of the three device scheme nibbles names a partition
+            keep = int.from_bytes(part_info[:2], "big") >> 4 != 0
+        if keep:
+            kept.append(block)
+        elif not _KEEP_TAGS.isdisjoint(tags):
+            # Without its image's OTA_PART_INFO the device would write this
+            # block to the partition of the image before it
+            return None
+    if len(kept) == count:
+        return None
+    out = bytearray(b"".join(kept))
+    for seq in range(len(kept)):
+        struct.pack_into("<2I", out, seq * UF2_BLOCK_SIZE + _SEQ_OFFSET, seq, len(kept))
+    return bytes(out)

esphome/components/ln882x/__init__.py

@@ -77,3 +77,7 @@ async def pin_to_code(config: ConfigType) -> MockObj:
 # Called by writer.py; delegates to the shared libretiny implementation.
 def copy_files() -> None:
     libretiny.copy_files()
+
+
+# Called by the OTA upload; the shared libretiny implementation.
+ota_image = libretiny.device_ota_image

esphome/components/rtl87xx/__init__.py

@@ -77,3 +77,7 @@ async def pin_to_code(config: ConfigType) -> MockObj:
 # Called by writer.py; delegates to the shared libretiny implementation.
 def copy_files() -> None:
     libretiny.copy_files()
+
+
+# Called by the OTA upload; the shared libretiny implementation.
+ota_image = libretiny.device_ota_image

esphome/platform_hooks.py

@@ -26,17 +26,20 @@ from esphome.const import (
     PLATFORM_BK72XX,
     PLATFORM_ESP32,
     PLATFORM_ESP8266,
+    PLATFORM_LN882X,
     PLATFORM_NRF52,
     PLATFORM_RP2,
+    PLATFORM_RTL87XX,
     Platform,
 )
 
 _LOGGER = logging.getLogger(__name__)
 
-# Hooks whose loss only degrades diagnostics; skipping one of these is
-# logged at debug, while skipping a hook that changes what the CLI does
-# (upload method, log transport) warns. A new hook is loud by default.
-COSMETIC_HOOKS: Final = frozenset({"process_stacktrace"})
+# Hooks whose loss only degrades diagnostics or upload size; skipping one
+# of these is logged at debug, while skipping a hook that changes what the
+# CLI does (upload method, log transport) warns. A new hook is loud by
+# default.
+COSMETIC_HOOKS: Final = frozenset({"process_stacktrace", "ota_image"})
 
 # Platforms whose firmware prints the crash handler report share one gate.
 _CRASH_HANDLER_GATE = r"0x[0-9a-fA-F]{3,}\b|CRASH DETECTED ON PREVIOUS BOOT"
@@ -71,6 +74,8 @@ PLATFORM_HOOKS: Final[dict[str, frozenset[str]]] = {
     "show_logs": frozenset({PLATFORM_NRF52}),
     "upload_program": frozenset({PLATFORM_NRF52}),
     "process_stacktrace": frozenset(STACKTRACE_GATES),
+    # The file an app update sends, when it is not the build output itself
+    "ota_image": frozenset({PLATFORM_BK72XX, PLATFORM_LN882X, PLATFORM_RTL87XX}),
 }
 
 

tests/unit_tests/components/test_libretiny_uf2.py

@@ -0,0 +1,125 @@
+"""Tests for trimming a LibreTiny UF2 to the images a device writes over OTA."""
+
+from pathlib import Path
+import struct
+
+import pytest
+
+from esphome.components.libretiny import device_ota_image
+from esphome.components.libretiny.uf2 import UF2_TAG_OTA_PART_INFO, device_ota_uf2
+from esphome.core import CORE
+
+# Written from the UF2 and LibreTiny formats, not from the module under test
+MAGIC_1, MAGIC_2, MAGIC_3 = 0x0A324655, 0x9E5D5157, 0x0AB16F30
+FLAG_FAMILY, FLAG_TAGS = 0x2000, 0x8000
+FAMILY = 0x7B3EF230
+TAG_OTA_FORMAT_2 = 0x6C8492
+TAG_BINPATCH = 0xB948DE
+# OTA_PART_INFO index nibbles: device single, dual 1, dual 2, flasher single, dual 1, dual 2
+DEVICE_SINGLE = b"\x10\x00\x00"
+FLASHER_SINGLE = b"\x00\x01\x00"
+DEVICE_DUAL_1 = b"\x01\x00\x00"
+DEVICE_DUAL_2 = b"\x00\x10\x00"
+
+
+def _tag(tag: int, value: bytes) -> bytes:
+    size = 4 + len(value)
+    return bytes([size]) + tag.to_bytes(3, "little") + value + bytes(-size % 4)
+
+
+def _block(payload: bytes = b"", tags: dict[int, bytes] | None = None) -> bytes:
+    """A block numbered 0 of 0; _package numbers it."""
+    flags = FLAG_FAMILY | (FLAG_TAGS if tags else 0)
+    header = struct.pack("<8I", MAGIC_1, MAGIC_2, flags, 0, len(payload), 0, 0, FAMILY)
+    body = payload + b"".join(_tag(t, v) for t, v in (tags or {}).items())
+    return header + body + bytes(512 - 32 - 4 - len(body)) + struct.pack("<I", MAGIC_3)
+
+
+def _image(part_info: bytes, blocks: int, fill: int) -> list[bytes]:
+    """An image: its first block names the partitions, the rest carry data."""
+    data = bytes([fill]) * 256
+    first = _block(data, {UF2_TAG_OTA_PART_INFO: part_info + b"part\x00"})
+    return [first] + [_block(data) for _ in range(blocks - 1)]
+
+
+def _package(*blocks: bytes) -> bytes:
+    """The blocks numbered in order, as a UF2 file."""
+    return b"".join(
+        block[:20] + struct.pack("<2I", seq, len(blocks)) + block[28:]
+        for seq, block in enumerate(blocks)
+    )
+
+
+HEADER = _block(tags={TAG_OTA_FORMAT_2: b"\x02"})
+FLASHER = _image(FLASHER_SINGLE, 5, 0xAA)
+DEVICE = _image(DEVICE_SINGLE, 3, 0xBB)
+SLOT_1 = _image(DEVICE_DUAL_1, 2, 0x11)
+SLOT_2 = _image(DEVICE_DUAL_2, 2, 0x22)
+PATCHED = _block(bytes([0xAA]) * 256, {TAG_BINPATCH: b"\x01\x02\x03\x04"})
+
+
+@pytest.mark.parametrize(
+    ("blocks", "expected"),
+    [
+        pytest.param([HEADER, *FLASHER, *DEVICE], [HEADER, *DEVICE], id="single_ota"),
+        pytest.param(
+            [HEADER, *SLOT_1, *FLASHER, *SLOT_2],
+            [HEADER, *SLOT_1, *SLOT_2],
+            id="dual_ota_keeps_both_slots",
+        ),
+        pytest.param(
+            [HEADER, *FLASHER, PATCHED, *DEVICE],
+            None,
+            id="dropped_image_carries_a_tag_the_device_acts_on",
+        ),
+        pytest.param([HEADER, *DEVICE], None, id="nothing_to_drop"),
+    ],
+)
+def test_device_ota_uf2(blocks: list[bytes], expected: list[bytes] | None) -> None:
+    trimmed = device_ota_uf2(_package(*blocks))
+    assert trimmed == (_package(*expected) if expected is not None else None)
+
+
+def test_data_that_is_not_a_well_formed_uf2_is_left_alone() -> None:
+    package = _package(HEADER, *FLASHER, *DEVICE)
+    bad_magic = package[:512] + b"\x00" + package[513:]
+    # A tag that runs past the data area of the first flasher block
+    at = 512 + 32 + 256
+    bad_tag = package[:at] + bytes([250]) + package[at + 1 :]
+
+    # A whole block missing, or repeated, keeps the original numbering
+    missing = package[:1024] + package[1536:]
+    repeated = package[:1536] + package[1024:]
+    # OTA_PART_INFO too short to hold the scheme map, which uf2ota rejects
+    short_selectors = [
+        _package(HEADER, _block(bytes(256), {UF2_TAG_OTA_PART_INFO: bytes(n)}), *DEVICE)
+        for n in range(3)
+    ]
+
+    for data in (
+        b"",
+        bytes(range(256)) * 4,
+        package[:-1],
+        bad_magic,
+        bad_tag,
+        missing,
+        repeated,
+        *short_selectors,
+    ):
+        assert device_ota_uf2(data) is None
+
+
+def test_device_ota_image(setup_core: Path) -> None:
+    tmp_path = setup_core
+    CORE.build_path = tmp_path / "build"
+    CORE.name = "test"
+    firmware = tmp_path / "firmware.uf2"
+
+    firmware.write_bytes(_package(HEADER, *FLASHER, *DEVICE))
+    image = device_ota_image(firmware)
+    assert image == CORE.relative_pioenvs_path("test", "firmware.ota.uf2")
+    assert image.read_bytes() == _package(HEADER, *DEVICE)
+
+    # Nothing to drop: the firmware itself is sent
+    firmware.write_bytes(_package(HEADER, *DEVICE))
+    assert device_ota_image(firmware) == firmware

tests/unit_tests/test_main.py

@@ -2135,6 +2135,118 @@ def test_upload_program_ota_success(
     )
 
 
+@pytest.mark.parametrize(
+    ("config", "uploader", "image_arg"),
+    [
+        pytest.param(
+            {CONF_OTA: [{CONF_PLATFORM: CONF_ESPHOME, CONF_PORT: 8892}]},
+            "mock_run_ota",
+            3,
+            id="native",
+        ),
+        pytest.param(
+            {
+                CONF_OTA: [{CONF_PLATFORM: CONF_WEB_SERVER}],
+                CONF_WEB_SERVER: {CONF_PORT: 80},
+            },
+            "mock_run_web_server_ota",
+            4,
+            id="web_server",
+        ),
+    ],
+)
+def test_upload_program_ota_sends_the_platform_ota_image(
+    config: ConfigType,
+    uploader: str,
+    image_arg: int,
+    mock_get_port_type: Mock,
+    tmp_path: Path,
+    request: pytest.FixtureRequest,
+) -> None:
+    """An app update sends the file the platform's ota_image hook picks."""
+    setup_core(platform=PLATFORM_BK72XX, tmp_path=tmp_path)
+    mock_get_port_type.return_value = "NETWORK"
+    mock_upload = request.getfixturevalue(uploader)
+    mock_upload.return_value = (0, "192.168.1.100")
+    image = tmp_path / "firmware.ota.uf2"
+
+    with patch(
+        "esphome.components.bk72xx.ota_image", return_value=image
+    ) as mock_ota_image:
+        exit_code, _ = upload_program(config, MockArgs(), ["192.168.1.100"])
+
+    assert exit_code == 0
+    mock_ota_image.assert_called_once_with(CORE.firmware_bin)
+    assert mock_upload.call_args.args[image_arg] == image
+
+
+def test_upload_program_ota_file_goes_through_the_ota_image_hook(
+    mock_run_ota: Mock,
+    mock_get_port_type: Mock,
+    tmp_path: Path,
+) -> None:
+    """--file is the firmware the platform's ota_image hook starts from."""
+    setup_core(platform=PLATFORM_BK72XX, tmp_path=tmp_path)
+    mock_get_port_type.return_value = "NETWORK"
+    mock_run_ota.return_value = (0, "192.168.1.100")
+    other = tmp_path / "other.uf2"
+    image = tmp_path / "firmware.ota.uf2"
+
+    config = {CONF_OTA: [{CONF_PLATFORM: CONF_ESPHOME, CONF_PORT: 8892}]}
+    with patch(
+        "esphome.components.bk72xx.ota_image", return_value=image
+    ) as mock_ota_image:
+        upload_program(config, MockArgs(file=str(other)), ["192.168.1.100"])
+
+    mock_ota_image.assert_called_once_with(other)
+    assert mock_run_ota.call_args.args[3] == image
+
+
+def test_upload_program_ota_partition_table_skips_the_ota_image_hook(
+    mock_run_ota: Mock,
+    mock_get_port_type: Mock,
+    tmp_path: Path,
+) -> None:
+    """Only app updates go through the ota_image hook."""
+    setup_core(platform=PLATFORM_BK72XX, tmp_path=tmp_path)
+    mock_get_port_type.return_value = "NETWORK"
+    mock_run_ota.return_value = (0, "192.168.1.100")
+    partition_file = tmp_path / "partitions.bin"
+    partition_file.write_bytes(_make_partition_table_bytes())
+
+    config = {
+        CONF_OTA: [
+            {
+                CONF_PLATFORM: CONF_ESPHOME,
+                CONF_PORT: 8892,
+                "allow_partition_access": True,
+            }
+        ]
+    }
+    args = MockArgs(file=str(partition_file), partition_table=True)
+    with patch("esphome.components.bk72xx.ota_image") as mock_ota_image:
+        upload_program(config, args, ["192.168.1.100"])
+
+    mock_ota_image.assert_not_called()
+    assert mock_run_ota.call_args.args[3] == partition_file
+
+
+def test_upload_program_ota_without_a_hook_sends_the_firmware(
+    mock_run_ota: Mock,
+    mock_get_port_type: Mock,
+    tmp_path: Path,
+) -> None:
+    """A platform without an ota_image hook sends its build output."""
+    setup_core(platform=PLATFORM_ESP32, tmp_path=tmp_path)
+    mock_get_port_type.return_value = "NETWORK"
+    mock_run_ota.return_value = (0, "192.168.1.100")
+
+    config = {CONF_OTA: [{CONF_PLATFORM: CONF_ESPHOME, CONF_PORT: 3232}]}
+    upload_program(config, MockArgs(), ["192.168.1.100"])
+
+    assert mock_run_ota.call_args.args[3].name == "firmware.bin"
+
+
 @pytest.mark.parametrize(
     "old_key",
     [None, "AQIDBAUGBwgJCgsMDQ4PEBESExQVFhcYGRobHB0eHyA="],