Replace --erase-otadata with --erase-parts and --erase-data-parts (#273)

maximeborges · web-flow · commit abd8208332e8 · 2022-10-26T11:43:32.000-07:00
* Replace --erase-otadata with --erase-parts and --erase-data-parts

* Update esp-idf-part to v0.1.1
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/cargo-espflash/Cargo.toml b/cargo-espflash/Cargo.toml
@@ -35,6 +35,7 @@ cargo_toml = "0.12.2"
 clap = { version = "4.0.14", features = ["derive"] }
 env_logger = "0.9.1"
 espflash = { version = "=2.0.0-dev", path = "../espflash" }
+esp-idf-part = "0.1.1"
 log = "0.4.17"
 miette = { version = "5.3.0", features = ["fancy"] }
 serde = { version = "1.0.145", features = ["derive"] }
diff --git a/cargo-espflash/src/main.rs b/cargo-espflash/src/main.rs
@@ -1,4 +1,5 @@
 use std::{
+    collections::HashMap,
     fs,
     path::PathBuf,
     process::{exit, Command, ExitStatus, Stdio},
@@ -8,10 +9,11 @@ use cargo_metadata::Message;
 use clap::{Args, Parser, Subcommand};
 use espflash::{
     cli::{
-        self, board_info, clap_enum_variants, config::Config, connect, flash_elf_image,
-        monitor::monitor, partition_table, save_elf_as_image, serial_monitor, ConnectArgs,
-        FlashConfigArgs, PartitionTableArgs,
+        self, board_info, clap_enum_variants, config::Config, connect, erase_partition,
+        flash_elf_image, monitor::monitor, parse_partition_table, partition_table,
+        save_elf_as_image, serial_monitor, ConnectArgs, FlashConfigArgs, PartitionTableArgs,
     },
+    error::{MissingPartition, MissingPartitionTable},
     image_format::ImageFormatKind,
     logging::initialize_logger,
     targets::Chip,
@@ -157,17 +159,11 @@ struct BuildContext {
 }
 
 fn flash(
-    mut args: FlashArgs,
+    args: FlashArgs,
     config: &Config,
     cargo_config: &CargoConfig,
     metadata: &CargoEspFlashMeta,
 ) -> Result<()> {
-    // The `erase_otadata` argument requires `use_stub`, which is implicitly
-    // enabled here.
-    if args.flash_args.erase_otadata {
-        args.connect_args.use_stub = true;
-    }
-
     let mut flasher = connect(&args.connect_args, config)?;
 
     let build_ctx = build(&args.build_args, cargo_config, flasher.chip())
@@ -198,6 +194,54 @@ fn flash(
             .or(metadata.partition_table.as_deref())
             .or(build_ctx.partition_table_path.as_deref());
 
+        let partition_table = match partition_table {
+            Some(path) => Some(parse_partition_table(path)?),
+            None => None,
+        };
+
+        if args.flash_args.erase_parts.is_some() || args.flash_args.erase_data_parts.is_some() {
+            let partition_table = match &partition_table {
+                Some(partition_table) => partition_table,
+                None => return Err((MissingPartitionTable {}).into()),
+            };
+
+            // Using a hashmap to deduplicate entries
+            let mut parts_to_erase = None;
+
+            // Look for any part with specific label
+            if let Some(part_labels) = args.flash_args.erase_parts {
+                for label in part_labels {
+                    let part = partition_table
+                        .find(label.as_str())
+                        .ok_or(MissingPartition::from(label))?;
+                    parts_to_erase
+                        .get_or_insert(HashMap::new())
+                        .insert(part.offset(), part);
+                }
+            }
+            // Look for any data partitions with specific data subtype
+            // There might be multiple partition of the same subtype, e.g. when using multiple FAT partitions
+            if let Some(partition_types) = args.flash_args.erase_data_parts {
+                for ty in partition_types {
+                    for part in partition_table.partitions() {
+                        if part.ty() == esp_idf_part::Type::Data
+                            && part.subtype() == esp_idf_part::SubType::Data(ty)
+                        {
+                            parts_to_erase
+                                .get_or_insert(HashMap::new())
+                                .insert(part.offset(), part);
+                        }
+                    }
+                }
+            }
+
+            if let Some(parts) = parts_to_erase {
+                parts
+                    .iter()
+                    .try_for_each(|(_, p)| erase_partition(&mut flasher, p))?;
+            }
+        }
+
         flash_elf_image(
             &mut flasher,
             &elf_data,
@@ -207,7 +251,6 @@ fn flash(
             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.erase_otadata,
         )?;
     }
 
diff --git a/espflash/Cargo.toml b/espflash/Cargo.toml
@@ -45,7 +45,7 @@ comfy-table = "6.1.0"
 crossterm = { version = "0.25.0", optional = true }
 dialoguer = { version = "0.10.2", optional = true }
 directories-next = "2.0.0"
-esp-idf-part = "0.1.0"
+esp-idf-part = "0.1.1"
 env_logger = { version = "0.9.1", optional = true }
 flate2 = "1.0.24"
 indicatif = "0.17.1"
diff --git a/espflash/src/bin/espflash.rs b/espflash/src/bin/espflash.rs
@@ -1,4 +1,5 @@
 use std::{
+    collections::HashMap,
     fs::{self, File},
     io::Read,
     num::ParseIntError,
@@ -8,10 +9,11 @@ use std::{
 use clap::{Args, Parser, Subcommand};
 use espflash::{
     cli::{
-        self, board_info, clap_enum_variants, config::Config, connect, flash_elf_image,
-        monitor::monitor, partition_table, save_elf_as_image, serial_monitor, ConnectArgs,
-        FlashConfigArgs, PartitionTableArgs,
+        self, board_info, clap_enum_variants, config::Config, connect, erase_partition,
+        flash_elf_image, monitor::monitor, parse_partition_table, partition_table,
+        save_elf_as_image, serial_monitor, ConnectArgs, FlashConfigArgs, PartitionTableArgs,
     },
+    error::{MissingPartition, MissingPartitionTable},
     image_format::ImageFormatKind,
     logging::initialize_logger,
     update::check_for_update,
@@ -114,13 +116,7 @@ fn main() -> Result<()> {
     }
 }
 
-fn flash(mut args: FlashArgs, config: &Config) -> Result<()> {
-    // The `erase_otadata` argument requires `use_stub`, which is implicitly
-    // enabled here.
-    if args.flash_args.erase_otadata {
-        args.connect_args.use_stub = true;
-    }
-
+fn flash(args: FlashArgs, config: &Config) -> Result<()> {
     let mut flasher = connect(&args.connect_args, config)?;
     flasher.board_info()?;
 
@@ -131,7 +127,53 @@ fn flash(mut args: FlashArgs, config: &Config) -> Result<()> {
         flasher.load_elf_to_ram(&elf_data)?;
     } else {
         let bootloader = args.flash_args.bootloader.as_deref();
-        let partition_table = args.flash_args.partition_table.as_deref();
+        let partition_table = match args.flash_args.partition_table.as_deref() {
+            Some(path) => Some(parse_partition_table(path)?),
+            None => None,
+        };
+
+        if args.flash_args.erase_parts.is_some() || args.flash_args.erase_data_parts.is_some() {
+            let partition_table = match &partition_table {
+                Some(partition_table) => partition_table,
+                None => return Err((MissingPartitionTable {}).into()),
+            };
+
+            // Using a hashmap to deduplicate entries
+            let mut parts_to_erase = None;
+
+            // Look for any part with specific label
+            if let Some(part_labels) = args.flash_args.erase_parts {
+                for label in part_labels {
+                    let part = partition_table
+                        .find(label.as_str())
+                        .ok_or(MissingPartition::from(label))?;
+                    parts_to_erase
+                        .get_or_insert(HashMap::new())
+                        .insert(part.offset(), part);
+                }
+            }
+            // Look for any data partitions with specific data subtype
+            // There might be multiple partition of the same subtype, e.g. when using multiple FAT partitions
+            if let Some(partition_types) = args.flash_args.erase_data_parts {
+                for ty in partition_types {
+                    for part in partition_table.partitions() {
+                        if part.ty() == esp_idf_part::Type::Data
+                            && part.subtype() == esp_idf_part::SubType::Data(ty)
+                        {
+                            parts_to_erase
+                                .get_or_insert(HashMap::new())
+                                .insert(part.offset(), part);
+                        }
+                    }
+                }
+            }
+
+            if let Some(parts) = parts_to_erase {
+                parts
+                    .iter()
+                    .try_for_each(|(_, p)| erase_partition(&mut flasher, p))?;
+            }
+        }
 
         flash_elf_image(
             &mut flasher,
@@ -142,7 +184,6 @@ fn flash(mut 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.erase_otadata,
         )?;
     }
 
diff --git a/espflash/src/cli/mod.rs b/espflash/src/cli/mod.rs
@@ -13,9 +13,9 @@ use std::{
     path::{Path, PathBuf},
 };
 
-use clap::Args;
+use clap::{builder::ArgPredicate, Args};
 use comfy_table::{modifiers, presets::UTF8_FULL, Attribute, Cell, Color, Table};
-use esp_idf_part::{DataType, PartitionTable, SubType, Type};
+use esp_idf_part::{DataType, Partition, PartitionTable};
 use log::{debug, info};
 use miette::{IntoDiagnostic, Result, WrapErr};
 use serialport::{SerialPortType, UsbPortInfo};
@@ -24,7 +24,6 @@ use strum::VariantNames;
 use self::{config::Config, monitor::monitor, serial::get_serial_port_info};
 use crate::{
     elf::ElfFirmwareImage,
-    error::{MissingPartitionTable, NoOtadataError},
     flasher::{FlashFrequency, FlashMode, FlashSize, Flasher},
     image_format::ImageFormatKind,
     interface::Interface,
@@ -74,7 +73,7 @@ pub struct ConnectArgs {
     #[cfg_attr(feature = "raspberry", clap(long))]
     pub rts: Option<u8>,
     /// Use RAM stub for loading
-    #[arg(long)]
+    #[arg(long, default_value_ifs([("erase_parts", ArgPredicate::IsPresent, Some("true")), ("erase_data_parts", ArgPredicate::IsPresent, Some("true"))]))]
     pub use_stub: bool,
 }
 
@@ -99,11 +98,17 @@ pub struct FlashArgs {
     /// Path to a binary (.bin) bootloader file
     #[arg(long, value_name = "FILE")]
     pub bootloader: Option<PathBuf>,
-    /// Erase the OTA data partition
-    /// This is useful when using multiple OTA partitions and still wanting to
-    /// be able to reflash via cargo-espflash or espflash
-    #[arg(long)]
-    pub erase_otadata: bool,
+    /// Erase partitions by label
+    #[arg(
+        long,
+        requires = "partition_table",
+        value_name = "LABELS",
+        value_delimiter = ','
+    )]
+    pub erase_parts: Option<Vec<String>>,
+    /// Erase specified data partitions
+    #[arg(long, requires = "partition_table", value_name = "PARTS", value_parser = clap_enum_variants!(DataType), value_delimiter = ',')]
+    pub erase_data_parts: Option<Vec<DataType>>,
     /// Image format to flash
     #[arg(long, value_parser = clap_enum_variants!(ImageFormatKind))]
     pub format: Option<ImageFormatKind>,
@@ -339,12 +344,11 @@ pub fn flash_elf_image(
     flasher: &mut Flasher,
     elf_data: &[u8],
     bootloader: Option<&Path>,
-    partition_table: Option<&Path>,
+    partition_table: Option<PartitionTable>,
     image_format: Option<ImageFormatKind>,
     flash_mode: Option<FlashMode>,
     flash_size: Option<FlashSize>,
     flash_freq: Option<FlashFrequency>,
-    erase_otadata: bool,
 ) -> Result<()> {
     // If the '--bootloader' option is provided, load the binary file at the
     // specified path.
@@ -357,39 +361,6 @@ pub fn flash_elf_image(
         None
     };
 
-    // If the '--partition-table' option is provided, load the partition table from
-    // the CSV or binary file at the specified path.
-    let partition_table = if let Some(path) = partition_table {
-        let path = fs::canonicalize(path).into_diagnostic()?;
-
-        let data = fs::read(path)
-            .into_diagnostic()
-            .wrap_err("Failed to open partition table")?;
-        let table = PartitionTable::try_from(data).into_diagnostic()?;
-
-        Some(table)
-    } else {
-        None
-    };
-
-    if erase_otadata {
-        let partition_table = match &partition_table {
-            Some(partition_table) => partition_table,
-            None => return Err((MissingPartitionTable {}).into()),
-        };
-
-        let otadata =
-            match partition_table.find_by_subtype(Type::Data, SubType::Data(DataType::Ota)) {
-                Some(otadata) => otadata,
-                None => return Err((NoOtadataError {}).into()),
-            };
-
-        let offset = otadata.offset();
-        let size = otadata.size();
-
-        flasher.erase_region(offset, size)?;
-    }
-
     // Load the ELF data, optionally using the provider bootloader/partition
     // table/image format, to the device's flash memory.
     flasher.load_elf_to_flash_with_format(
@@ -406,11 +377,24 @@ pub fn flash_elf_image(
     Ok(())
 }
 
+pub fn parse_partition_table(path: &Path) -> Result<PartitionTable> {
+    let data = fs::read(path)
+        .into_diagnostic()
+        .wrap_err("Failed to open partition table")?;
+    PartitionTable::try_from(data).into_diagnostic()
+}
+
+pub fn erase_partition(flasher: &mut Flasher, part: &Partition) -> Result<()> {
+    log::info!("Erasing {} ({:?})...", part.name(), part.subtype());
+    let offset = part.offset();
+    let size = part.size();
+    flasher.erase_region(offset, size).into_diagnostic()
+}
+
 /// Convert and display CSV and binary partition tables
 pub fn partition_table(args: PartitionTableArgs) -> Result<()> {
     if args.to_binary {
-        let input = fs::read_to_string(&args.partition_table).into_diagnostic()?;
-        let table = PartitionTable::try_from_str(input).into_diagnostic()?;
+        let table = parse_partition_table(&args.partition_table)?;
 
         // Use either stdout or a file if provided for the output.
         let mut writer: Box<dyn Write> = if let Some(output) = args.output {
diff --git a/espflash/src/error.rs b/espflash/src/error.rs
@@ -395,13 +395,18 @@ impl<T> ResultExt for Result<T, Error> {
 }
 
 #[derive(Debug, Error, Diagnostic)]
-#[error("No otadata partition was found")]
+#[error("Missing partition")]
 #[diagnostic(
-    code(espflash::partition_table::no_otadata),
-    help("Partition table must contain an otadata partition when trying to erase it")
+    code(espflash::partition_table::missing_partition),
+    help("Partition table must contain the partition of type `{0}` that is going to be erased")
 )]
+pub struct MissingPartition(String);
 
-pub struct NoOtadataError;
+impl From<String> for MissingPartition {
+    fn from(part: String) -> Self {
+        MissingPartition(part)
+    }
+}
 
 #[derive(Debug, Error, Diagnostic)]
 #[error("Missing partition table")]
