From 51c5500809abbf47ac718ff1b2213f79c8d53ca7 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sun, 3 May 2026 09:15:28 -0500 Subject: [PATCH] [ota] Validate partition-table binary host-side before OTA Read the resolved partition-table file in upload_program before opening a network connection. Reject anything that isn't 0xC00 bytes, doesn't start with ESP_PARTITION_MAGIC, or is missing the MD5 checksum entry, so mistakes (wrong file, swapped --file path) surface as a local error instead of a post-handshake OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY. Includes unit tests covering size, magic, md5-presence, missing-file, and end-to-end upload_program rejection, plus three real partition tables checked in as fixtures (ESPHome build, ESP-IDF Hello-world, esphome_dashboard prebuilt). --- esphome/__main__.py | 56 +++++++++ .../partition_tables/esp_idf_hello_world.bin | Bin 0 -> 3072 bytes .../esphome_dashboard_firmware.bin | Bin 0 -> 3072 bytes .../partition_tables/esphome_default.bin | Bin 0 -> 3072 bytes tests/unit_tests/test_main.py | 110 +++++++++++++++++- 5 files changed, 164 insertions(+), 2 deletions(-) create mode 100644 tests/unit_tests/fixtures/partition_tables/esp_idf_hello_world.bin create mode 100644 tests/unit_tests/fixtures/partition_tables/esphome_dashboard_firmware.bin create mode 100644 tests/unit_tests/fixtures/partition_tables/esphome_default.bin diff --git a/esphome/__main__.py b/esphome/__main__.py index 9292e98043..c88efd6750 100644 --- a/esphome/__main__.py +++ b/esphome/__main__.py @@ -1151,9 +1151,65 @@ def upload_program( if getattr(args, "file", None) is not None: binary = Path(args.file) + if ota_type == espota2.OTA_TYPE_UPDATE_PARTITION_TABLE: + _validate_partition_table_binary(binary) + return espota2.run_ota(network_devices, remote_port, password, binary, ota_type) +# Layout of esp_partition_info_t on flash. Each entry is 32 bytes, leading with a +# 16-bit little-endian magic. ESP-IDF defines ESP_PARTITION_MAGIC = 0x50AA (stored as +# bytes 0xAA, 0x50) for partition entries and ESP_PARTITION_MAGIC_MD5 = 0xEBEB for the +# trailing checksum entry. Padding past the last entry is 0xFF. The full table is +# exactly ESP_PARTITION_TABLE_MAX_LEN bytes. +_PARTITION_TABLE_MAX_LEN = 0xC00 +_ESP_PARTITION_MAGIC = 0x50AA +_ESP_PARTITION_MAGIC_MD5 = 0xEBEB + + +def _validate_partition_table_binary(binary: Path) -> None: + """Validate that ``binary`` looks like an ESP32 partition table image. + + Catches common mistakes (wrong file, truncated build output, swapped --file path) + before opening a network connection so the failure mode is a clear local error + instead of a post-handshake device rejection. + """ + try: + data = binary.read_bytes() + except OSError as err: + raise EsphomeError( + f"Cannot read partition table file '{binary}': {err}" + ) from err + + if len(data) != _PARTITION_TABLE_MAX_LEN: + raise EsphomeError( + f"Partition table file '{binary}' has wrong size: expected " + f"{_PARTITION_TABLE_MAX_LEN} bytes, got {len(data)}. " + "Pass the partition table image (e.g. partitions.bin / partition-table.bin), " + "not the firmware image." + ) + + first_magic = data[0] | (data[1] << 8) + if first_magic != _ESP_PARTITION_MAGIC: + raise EsphomeError( + f"Partition table file '{binary}' does not start with the expected " + f"partition magic 0x{_ESP_PARTITION_MAGIC:04X} (got 0x{first_magic:04X}). " + "This file does not look like an ESP32 partition table." + ) + + # The MD5 checksum entry is required: without it the device-side + # esp_partition_table_verify will accept the table but the bootloader will + # refuse to boot from it. Scan the 32-byte entries for the MD5 magic. + if not any( + (data[off] | (data[off + 1] << 8)) == _ESP_PARTITION_MAGIC_MD5 + for off in range(0, _PARTITION_TABLE_MAX_LEN, 32) + ): + raise EsphomeError( + f"Partition table file '{binary}' is missing the MD5 checksum entry. " + "Regenerate the partition table with gen_esp32part.py or rebuild the project." + ) + + def show_logs(config: ConfigType, args: ArgsProtocol, devices: list[str]) -> int | None: try: module = importlib.import_module("esphome.components." + CORE.target_platform) diff --git a/tests/unit_tests/fixtures/partition_tables/esp_idf_hello_world.bin b/tests/unit_tests/fixtures/partition_tables/esp_idf_hello_world.bin new file mode 100644 index 0000000000000000000000000000000000000000..b8fa03b4b3536b1f4d0def4c1fed550e8fc2acc7 GIT binary patch literal 3072 zcmZ1#z{tcffq{V`fq@~fte62EtO{UcWca|qz#zcDP>@j>pP83gf~;m$0Eov3R*;sM zT#{c@2@-(g*RTJhfG=zPT`j`AV@pi8>6C7ps)8ap${7uT(GVC7fzc2c4S~@R7!85Z O5Eu=C(GZ|%2mk;_=9q5) literal 0 HcmV?d00001 diff --git a/tests/unit_tests/fixtures/partition_tables/esphome_dashboard_firmware.bin b/tests/unit_tests/fixtures/partition_tables/esphome_dashboard_firmware.bin new file mode 100644 index 0000000000000000000000000000000000000000..e648fa32709414410fbd9be592c84659b7ebb3ad GIT binary patch literal 3072 zcmZ1#z{tcffq{V`fPo>ete62EtO{UcV0gg5z@WgukYAFRl30?6qGVM7g8%~qBLf42 zBtv3BfdPsn0|UdV00u#@W{A8Yraa?J1_nz8kSVFD1x5L}s47+kFg7s=STZntU|=XN z$V^K^bK>jQ|51SEf~h|;7*|G0zA!t#)$$w%isUF~Gz3ONU^E0qLtr!nMnhmU1V%$( KGz3O?2mk;I)@j>pP83gf+`P^VPs%n zkYPwHC@?^l1F=^HFbFa*$in0eL1M^wRRALs1A`?40|PrlURg1+6qx<`^?wv_Dy30A wx{o#F{@PPlP82`;j3PP884ZEa5Eu=C(GVC7fzc2c4S~@R7!85Z5WpJ(0C1>a2LJ#7 literal 0 HcmV?d00001 diff --git a/tests/unit_tests/test_main.py b/tests/unit_tests/test_main.py index 5f80a4217b..e5564b6933 100644 --- a/tests/unit_tests/test_main.py +++ b/tests/unit_tests/test_main.py @@ -24,6 +24,7 @@ from esphome.__main__ import ( _get_configured_xtal_freq, _make_crystal_freq_callback, _resolve_network_devices, + _validate_partition_table_binary, choose_upload_log_host, command_analyze_memory, command_bundle, @@ -1630,6 +1631,21 @@ def test_upload_program_ota_with_file_arg( ) +_PARTITION_TABLE_LEN = 0xC00 + + +def _make_partition_table_bytes() -> bytes: + """Build a minimal partition table image accepted by _validate_partition_table_binary.""" + table = bytearray(b"\xff" * _PARTITION_TABLE_LEN) + # First entry: ESP_PARTITION_MAGIC (0x50AA) little-endian -> bytes 0xAA, 0x50. + table[0] = 0xAA + table[1] = 0x50 + # MD5 checksum entry at offset 32: ESP_PARTITION_MAGIC_MD5 (0xEBEB) little-endian. + table[32] = 0xEB + table[33] = 0xEB + return bytes(table) + + def test_upload_program_ota_partition_table_with_file_arg( mock_run_ota: Mock, mock_get_port_type: Mock, @@ -1641,6 +1657,9 @@ def test_upload_program_ota_partition_table_with_file_arg( 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: [ { @@ -1650,7 +1669,7 @@ def test_upload_program_ota_partition_table_with_file_arg( } ] } - args = MockArgs(file="partitions.bin", partition_table=True) + args = MockArgs(file=str(partition_file), partition_table=True) devices = ["192.168.1.100"] exit_code, host = upload_program(config, args, devices) @@ -1661,7 +1680,7 @@ def test_upload_program_ota_partition_table_with_file_arg( ["192.168.1.100"], 3232, None, - Path("partitions.bin"), + partition_file, OTA_TYPE_UPDATE_PARTITION_TABLE, ) @@ -1686,6 +1705,93 @@ def test_upload_program_serial_partition_table( upload_program(config, args, devices) +def test_validate_partition_table_binary_accepts_valid(tmp_path: Path) -> None: + f = tmp_path / "partitions.bin" + f.write_bytes(_make_partition_table_bytes()) + _validate_partition_table_binary(f) + + +_PARTITION_FIXTURE_DIR = Path(__file__).parent / "fixtures" / "partition_tables" + + +@pytest.mark.parametrize( + "fixture", + [ + # Stock ESP-IDF gen_esp32part.py output for an ESPHome build. + "esphome_default.bin", + # ESP-IDF Hello-world example partition table (vendored from espressif/esp-serial-flasher). + "esp_idf_hello_world.bin", + # Partition table shipped with esphome_dashboard's prebuilt firmware. + "esphome_dashboard_firmware.bin", + ], +) +def test_validate_partition_table_binary_accepts_real_binaries(fixture: str) -> None: + """Real-world partition-table binaries from ESP-IDF / ESPHome tooling pass validation.""" + _validate_partition_table_binary(_PARTITION_FIXTURE_DIR / fixture) + + +def test_validate_partition_table_binary_rejects_wrong_size(tmp_path: Path) -> None: + f = tmp_path / "partitions.bin" + f.write_bytes(b"\xaa\x50" + b"\xff" * 100) + with pytest.raises(EsphomeError, match="wrong size"): + _validate_partition_table_binary(f) + + +def test_validate_partition_table_binary_rejects_wrong_magic(tmp_path: Path) -> None: + data = bytearray(_make_partition_table_bytes()) + data[0] = 0x00 + data[1] = 0x00 + f = tmp_path / "partitions.bin" + f.write_bytes(bytes(data)) + with pytest.raises(EsphomeError, match="partition magic"): + _validate_partition_table_binary(f) + + +def test_validate_partition_table_binary_rejects_missing_md5(tmp_path: Path) -> None: + data = bytearray(_make_partition_table_bytes()) + data[32] = 0xFF + data[33] = 0xFF + f = tmp_path / "partitions.bin" + f.write_bytes(bytes(data)) + with pytest.raises(EsphomeError, match="missing the MD5 checksum entry"): + _validate_partition_table_binary(f) + + +def test_validate_partition_table_binary_missing_file(tmp_path: Path) -> None: + with pytest.raises(EsphomeError, match="Cannot read partition table file"): + _validate_partition_table_binary(tmp_path / "does-not-exist.bin") + + +def test_upload_program_ota_partition_table_invalid_file( + mock_run_ota: Mock, + mock_get_port_type: Mock, + tmp_path: Path, +) -> None: + """--partition-table must fail before calling run_ota when the file is not a partition table.""" + setup_core(platform=PLATFORM_ESP32, tmp_path=tmp_path) + + mock_get_port_type.return_value = "NETWORK" + + bad_file = tmp_path / "firmware.bin" + bad_file.write_bytes(b"\x00" * 4096) + + config = { + CONF_OTA: [ + { + CONF_PLATFORM: CONF_ESPHOME, + CONF_PORT: 3232, + "allow_partition_access": True, + } + ] + } + args = MockArgs(file=str(bad_file), partition_table=True) + devices = ["192.168.1.100"] + + with pytest.raises(EsphomeError, match="wrong size"): + upload_program(config, args, devices) + mock_run_ota.assert_not_called() + + def test_upload_program_ota_partition_table_without_allow_flag( mock_run_ota: Mock, mock_get_port_type: Mock,