Add support for providing custom partition tables

cargo-espflash/src/main.rs

@@ -1,11 +1,11 @@
 use anyhow::{anyhow, bail, Context};
 use cargo_metadata::Message;
 use clap::{App, Arg, SubCommand};
-use espflash::{Config, Flasher};
+use espflash::{Config, Flasher, PartitionTable};
 use serial::{BaudRate, SerialPort};
 
 use std::{
-    fs::read,
+    fs,
     path::PathBuf,
     process::{exit, Command, ExitStatus, Stdio},
     string::ToString,
@@ -51,6 +51,13 @@ fn main() -> anyhow::Result<()> {
                         .value_name("FEATURES")
                         .help("Comma delimited list of build features"),
                 )
+                .arg(
+                    Arg::with_name("partition_table")
+                        .long("partition-table")
+                        .takes_value(true)
+                        .value_name("PATH")
+                        .help("Path to a CSV file containing partition table"),
+                )
                 .arg(
                     Arg::with_name("speed")
                         .long("speed")
@@ -128,12 +135,26 @@ fn main() -> anyhow::Result<()> {
         return Ok(());
     }
 
+    // If the '--partition-table' option is provided, load the partition table from
+    // the CSV at the specified path.
+    let partition_table = if let Some(path) = matches.value_of("partition_table") {
+        let path = fs::canonicalize(path)?;
+        let data = fs::read_to_string(path)?;
+
+        match PartitionTable::try_from_str(data) {
+            Ok(t) => Some(t),
+            Err(e) => bail!("{}", e),
+        }
+    } else {
+        None
+    };
+
     // Read the ELF data from the build path and load it to the target.
-    let elf_data = read(path.unwrap())?;
+    let elf_data = fs::read(path.unwrap())?;
     if matches.is_present("ram") {
         flasher.load_elf_to_ram(&elf_data)?;
     } else {
-        flasher.load_elf_to_flash(&elf_data)?;
+        flasher.load_elf_to_flash(&elf_data, partition_table)?;
     }
 
     // We're all done!

espflash/Cargo.toml

@@ -30,6 +30,8 @@ directories-next = "2.0.0"
 color-eyre = "0.5"
 strum = "0.21.0"
 strum_macros = "0.21.1"
+csv = "1.1.6"
+regex = "1.5.4"
 
 [dev-dependencies]
 pretty_assertions = "0.7.1"

espflash/src/chip/esp32.rs

@@ -8,8 +8,7 @@ use crate::{
         WP_PIN_DISABLED,
     },
     elf::{FirmwareImage, RomSegment, ESP_CHECKSUM_MAGIC},
-    partition_table::PartitionTable,
-    Error,
+    Error, PartitionTable,
 };
 
 use std::{borrow::Cow, io::Write, iter::once};
@@ -52,18 +51,23 @@ impl ChipType for Esp32 {
 
     fn get_flash_segments<'a>(
         image: &'a FirmwareImage,
+        partition_table: Option<PartitionTable>,
     ) -> Box<dyn Iterator<Item = Result<RomSegment<'a>, Error>> + 'a> {
         let bootloader = include_bytes!("../../bootloader/esp32-bootloader.bin");
 
-        let partition_table = PartitionTable::basic(
-            NVS_ADDR,
-            NVS_SIZE,
-            PHY_INIT_DATA_ADDR,
-            PHY_INIT_DATA_SIZE,
-            APP_ADDR,
-            APP_SIZE,
-        )
-        .to_bytes();
+        let partition_table = if let Some(table) = partition_table {
+            table
+        } else {
+            PartitionTable::basic(
+                NVS_ADDR,
+                NVS_SIZE,
+                PHY_INIT_DATA_ADDR,
+                PHY_INIT_DATA_SIZE,
+                APP_ADDR,
+                APP_SIZE,
+            )
+        };
+        let partition_table = partition_table.to_bytes();
 
         fn get_data<'a>(image: &'a FirmwareImage) -> Result<RomSegment<'a>, Error> {
             let mut data = Vec::new();
@@ -177,7 +181,7 @@ fn test_esp32_rom() {
 
     let image = FirmwareImage::from_data(&input_bytes).unwrap();
 
-    let segments = Esp32::get_flash_segments(&image)
+    let segments = Esp32::get_flash_segments(&image, None)
         .collect::<Result<Vec<_>, Error>>()
         .unwrap();
 

espflash/src/chip/esp32c3.rs

@@ -8,8 +8,7 @@ use crate::{
         WP_PIN_DISABLED,
     },
     elf::{FirmwareImage, RomSegment, ESP_CHECKSUM_MAGIC},
-    partition_table::PartitionTable,
-    Error,
+    Error, PartitionTable,
 };
 
 use std::{borrow::Cow, io::Write, iter::once};
@@ -53,18 +52,23 @@ impl ChipType for Esp32c3 {
 
     fn get_flash_segments<'a>(
         image: &'a FirmwareImage,
+        partition_table: Option<PartitionTable>,
     ) -> Box<dyn Iterator<Item = Result<RomSegment<'a>, Error>> + 'a> {
         let bootloader = include_bytes!("../../bootloader/esp32c3-bootloader.bin");
 
-        let partition_table = PartitionTable::basic(
-            NVS_ADDR,
-            NVS_SIZE,
-            PHY_INIT_DATA_ADDR,
-            PHY_INIT_DATA_SIZE,
-            APP_ADDR,
-            APP_SIZE,
-        )
-        .to_bytes();
+        let partition_table = if let Some(table) = partition_table {
+            table
+        } else {
+            PartitionTable::basic(
+                NVS_ADDR,
+                NVS_SIZE,
+                PHY_INIT_DATA_ADDR,
+                PHY_INIT_DATA_SIZE,
+                APP_ADDR,
+                APP_SIZE,
+            )
+        };
+        let partition_table = partition_table.to_bytes();
 
         fn get_data<'a>(image: &'a FirmwareImage) -> Result<RomSegment<'a>, Error> {
             let mut data = Vec::new();

espflash/src/chip/esp8266.rs

@@ -5,7 +5,7 @@ use crate::{
     chip::{Chip, SpiRegisters},
     elf::{update_checksum, CodeSegment, FirmwareImage, RomSegment, ESP_CHECKSUM_MAGIC},
     flasher::FlashSize,
-    Error,
+    Error, PartitionTable,
 };
 
 use std::{borrow::Cow, io::Write, iter::once, mem::size_of};
@@ -34,6 +34,7 @@ impl ChipType for Esp8266 {
 
     fn get_flash_segments<'a>(
         image: &'a FirmwareImage,
+        _partition_table: Option<PartitionTable>,
     ) -> Box<dyn Iterator<Item = Result<RomSegment<'a>, Error>> + 'a> {
         // irom goes into a separate plain bin
         let irom_data = merge_rom_segments(image.rom_segments(Chip::Esp8266))
@@ -143,7 +144,7 @@ fn test_esp8266_rom() {
 
     let image = FirmwareImage::from_data(&input_bytes).unwrap();
 
-    let segments = Esp8266::get_flash_segments(&image)
+    let segments = Esp8266::get_flash_segments(&image, None)
         .collect::<Result<Vec<_>, Error>>()
         .unwrap();
 

espflash/src/chip/mod.rs

@@ -4,7 +4,7 @@ use strum_macros::Display;
 use crate::{
     elf::{update_checksum, CodeSegment, FirmwareImage, RomSegment},
     flasher::FlashSize,
-    Error,
+    Error, PartitionTable,
 };
 
 use std::{io::Write, str::FromStr};
@@ -29,6 +29,7 @@ pub trait ChipType {
     /// Get the firmware segments for writing an image to flash
     fn get_flash_segments<'a>(
         image: &'a FirmwareImage,
+        partition_table: Option<PartitionTable>,
     ) -> Box<dyn Iterator<Item = Result<RomSegment<'a>, Error>> + 'a>;
 
     fn addr_is_flash(addr: u32) -> bool;
@@ -112,11 +113,12 @@ impl Chip {
     pub fn get_flash_segments<'a>(
         &self,
         image: &'a FirmwareImage,
+        partition_table: Option<PartitionTable>,
     ) -> Box<dyn Iterator<Item = Result<RomSegment<'a>, Error>> + 'a> {
         match self {
-            Chip::Esp32 => Esp32::get_flash_segments(image),
-            Chip::Esp32c3 => Esp32c3::get_flash_segments(image),
-            Chip::Esp8266 => Esp8266::get_flash_segments(image),
+            Chip::Esp32 => Esp32::get_flash_segments(image, partition_table),
+            Chip::Esp32c3 => Esp32c3::get_flash_segments(image, partition_table),
+            Chip::Esp8266 => Esp8266::get_flash_segments(image, None),
         }
     }
 

espflash/src/flasher.rs

@@ -7,7 +7,7 @@ use std::{mem::size_of, thread::sleep};
 
 use crate::{
     chip::Chip, connection::Connection, elf::FirmwareImage, encoder::SlipEncoder, error::RomError,
-    Error,
+    Error, PartitionTable,
 };
 
 type Encoder<'a> = SlipEncoder<'a, Box<dyn SerialPort>>;
@@ -551,12 +551,16 @@ impl Flasher {
     }
 
     /// Load an elf image to flash and execute it
-    pub fn load_elf_to_flash(&mut self, elf_data: &[u8]) -> Result<(), Error> {
+    pub fn load_elf_to_flash(
+        &mut self,
+        elf_data: &[u8],
+        partition_table: Option<PartitionTable>,
+    ) -> Result<(), Error> {
         self.enable_flash(self.spi_params)?;
         let mut image = FirmwareImage::from_data(elf_data).map_err(|_| Error::InvalidElf)?;
         image.flash_size = self.flash_size();
 
-        for segment in self.chip.get_flash_segments(&image) {
+        for segment in self.chip.get_flash_segments(&image, partition_table) {
             let segment = segment?;
             let addr = segment.addr;
             let block_count = (segment.data.len() + FLASH_WRITE_SIZE - 1) / FLASH_WRITE_SIZE;

espflash/src/lib.rs

@@ -11,3 +11,4 @@ pub use chip::Chip;
 pub use config::Config;
 pub use error::Error;
 pub use flasher::Flasher;
+pub use partition_table::PartitionTable;

espflash/src/main.rs

@@ -62,7 +62,7 @@ fn main() -> Result<()> {
     if ram {
         flasher.load_elf_to_ram(&input_bytes)?;
     } else {
-        flasher.load_elf_to_flash(&input_bytes)?;
+        flasher.load_elf_to_flash(&input_bytes, None)?;
     }
 
     Ok(())