Skip to content

Commit

Permalink
Extract common functionality for erasing partition tables into a func…
Browse files Browse the repository at this point in the history 
…tion
  • Loading branch information
@jessebraham
jessebraham committed Oct 27, 2022
1 parent abd8208 commit d3e3426
Show file tree
Hide file tree
Showing 3 changed files with 74 additions and 87 deletions.
50 changes: 7 additions & 43 deletions cargo-espflash/src/main.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
use std::{
collections::HashMap,
fs,
path::PathBuf,
process::{exit, Command, ExitStatus, Stdio},
Expand All @@ -9,11 +8,10 @@ use cargo_metadata::Message;
use clap::{Args, Parser, Subcommand};
use espflash::{
cli::{
self, board_info, clap_enum_variants, config::Config, connect, erase_partition,
self, board_info, clap_enum_variants, config::Config, connect, erase_partitions,
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,
Expand Down Expand Up @@ -200,46 +198,12 @@ fn flash(
};

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))?;
}
erase_partitions(
&mut flasher,
partition_table.clone(),
args.flash_args.erase_parts,
args.flash_args.erase_data_parts,
)?;
}

flash_elf_image(
Expand Down
50 changes: 7 additions & 43 deletions espflash/src/bin/espflash.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
use std::{
collections::HashMap,
fs::{self, File},
io::Read,
num::ParseIntError,
Expand All @@ -9,11 +8,10 @@ use std::{
use clap::{Args, Parser, Subcommand};
use espflash::{
cli::{
self, board_info, clap_enum_variants, config::Config, connect, erase_partition,
self, board_info, clap_enum_variants, config::Config, connect, erase_partitions,
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,
Expand Down Expand Up @@ -133,46 +131,12 @@ fn flash(args: FlashArgs, config: &Config) -> Result<()> {
};

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))?;
}
erase_partitions(
&mut flasher,
partition_table.clone(),
args.flash_args.erase_parts,
args.flash_args.erase_data_parts,
)?;
}

flash_elf_image(
Expand Down
61 changes: 60 additions & 1 deletion espflash/src/cli/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
//! [espflash]: https://crates.io/crates/espflash

use std::{
collections::HashMap,
fs,
io::Write,
path::{Path, PathBuf},
Expand All @@ -24,6 +25,7 @@ use strum::VariantNames;
use self::{config::Config, monitor::monitor, serial::get_serial_port_info};
use crate::{
elf::ElfFirmwareImage,
error::{MissingPartition, MissingPartitionTable},
flasher::{FlashFrequency, FlashMode, FlashSize, Flasher},
image_format::ImageFormatKind,
interface::Interface,
Expand Down Expand Up @@ -377,17 +379,74 @@ pub fn flash_elf_image(
Ok(())
}

/// Parse a [PartitionTable] from the provided path
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<()> {
/// Erase one or more partitions by label or [DataType]
pub fn erase_partitions(
flasher: &mut Flasher,
partition_table: Option<PartitionTable>,
erase_parts: Option<Vec<String>>,
erase_data_parts: Option<Vec<DataType>>,
) -> Result<()> {
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 partitions with specific labels
if let Some(part_labels) = 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) = 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(flasher, p))?;
}

Ok(())
}

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()
}

Expand Down

0 comments on commit d3e3426

Please sign in to comment.