Add handling for custom partition table offsets (#516)

* read flash locations from partition table

* add partition table offset option

* fix tests

* fix cargo-espflash build

* update changelog

Author: thetek42

CHANGELOG.md

@@ -13,11 +13,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add --target-app-partition argument to flash command (#461)
 - Add --confirm-port argument to flash command (#455)
 - Add --chip argument for flash and write-bin commands (#514)
+- Add --partition-table-offset argument for specifying the partition table offset (#516)
 
 ### Fixed
 
 - Fixed printing panic backtraces when using `esp-println` and `defmt` (#496)
 - Fixed defmt parsing when data is read in parts (#503)
+- Use partition table instead of hard-coded values for the location of partitions (#516)
 
 ### Changed
 

cargo-espflash/src/main.rs

@@ -328,6 +328,7 @@ fn flash(args: FlashArgs, config: &Config) -> Result<()> {
             args.build_args.flash_config_args.flash_mode,
             args.build_args.flash_config_args.flash_size,
             args.build_args.flash_config_args.flash_freq,
+            args.flash_args.partition_table_offset,
         )?;
     }
 
@@ -562,6 +563,7 @@ fn save_image(args: SaveImageArgs) -> Result<()> {
         args.build_args.flash_config_args.flash_mode,
         args.build_args.flash_config_args.flash_size,
         args.build_args.flash_config_args.flash_freq,
+        args.save_image_args.partition_table_offset,
         args.save_image_args.merge,
         bootloader,
         partition_table,

espflash/src/bin/espflash.rs

@@ -250,6 +250,7 @@ fn flash(args: FlashArgs, config: &Config) -> Result<()> {
             args.flash_config_args.flash_mode,
             args.flash_config_args.flash_size,
             args.flash_config_args.flash_freq,
+            args.flash_args.partition_table_offset,
         )?;
     }
 
@@ -305,6 +306,7 @@ fn save_image(args: SaveImageArgs) -> Result<()> {
         args.flash_config_args.flash_mode,
         args.flash_config_args.flash_size,
         args.flash_config_args.flash_freq,
+        args.save_image_args.partition_table_offset,
         args.save_image_args.merge,
         args.save_image_args.bootloader,
         args.save_image_args.partition_table,

espflash/src/cli/mod.rs

@@ -145,6 +145,9 @@ pub struct FlashArgs {
     /// Label of target app partition
     #[arg(long, value_name = "LABEL")]
     pub target_app_partition: Option<String>,
+    /// Partition table offset
+    #[arg(long, value_name = "OFFSET")]
+    pub partition_table_offset: Option<u32>,
     /// Load the application to RAM instead of Flash
     #[arg(long)]
     pub ram: bool,
@@ -185,6 +188,9 @@ pub struct SaveImageArgs {
     /// Custom partition table for merging
     #[arg(long, short = 'T', requires = "merge", value_name = "FILE")]
     pub partition_table: Option<PathBuf>,
+    /// Partition table offset
+    #[arg(long, value_name = "OFFSET")]
+    pub partition_table_offset: Option<u32>,
     /// Label of target app partition
     #[arg(long, value_name = "LABEL")]
     pub target_app_partition: Option<String>,
@@ -333,6 +339,7 @@ pub fn save_elf_as_image(
     flash_mode: Option<FlashMode>,
     flash_size: Option<FlashSize>,
     flash_freq: Option<FlashFrequency>,
+    partition_table_offset: Option<u32>,
     merge: bool,
     bootloader_path: Option<PathBuf>,
     partition_table_path: Option<PathBuf>,
@@ -384,6 +391,7 @@ pub fn save_elf_as_image(
             flash_mode,
             flash_size,
             flash_freq,
+            partition_table_offset,
         )?;
 
         display_image_size(image.app_size(), image.part_size());
@@ -425,6 +433,7 @@ pub fn save_elf_as_image(
             flash_mode,
             flash_size,
             flash_freq,
+            partition_table_offset,
         )?;
 
         display_image_size(image.app_size(), image.part_size());
@@ -538,6 +547,7 @@ pub fn flash_elf_image(
     flash_mode: Option<FlashMode>,
     flash_size: Option<FlashSize>,
     flash_freq: Option<FlashFrequency>,
+    partition_table_offset: Option<u32>,
 ) -> Result<()> {
     // If the '--bootloader' option is provided, load the binary file at the
     // specified path.
@@ -561,6 +571,7 @@ pub fn flash_elf_image(
         flash_mode,
         flash_size,
         flash_freq,
+        partition_table_offset,
         Some(&mut EspflashProgress::default()),
     )?;
     info!("Flashing has completed!");

espflash/src/flasher/mod.rs

@@ -783,6 +783,7 @@ impl Flasher {
         flash_mode: Option<FlashMode>,
         flash_size: Option<FlashSize>,
         flash_freq: Option<FlashFrequency>,
+        partition_table_offset: Option<u32>,
         mut progress: Option<&mut dyn ProgressCallbacks>,
     ) -> Result<(), Error> {
         let image = ElfFirmwareImage::try_from(elf_data)?;
@@ -812,6 +813,7 @@ impl Flasher {
             flash_mode,
             flash_size.or(Some(self.flash_size)),
             flash_freq,
+            partition_table_offset,
         )?;
 
         // When the `cli` feature is enabled, display the image size information.
@@ -869,6 +871,7 @@ impl Flasher {
         flash_mode: Option<FlashMode>,
         flash_size: Option<FlashSize>,
         flash_freq: Option<FlashFrequency>,
+        partition_table_offset: Option<u32>,
         progress: Option<&mut dyn ProgressCallbacks>,
     ) -> Result<(), Error> {
         self.load_elf_to_flash_with_format(
@@ -880,6 +883,7 @@ impl Flasher {
             flash_mode,
             flash_size,
             flash_freq,
+            partition_table_offset,
             progress,
         )
     }

espflash/src/image_format/idf_bootloader.rs

@@ -27,6 +27,7 @@ pub struct IdfBootloaderFormat<'a> {
     flash_segment: RomSegment<'a>,
     app_size: u32,
     part_size: u32,
+    partition_table_offset: u32,
 }
 
 impl<'a> IdfBootloaderFormat<'a> {
@@ -40,6 +41,7 @@ impl<'a> IdfBootloaderFormat<'a> {
         flash_mode: Option<FlashMode>,
         flash_size: Option<FlashSize>,
         flash_freq: Option<FlashFrequency>,
+        partition_table_offset: Option<u32>,
     ) -> Result<Self, Error> {
         let partition_table = partition_table
             .unwrap_or_else(|| params.default_partition_table(flash_size.map(|v| v.size())));
@@ -178,13 +180,25 @@ impl<'a> IdfBootloaderFormat<'a> {
             data: Cow::Owned(data),
         };
 
+        // If the user did not specify a partition offset, we need to assume that the partition
+        // offset is (first partition offset) - 0x1000, since this is the most common case.
+        let partition_table_offset = partition_table_offset.unwrap_or_else(|| {
+            let partitions = partition_table.partitions();
+            let first_partition = partitions
+                .iter()
+                .min_by(|a, b| a.offset().cmp(&b.offset()))
+                .unwrap();
+            first_partition.offset() - 0x1000
+        });
+
         Ok(Self {
             params,
             bootloader,
             partition_table,
             flash_segment,
             app_size,
             part_size,
+            partition_table_offset,
         })
     }
 }
@@ -194,16 +208,39 @@ impl<'a> ImageFormat<'a> for IdfBootloaderFormat<'a> {
     where
         'a: 'b,
     {
+        let bootloader_segment = RomSegment {
+            addr: self.params.boot_addr,
+            data: Cow::Borrowed(&self.bootloader),
+        };
+
+        let partition_table_segment = RomSegment {
+            addr: self.partition_table_offset,
+            data: Cow::Owned(self.partition_table.to_bin().unwrap()),
+        };
+
+        let app_partition = self
+            .partition_table
+            .find("factory")
+            .or_else(|| self.partition_table.find_by_type(Type::App))
+            .expect("no application partition found");
+
+        if self.flash_segment.data.len() > app_partition.size() as usize {
+            panic!(
+                "image size ({} bytes) is larger partition size ({} bytes)",
+                self.flash_segment.data.len(),
+                app_partition.size()
+            );
+        }
+
+        let app_segment = RomSegment {
+            addr: app_partition.offset(),
+            data: Cow::Borrowed(&self.flash_segment.data),
+        };
+
         Box::new(
-            once(RomSegment {
-                addr: self.params.boot_addr,
-                data: Cow::Borrowed(&self.bootloader),
-            })
-            .chain(once(RomSegment {
-                addr: self.params.partition_addr,
-                data: Cow::Owned(self.partition_table.to_bin().unwrap()),
-            }))
-            .chain(once(self.flash_segment.borrow())),
+            once(bootloader_segment)
+                .chain(once(partition_table_segment))
+                .chain(once(app_segment)),
         )
     }
 
@@ -334,6 +371,7 @@ pub mod tests {
             None,
             None,
             None,
+            None,
         )
         .unwrap();
 

espflash/src/targets/esp32.rs

@@ -161,6 +161,7 @@ impl Target for Esp32 {
         flash_mode: Option<FlashMode>,
         flash_size: Option<FlashSize>,
         flash_freq: Option<FlashFrequency>,
+        partition_table_offset: Option<u32>,
     ) -> Result<Box<dyn ImageFormat<'a> + 'a>, Error> {
         let image_format = image_format.unwrap_or(ImageFormatKind::EspBootloader);
 
@@ -175,6 +176,7 @@ impl Target for Esp32 {
                 flash_mode,
                 flash_size,
                 flash_freq,
+                partition_table_offset,
             )?)),
             _ => Err(UnsupportedImageFormatError::new(image_format, Chip::Esp32, None).into()),
         }

espflash/src/targets/esp32c2.rs

@@ -95,6 +95,7 @@ impl Target for Esp32c2 {
         flash_mode: Option<FlashMode>,
         flash_size: Option<FlashSize>,
         flash_freq: Option<FlashFrequency>,
+        partition_table_offset: Option<u32>,
     ) -> Result<Box<dyn ImageFormat<'a> + 'a>, Error> {
         let image_format = image_format.unwrap_or(ImageFormatKind::EspBootloader);
 
@@ -109,6 +110,7 @@ impl Target for Esp32c2 {
                 flash_mode,
                 flash_size,
                 flash_freq,
+                partition_table_offset,
             )?)),
             ImageFormatKind::DirectBoot => Ok(Box::new(DirectBootFormat::new(image, 0)?)),
         }

espflash/src/targets/esp32c3.rs

@@ -82,6 +82,7 @@ impl Target for Esp32c3 {
         flash_mode: Option<FlashMode>,
         flash_size: Option<FlashSize>,
         flash_freq: Option<FlashFrequency>,
+        partition_table_offset: Option<u32>,
     ) -> Result<Box<dyn ImageFormat<'a> + 'a>, Error> {
         let image_format = image_format.unwrap_or(ImageFormatKind::EspBootloader);
 
@@ -96,6 +97,7 @@ impl Target for Esp32c3 {
                 flash_mode,
                 flash_size,
                 flash_freq,
+                partition_table_offset,
             )?)),
             (ImageFormatKind::DirectBoot, None | Some((_, 3..))) => {
                 Ok(Box::new(DirectBootFormat::new(image, 0)?))

espflash/src/targets/esp32c6.rs

@@ -79,6 +79,7 @@ impl Target for Esp32c6 {
         flash_mode: Option<FlashMode>,
         flash_size: Option<FlashSize>,
         flash_freq: Option<FlashFrequency>,
+        partition_table_offset: Option<u32>,
     ) -> Result<Box<dyn ImageFormat<'a> + 'a>, Error> {
         let image_format = image_format.unwrap_or(ImageFormatKind::EspBootloader);
 
@@ -93,6 +94,7 @@ impl Target for Esp32c6 {
                 flash_mode,
                 flash_size,
                 flash_freq,
+                partition_table_offset,
             )?)),
             ImageFormatKind::DirectBoot => Ok(Box::new(DirectBootFormat::new(image, 0x0)?)),
         }