From bf8fca5fbf82890544acfc211d7534660938b8a5 Mon Sep 17 00:00:00 2001 From: Peter M Date: Tue, 9 Jun 2026 09:22:04 +0200 Subject: [PATCH] mix atomvm.esp32.expand expands main.avm partition size, so full flash can be utilised. Signed-off-by: Peter M --- README.md | 18 +++ lib/esp32_image_header.ex | 33 ++++ lib/esp32_partition_table.ex | 229 ++++++++++++++++++++++++++++ lib/esptool_helper.ex | 196 ++++++++++++++++++++++++ lib/mix/tasks/esp32.expand.ex | 222 +++++++++++++++++++++++++++ test/esp32_image_header_test.exs | 24 +++ test/esp32_partition_table_test.exs | 108 +++++++++++++ 7 files changed, 830 insertions(+) create mode 100644 lib/esp32_image_header.ex create mode 100644 lib/esp32_partition_table.ex create mode 100644 lib/mix/tasks/esp32.expand.ex create mode 100644 test/esp32_image_header_test.exs create mode 100644 test/esp32_partition_table_test.exs diff --git a/README.md b/README.md index 1078ea8..8a3ac2a 100644 --- a/README.md +++ b/README.md @@ -364,6 +364,24 @@ Example: Leaving... Hard resetting via RTS pin... +### The `atomvm.esp32.expand` task + +The `atomvm.esp32.expand` task expands a final `main.avm` partition to use the +rest of the ESP32's detected physical flash. It updates the flash size encoded +in the bootloader image header, reads the existing partition table at `0x8000`, +preserves all offsets and other partitions, updates the partition table +checksum, flashes both updates, and verifies them. + +This is useful when an AtomVM image contains a partition table sized for a +smaller flash chip. Detected 8 MB, 16 MB, and 32 MB flash sizes are supported. + + shell$ mix atomvm.esp32.expand --port /dev/tty.usbserial + +Omit `--port` to use the configured AtomVM port or select a connected device +automatically. The task refuses to make changes when another partition follows +`main.avm`, because expanding it would overwrite that partition. It also +refuses to modify devices with secure boot or secure download mode enabled. + ### The `atomvm.stm32.flash` task The `atomvm.stm32.flash` task is used to flash your application to a micro-controller and executed by the AtomVM virtual machine. diff --git a/lib/esp32_image_header.ex b/lib/esp32_image_header.ex new file mode 100644 index 0000000..1a4b8e1 --- /dev/null +++ b/lib/esp32_image_header.ex @@ -0,0 +1,33 @@ +defmodule ExAtomVM.Esp32ImageHeader do + @moduledoc false + + import Bitwise + + @image_magic 0xE9 + @flash_sizes %{ + 0x00 => 1 * 1024 * 1024, + 0x10 => 2 * 1024 * 1024, + 0x20 => 4 * 1024 * 1024, + 0x30 => 8 * 1024 * 1024, + 0x40 => 16 * 1024 * 1024, + 0x50 => 32 * 1024 * 1024, + 0x60 => 64 * 1024 * 1024, + 0x70 => 128 * 1024 * 1024 + } + + def flash_size_id(<<@image_magic, _segments, _mode, size_frequency, _rest::binary>>) do + {:ok, band(size_frequency, 0xF0)} + end + + def flash_size_id(_image), do: {:error, :invalid_image_header} + + def flash_size(image) do + with {:ok, size_id} <- flash_size_id(image), + {:ok, size} <- Map.fetch(@flash_sizes, size_id) do + {:ok, size} + else + :error -> {:error, :unsupported_flash_size} + error -> error + end + end +end diff --git a/lib/esp32_partition_table.ex b/lib/esp32_partition_table.ex new file mode 100644 index 0000000..93ac07b --- /dev/null +++ b/lib/esp32_partition_table.ex @@ -0,0 +1,229 @@ +defmodule ExAtomVM.Esp32PartitionTable do + @moduledoc false + + @entry_size 32 + @data_partition_type 0x01 + @erased_entry :binary.copy(<<0xFF>>, @entry_size) + @md5_prefix <<0xEB, 0xEB>> <> :binary.copy(<<0xFF>>, 14) + + def expand_partition(partition_table, partition_name, flash_size) + when is_binary(partition_table) and is_binary(partition_name) and is_integer(flash_size) do + with {:ok, records, partitions} <- parse_records(partition_table), + {:ok, partition} <- find_partition(partitions, partition_name), + :ok <- validate_partition(partition, partitions, flash_size) do + new_size = flash_size - partition.offset + updated_partition = %{partition | size: new_size} + + {:ok, + %{ + changed?: new_size != partition.size, + partition: partition, + updated_partition: updated_partition, + partition_table: rebuild(records, partition.entry_offset, new_size) + }} + end + end + + def expand_partition(_partition_table, _partition_name, _flash_size) do + {:error, :invalid_arguments} + end + + def parse(partition_table) when is_binary(partition_table) do + with {:ok, _records, partitions} <- parse_records(partition_table) do + {:ok, partitions} + end + end + + defp parse_records(partition_table) do + if rem(byte_size(partition_table), @entry_size) == 0 do + parse_records(partition_table, 0, [], [], []) + else + {:error, :invalid_partition_table} + end + end + + defp parse_records(<<>>, _entry_offset, _segment, _records, _partitions) do + {:error, :invalid_partition_table} + end + + defp parse_records( + <> = remaining, + entry_offset, + segment, + records, + partitions + ) do + cond do + entry == @erased_entry -> + {:ok, Enum.reverse([{:tail, remaining} | records]), Enum.reverse(partitions)} + + md5_entry?(entry) -> + with :ok <- verify_md5(entry, segment) do + parse_records( + rest, + entry_offset + @entry_size, + [], + [{:md5, entry} | records], + partitions + ) + end + + true -> + with {:ok, partition} <- parse_partition(entry, entry_offset) do + parse_records( + rest, + entry_offset + @entry_size, + [entry | segment], + [{:partition, partition, entry} | records], + [partition | partitions] + ) + end + end + end + + defp parse_partition( + <<0xAA, 0x50, type, subtype, offset::little-unsigned-32, size::little-unsigned-32, + label::binary-size(16), flags::little-unsigned-32>>, + entry_offset + ) do + {:ok, + %{ + entry_offset: entry_offset, + flags: flags, + name: decode_label(label), + offset: offset, + size: size, + subtype: subtype, + type: type + }} + end + + defp parse_partition(_entry, _entry_offset), do: {:error, :corrupt_partition_data} + + defp md5_entry?(<<@md5_prefix::binary, _digest::binary-size(16)>>), do: true + defp md5_entry?(_entry), do: false + + defp verify_md5(<<@md5_prefix::binary, digest::binary-size(16)>>, segment) do + if digest == segment_digest(segment) do + :ok + else + {:error, :invalid_partition_table} + end + end + + defp find_partition(partitions, partition_name) do + case Enum.filter(partitions, &(&1.name == partition_name)) do + [partition] -> {:ok, partition} + [] -> {:error, {:partition_not_found, partition_name}} + _partitions -> {:error, {:duplicate_partition, partition_name}} + end + end + + defp validate_partition(partition, partitions, flash_size) do + with :ok <- validate_data_partition(partition), + :ok <- validate_flash_size(flash_size), + :ok <- validate_layout(partitions, flash_size), + :ok <- validate_last_partition(partition, partitions), + :ok <- validate_expansion(partition, flash_size) do + :ok + end + end + + defp validate_data_partition(%{type: @data_partition_type}), do: :ok + + defp validate_data_partition(%{name: name}) do + {:error, {:invalid_partition_type, name}} + end + + defp validate_flash_size(flash_size) when flash_size > 0 and flash_size <= 0xFFFFFFFF, + do: :ok + + defp validate_flash_size(_flash_size), do: {:error, :invalid_flash_size} + + defp validate_layout(partitions, flash_size) do + partitions + |> Enum.sort_by(& &1.offset) + |> Enum.reduce_while(nil, fn partition, previous -> + partition_end = partition.offset + partition.size + + cond do + partition_end > flash_size -> + {:halt, {:error, {:partition_exceeds_flash, partition.name}}} + + previous && previous.offset + previous.size > partition.offset -> + {:halt, {:error, {:overlapping_partitions, previous.name, partition.name}}} + + true -> + {:cont, partition} + end + end) + |> case do + {:error, reason} -> {:error, reason} + _partition -> :ok + end + end + + defp validate_last_partition(partition, partitions) do + case Enum.find(partitions, &(&1.offset > partition.offset)) do + nil -> :ok + next_partition -> {:error, {:partition_not_last, partition.name, next_partition.name}} + end + end + + defp validate_expansion(partition, flash_size) do + current_end = partition.offset + partition.size + + cond do + partition.offset >= flash_size -> + {:error, {:partition_exceeds_flash, partition.name}} + + current_end > flash_size -> + {:error, {:partition_exceeds_flash, partition.name}} + + true -> + :ok + end + end + + defp rebuild(records, target_entry_offset, new_size) do + {iodata, _segment} = + Enum.map_reduce(records, [], fn + {:partition, %{entry_offset: ^target_entry_offset}, entry}, segment -> + updated_entry = update_size(entry, new_size) + {updated_entry, [updated_entry | segment]} + + {:partition, _partition, entry}, segment -> + {entry, [entry | segment]} + + {:md5, _entry}, segment -> + {@md5_prefix <> segment_digest(segment), []} + + {:tail, remaining}, segment -> + {remaining, segment} + end) + + IO.iodata_to_binary(iodata) + end + + defp update_size( + <>, + new_size + ) do + <> + end + + defp segment_digest(segment) do + data = + segment + |> Enum.reverse() + |> IO.iodata_to_binary() + + :crypto.hash(:md5, data) + end + + defp decode_label(label) do + label + |> :binary.split(<<0>>, [:global]) + |> hd() + end +end diff --git a/lib/esptool_helper.ex b/lib/esptool_helper.ex index 8bd3e0d..d54f985 100644 --- a/lib/esptool_helper.ex +++ b/lib/esptool_helper.ex @@ -91,6 +91,202 @@ defmodule ExAtomVM.EsptoolHelper do Pythonx.decode(globals["result"]) end + def read_flash_with_size(port, address, size, reset_after \\ false) do + case (try do + Pythonx.eval( + """ + from esptool.cmds import ( + attach_flash, + detect_chip, + detect_flash_size, + read_flash, + reset_chip, + ) + from esptool.util import flash_size_bytes + + port = port.decode("utf-8") + + with detect_chip(port) as esp: + attach_flash(esp) + try: + flash_size_name = detect_flash_size(esp) + if flash_size_name is None: + raise RuntimeError("Unable to detect flash size") + + data = read_flash( + esp, + address, + size, + None, + flash_size=flash_size_name, + no_progress=True, + ) + result = { + "bootloader_offset": esp.BOOTLOADER_FLASH_OFFSET, + "chip_name": esp.CHIP_NAME, + "data": data, + "flash_size": flash_size_bytes(flash_size_name), + "flash_size_id": esp.parse_flash_size_arg(flash_size_name), + "flash_size_name": flash_size_name, + } + finally: + if reset_after: + reset_chip(esp, "hard-reset") + + if not reset_after: + # Keep USB-OTG devices such as ESP32-S2 in a usable state. + esp.run_stub() + """, + %{ + "address" => address, + "port" => port, + "reset_after" => reset_after, + "size" => size + } + ) + rescue + e in Pythonx.Error -> + {:error, {:pythonx_error, "Pythonx error occurred: #{inspect(e)}"}} + end) do + {_result, %{"result" => result}} -> + {:ok, Pythonx.decode(result)} + + {:error, reason} -> + {:error, reason} + + _ -> + {:error, :flash_read_failed} + end + end + + def write_flash_data(port, address, data) when is_binary(data) do + case (try do + Pythonx.eval( + """ + from esptool.cmds import attach_flash, detect_chip, reset_chip, write_flash + + port = port.decode("utf-8") + + with detect_chip(port) as esp: + attach_flash(esp) + try: + write_flash(esp, [(address, data)], flash_size="keep") + result = True + finally: + reset_chip(esp, "hard-reset") + """, + %{"address" => address, "data" => data, "port" => port} + ) + rescue + e in Pythonx.Error -> + {:error, {:pythonx_error, "Pythonx error occurred: #{inspect(e)}"}} + end) do + {_result, %{"result" => result}} -> + {:ok, Pythonx.decode(result)} + + {:error, reason} -> + {:error, reason} + + _ -> + {:error, :flash_write_failed} + end + end + + def write_flash_size_and_partition( + port, + bootloader_offset, + bootloader, + partition_table_offset, + partition_table, + flash_size_name + ) + when is_binary(bootloader) and is_binary(partition_table) do + case (try do + Pythonx.eval( + """ + from esptool.cmds import ( + _update_image_flash_params, + attach_flash, + detect_chip, + detect_flash_size, + reset_chip, + write_flash, + ) + + port = port.decode("utf-8") + flash_size_name = flash_size_name.decode("utf-8") + + with detect_chip(port) as esp: + attach_flash(esp) + try: + if esp.BOOTLOADER_FLASH_OFFSET != bootloader_offset: + raise RuntimeError( + f"Unexpected bootloader offset {bootloader_offset:#x}; " + f"{esp.CHIP_NAME} uses {esp.BOOTLOADER_FLASH_OFFSET:#x}" + ) + if esp.secure_download_mode or esp.get_secure_boot_enabled(): + raise RuntimeError( + "Cannot update the flash-size header when secure boot " + "or secure download mode is enabled" + ) + + detected_size = detect_flash_size(esp) + if detected_size != flash_size_name: + raise RuntimeError( + f"Flash size changed from {flash_size_name} " + f"to {detected_size or 'unknown'}" + ) + + updated_bootloader = _update_image_flash_params( + esp, + bootloader_offset, + "keep", + "keep", + flash_size_name, + bootloader, + ) + expected_size_id = esp.parse_flash_size_arg(flash_size_name) + if updated_bootloader[0] != esp.ESP_IMAGE_MAGIC: + raise RuntimeError("Invalid bootloader image header") + if updated_bootloader[3] & 0xF0 != expected_size_id: + raise RuntimeError("Failed to update bootloader flash-size header") + + write_flash( + esp, + [ + (bootloader_offset, updated_bootloader), + (partition_table_offset, partition_table), + ], + flash_size="keep", + ) + result = True + finally: + reset_chip(esp, "hard-reset") + """, + %{ + "bootloader" => bootloader, + "bootloader_offset" => bootloader_offset, + "flash_size_name" => flash_size_name, + "partition_table" => partition_table, + "partition_table_offset" => partition_table_offset, + "port" => port + } + ) + rescue + e in Pythonx.Error -> + {:error, {:pythonx_error, "Pythonx error occurred: #{inspect(e)}"}} + end) do + {_result, %{"result" => result}} -> + {:ok, Pythonx.decode(result)} + + {:error, reason} -> + {:error, reason} + + _ -> + {:error, :flash_write_failed} + end + end + @doc """ Erases flash of an ESP32 device. --after "no-reset" is needed for keeping USB-OTG devices like esp32-S2 in a good state. diff --git a/lib/mix/tasks/esp32.expand.ex b/lib/mix/tasks/esp32.expand.ex new file mode 100644 index 0000000..af5f9f1 --- /dev/null +++ b/lib/mix/tasks/esp32.expand.ex @@ -0,0 +1,222 @@ +defmodule Mix.Tasks.Atomvm.Esp32.Expand do + @moduledoc """ + Expands the `main.avm` partition to the end of the detected ESP32 flash. + + The task updates the flash size in the bootloader image header, preserves all + partition offsets and data, and updates the size of a final data partition + named `main.avm` and the partition table checksum. + + ## Options + + * `--port` - Serial port to use. Defaults to the configured AtomVM port, + or automatic device selection when no port is configured. + + ## Example + + mix atomvm.esp32.expand + mix atomvm.esp32.expand --port /dev/tty.usbserial-0001 + """ + + use Mix.Task + + alias ExAtomVM.Esp32ImageHeader + alias ExAtomVM.Esp32PartitionTable + alias ExAtomVM.EsptoolHelper + + @shortdoc "Expand the ESP32 main.avm partition to fill flash" + @partition_name "main.avm" + @partition_table_offset 0x8000 + @partition_table_size 0xC00 + + @impl Mix.Task + def run(args) do + {opts, remaining, invalid} = OptionParser.parse(args, strict: [port: :string]) + + if remaining != [] or invalid != [] do + Mix.raise("Usage: mix atomvm.esp32.expand [--port PORT]") + end + + with :ok <- EsptoolHelper.setup(), + port <- resolve_port(Keyword.get(opts, :port, configured_port())), + {:ok, flash_info} <- + EsptoolHelper.read_flash_with_size( + port, + @partition_table_offset, + @partition_table_size, + true + ), + {:ok, expansion} <- + Esp32PartitionTable.expand_partition( + flash_info["data"], + @partition_name, + flash_info["flash_size"] + ), + {:ok, bootloader_info} <- read_bootloader(port, flash_info), + {:ok, bootloader_flash_size} <- + Esp32ImageHeader.flash_size(bootloader_info["data"]) do + print_expansion(port, flash_info, bootloader_flash_size, expansion) + apply_expansion(port, flash_info, bootloader_info, bootloader_flash_size, expansion) + else + {:error, :pythonx_not_available, message} -> + Mix.raise(message) + + {:error, reason} -> + raise_expand_error(reason) + end + end + + defp configured_port do + Mix.Project.config() + |> Keyword.get(:atomvm, []) + |> Keyword.get(:port, "auto") + end + + defp resolve_port("auto") do + EsptoolHelper.select_device() + |> Map.fetch!("port") + end + + defp resolve_port(port), do: port + + defp read_bootloader(port, flash_info) do + bootloader_offset = flash_info["bootloader_offset"] + bootloader_size = @partition_table_offset - bootloader_offset + + if bootloader_size <= 0 do + {:error, :invalid_bootloader_offset} + else + EsptoolHelper.read_flash_with_size(port, bootloader_offset, bootloader_size, true) + end + end + + defp print_expansion(port, flash_info, bootloader_flash_size, expansion) do + partition = expansion.partition + updated_partition = expansion.updated_partition + + IO.puts(""" + + ESP32: #{flash_info["chip_name"]} on #{port} + Detected flash: #{flash_info["flash_size_name"]} (#{hex(flash_info["flash_size"])}) + Bootloader flash size: #{format_size(bootloader_flash_size)} + #{@partition_name} offset: #{hex(partition.offset)} + Current size: #{format_size(partition.size)} + Expanded size: #{format_size(updated_partition.size)} + """) + end + + defp apply_expansion( + _port, + %{"flash_size" => flash_size}, + _bootloader_info, + flash_size, + %{changed?: false} + ) do + IO.puts("Bootloader and #{@partition_name} already use the detected flash size.") + :ok + end + + defp apply_expansion(port, flash_info, bootloader_info, _bootloader_flash_size, expansion) do + with {:ok, true} <- + EsptoolHelper.write_flash_size_and_partition( + port, + flash_info["bootloader_offset"], + bootloader_info["data"], + @partition_table_offset, + expansion.partition_table, + flash_info["flash_size_name"] + ), + {:ok, bootloader_verification} <- + EsptoolHelper.read_flash_with_size( + port, + flash_info["bootloader_offset"], + 24, + true + ), + {:ok, partition_verification} <- + EsptoolHelper.read_flash_with_size( + port, + @partition_table_offset, + @partition_table_size, + true + ), + {:ok, flash_size_id} <- + Esp32ImageHeader.flash_size_id(bootloader_verification["data"]), + true <- flash_size_id == flash_info["flash_size_id"], + true <- partition_verification["data"] == expansion.partition_table do + IO.puts( + "Updated the bootloader flash size, expanded #{@partition_name}, and verified both." + ) + + :ok + else + {:error, reason} -> raise_expand_error(reason) + false -> Mix.raise("Bootloader or partition table verification failed after flashing.") + other -> Mix.raise("Failed to update the partition table: #{inspect(other)}") + end + end + + defp raise_expand_error(:invalid_partition_table) do + Mix.raise("The ESP32 returned an invalid partition table from flash offset 0x8000.") + end + + defp raise_expand_error(:corrupt_partition_data) do + Mix.raise("The partition table at flash offset 0x8000 contains corrupt data.") + end + + defp raise_expand_error({:partition_not_found, @partition_name}) do + Mix.raise("The device partition table does not contain a #{@partition_name} partition.") + end + + defp raise_expand_error({:duplicate_partition, @partition_name}) do + Mix.raise("The device partition table contains more than one #{@partition_name} partition.") + end + + defp raise_expand_error({:invalid_partition_type, @partition_name}) do + Mix.raise("The #{@partition_name} entry is not a data partition.") + end + + defp raise_expand_error({:partition_not_last, @partition_name, next_partition}) do + Mix.raise(""" + Cannot expand #{@partition_name} because partition #{next_partition} follows it. + Expanding it would overwrite another partition. + """) + end + + defp raise_expand_error({:partition_exceeds_flash, partition}) do + Mix.raise("Partition #{partition} extends beyond the detected physical flash.") + end + + defp raise_expand_error({:overlapping_partitions, first, second}) do + Mix.raise("Partitions #{first} and #{second} overlap; refusing to modify the table.") + end + + defp raise_expand_error(:invalid_flash_size) do + Mix.raise("Esptool returned an invalid physical flash size.") + end + + defp raise_expand_error(:invalid_bootloader_offset) do + Mix.raise("Esptool returned an invalid bootloader offset.") + end + + defp raise_expand_error(:invalid_image_header) do + Mix.raise("The ESP32 bootloader image header is invalid.") + end + + defp raise_expand_error(:unsupported_flash_size) do + Mix.raise("The ESP32 bootloader declares an unsupported flash size.") + end + + defp raise_expand_error({:pythonx_error, message}), do: Mix.raise(message) + defp raise_expand_error(:flash_read_failed), do: Mix.raise("Failed to read ESP32 flash.") + defp raise_expand_error(:flash_write_failed), do: Mix.raise("Failed to write ESP32 flash.") + + defp raise_expand_error(reason) do + Mix.raise("Unable to expand #{@partition_name}: #{inspect(reason)}") + end + + defp format_size(bytes) do + "#{bytes} bytes (#{hex(bytes)})" + end + + defp hex(value), do: "0x" <> String.upcase(Integer.to_string(value, 16)) +end diff --git a/test/esp32_image_header_test.exs b/test/esp32_image_header_test.exs new file mode 100644 index 0000000..71e42c0 --- /dev/null +++ b/test/esp32_image_header_test.exs @@ -0,0 +1,24 @@ +defmodule ExAtomVM.Esp32ImageHeaderTest do + use ExUnit.Case, async: true + + alias ExAtomVM.Esp32ImageHeader + + test "reads flash size without changing the frequency nibble" do + assert {:ok, 0x20} = Esp32ImageHeader.flash_size_id(<<0xE9, 4, 2, 0x2F, 0, 0, 0, 0>>) + assert {:ok, 0x400000} = Esp32ImageHeader.flash_size(<<0xE9, 4, 2, 0x2F>>) + + assert {:ok, 0x30} = Esp32ImageHeader.flash_size_id(<<0xE9, 4, 2, 0x3F, 0, 0, 0, 0>>) + assert {:ok, 0x800000} = Esp32ImageHeader.flash_size(<<0xE9, 4, 2, 0x3F>>) + + assert {:ok, 0x40} = Esp32ImageHeader.flash_size_id(<<0xE9, 4, 2, 0x4F, 0, 0, 0, 0>>) + assert {:ok, 0x1000000} = Esp32ImageHeader.flash_size(<<0xE9, 4, 2, 0x4F>>) + + assert {:ok, 0x50} = Esp32ImageHeader.flash_size_id(<<0xE9, 4, 2, 0x5F, 0, 0, 0, 0>>) + assert {:ok, 0x2000000} = Esp32ImageHeader.flash_size(<<0xE9, 4, 2, 0x5F>>) + end + + test "rejects invalid image headers" do + assert {:error, :invalid_image_header} = Esp32ImageHeader.flash_size(<<0xFF, 0, 0, 0>>) + assert {:error, :invalid_image_header} = Esp32ImageHeader.flash_size(<<0xE9, 0, 0>>) + end +end diff --git a/test/esp32_partition_table_test.exs b/test/esp32_partition_table_test.exs new file mode 100644 index 0000000..286db01 --- /dev/null +++ b/test/esp32_partition_table_test.exs @@ -0,0 +1,108 @@ +defmodule ExAtomVM.Esp32PartitionTableTest do + use ExUnit.Case, async: true + + alias ExAtomVM.Esp32PartitionTable + + test "expands a final main.avm partition to the detected flash size" do + partition_table = + build_partition_table([ + partition("nvs", 0x01, 0x02, 0x9000, 0x4000), + partition("phy_init", 0x01, 0x01, 0xF000, 0x1000), + partition("factory", 0x00, 0x00, 0x10000, 0x1F0000), + partition("boot.avm", 0x01, 0x60, 0x200000, 0x50000), + partition("main.avm", 0x01, 0x60, 0x250000, 0x100000) + ]) + + assert {:ok, expansion} = + Esp32PartitionTable.expand_partition(partition_table, "main.avm", 0x1000000) + + assert expansion.changed? + assert expansion.partition.size == 0x100000 + assert expansion.updated_partition.size == 0xDB0000 + assert byte_size(expansion.partition_table) == 0xC00 + + assert {:ok, partitions} = Esp32PartitionTable.parse(expansion.partition_table) + + assert %{offset: 0x250000, size: 0xDB0000} = + Enum.find(partitions, &(&1.name == "main.avm")) + end + + test "expands main.avm for 8 MB and 32 MB flash" do + partition_table = + build_partition_table([ + partition("factory", 0x00, 0x00, 0x10000, 0x1F0000), + partition("boot.avm", 0x01, 0x60, 0x200000, 0x50000), + partition("main.avm", 0x01, 0x60, 0x250000, 0x100000) + ]) + + assert {:ok, %{updated_partition: %{size: 0x5B0000}}} = + Esp32PartitionTable.expand_partition(partition_table, "main.avm", 0x800000) + + assert {:ok, %{updated_partition: %{size: 0x1DB0000}}} = + Esp32PartitionTable.expand_partition(partition_table, "main.avm", 0x2000000) + end + + test "reports an already expanded partition without changing the table" do + partition_table = + build_partition_table([ + partition("main.avm", 0x01, 0x60, 0x250000, 0xDB0000) + ]) + + assert {:ok, expansion} = + Esp32PartitionTable.expand_partition(partition_table, "main.avm", 0x1000000) + + refute expansion.changed? + assert expansion.partition_table == partition_table + end + + test "refuses to overwrite a partition following main.avm" do + partition_table = + build_partition_table([ + partition("main.avm", 0x01, 0x60, 0x250000, 0x100000), + partition("storage", 0x01, 0x40, 0x350000, 0x100000) + ]) + + assert {:error, {:partition_not_last, "main.avm", "storage"}} = + Esp32PartitionTable.expand_partition(partition_table, "main.avm", 0x1000000) + end + + test "refuses a table that extends beyond physical flash" do + partition_table = + build_partition_table([ + partition("factory", 0x00, 0x00, 0x10000, 0xF0000), + partition("main.avm", 0x01, 0x60, 0x100000, 0x400000) + ]) + + assert {:error, {:partition_exceeds_flash, "main.avm"}} = + Esp32PartitionTable.expand_partition(partition_table, "main.avm", 0x400000) + end + + test "refuses a corrupt checksum" do + partition_table = + build_partition_table([ + partition("main.avm", 0x01, 0x60, 0x250000, 0x100000) + ]) + + <> = partition_table + corrupt_table = head <> <<0>> <> rest + + assert {:error, :invalid_partition_table} = + Esp32PartitionTable.expand_partition(corrupt_table, "main.avm", 0x1000000) + end + + defp build_partition_table(entries) do + data = IO.iodata_to_binary(entries) + md5_entry = <<0xEB, 0xEB>> <> :binary.copy(<<0xFF>>, 14) <> :crypto.hash(:md5, data) + + data <> + md5_entry <> + :binary.copy(<<0xFF>>, 0xC00 - byte_size(data) - byte_size(md5_entry)) + end + + defp partition(name, type, subtype, offset, size) do + label = name <> :binary.copy(<<0>>, 16 - byte_size(name)) + + <<0xAA, 0x50, type, subtype, offset::little-unsigned-32, size::little-unsigned-32, + label::binary-size(16), 0::little-unsigned-32>> + end +end