Bail from flashing a partition table to too-small flash (#830)

* Don't flash unbootable partition table

* Attempt to test

Author: bugadani

CHANGELOG.md

@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add `--output-format` option to monitor (#818)
 - Added chip detection based on security info, where supported (#814)
 - `espflash` can detect the chip from ESP-HAL metadata to prevent flashing firmware built for a different device. Reqires `esp-hal` 1.0.0-beta.0 (presumably, yet to be released) (#816)
+- `espflash` no longer allows flashing a too-big partition table (#830)
 
 ### Changed
 

espflash/src/error.rs

@@ -228,6 +228,13 @@ pub enum Error {
         help("Ensure that the device is connected and your host recognizes the serial adapter")
     )]
     FirmwareChipMismatch { elf: String, detected: Chip },
+
+    #[error("The partition table does not fit into the flash ({0})")]
+    #[diagnostic(
+        code(espflash::partition_table::does_not_fit),
+        help("Make sure you set the correct flash size with the `--flash-size` option")
+    )]
+    PartitionTableDoesNotFit(FlashSize),
 }
 
 #[cfg(feature = "serialport")]

espflash/src/image_format/esp_idf.rs

@@ -130,6 +130,18 @@ impl<'a> IdfBootloaderFormat<'a> {
             default_partition_table(&params, flash_data.flash_settings.size.map(|v| v.size()))
         });
 
+        if partition_table
+            .partitions()
+            .iter()
+            .map(|p| p.size())
+            .sum::<u32>()
+            > flash_data.flash_settings.size.unwrap_or_default().size()
+        {
+            return Err(Error::PartitionTableDoesNotFit(
+                flash_data.flash_settings.size.unwrap_or_default(),
+            ));
+        }
+
         let mut bootloader = if let Some(bytes) = flash_data.bootloader {
             Cow::Owned(bytes)
         } else {
@@ -406,7 +418,7 @@ fn segment_padding(offset: usize, segment: &Segment<'_>) -> u32 {
     let align_past = (segment.addr - SEG_HEADER_LEN) % IROM_ALIGN;
     let pad_len = ((IROM_ALIGN - ((offset as u32) % IROM_ALIGN)) + align_past) % IROM_ALIGN;
 
-    if pad_len == 0 || pad_len == IROM_ALIGN {
+    if pad_len % IROM_ALIGN == 0 {
         0
     } else if pad_len > SEG_HEADER_LEN {
         pad_len - SEG_HEADER_LEN

espflash/tests/data/partitions.csv

@@ -0,0 +1,4 @@
+# Name,    Type, SubType, Offset,  Size, Flags
+nvs,       data, nvs,     0x9000,  0x6000,
+phy_init,  data, phy,     0xf000,  0x1000,
+factory,   app,  factory, 0x10000, 4M,

espflash/tests/scripts/flash.sh

@@ -1,5 +1,20 @@
 #!/usr/bin/env bash
 app="espflash/tests/data/$1"
+part_table="espflash/tests/data/partitions.csv"
+
+# espflash should not flash a partition table that is too big
+result=$(timeout 15s espflash flash --no-skip --monitor --non-interactive $app --flash-size 2mb --partition-table $part_table 2>&1)
+echo "$result"
+if [[ $result =~ "Flashing has completed!" ]]; then
+    echo "Flashing should have failed!"
+    exit 1
+fi
+if [[ $result =~ "espflash::partition_table::does_not_fit" ]]; then
+    echo "Flashing failed as expected!"
+else
+    echo "Flashing has failed but not with the expected error!"
+    exit 1
+fi
 
 if [[ "$1" == "esp32c6" ]]; then
     # With manual log-format