From e66434e8b8d2d6a47e4f037146ff657567fca2bd Mon Sep 17 00:00:00 2001 From: Sven-Hendrik Haase Date: Sat, 30 Apr 2022 19:47:01 +0200 Subject: [PATCH] Add partition-table subcommand --- .github/workflows/rust.yml | 2 +- Cargo.lock | 55 ++++++++- espflash/Cargo.toml | 3 +- espflash/src/error.rs | 32 +++++ espflash/src/lib.rs | 2 +- espflash/src/main.rs | 77 +++++++++++- espflash/src/partition_table.rs | 200 ++++++++++++++++++++++++++++---- 7 files changed, 339 insertions(+), 32 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index dc4226a..5554721 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -37,7 +37,7 @@ jobs: - uses: actions-rs/toolchain@v1 with: profile: minimal - toolchain: 1.56 + toolchain: 1.58 override: true - uses: Swatinem/rust-cache@v1 - uses: actions-rs/cargo@v1 diff --git a/Cargo.lock b/Cargo.lock index e78f847..dd20cdd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -253,7 +253,7 @@ version = "3.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a3aab4734e083b809aaf5794e14e756d1c798d2c69c7f7de7a09a2f5214993c1" dependencies = [ - "heck", + "heck 0.4.0", "proc-macro-error", "proc-macro2", "quote", @@ -269,6 +269,18 @@ dependencies = [ "os_str_bytes", ] +[[package]] +name = "comfy-table" +version = "5.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b103d85ca6e209388771bfb7aa6b68a7aeec4afbf6f0a0264bfbf50360e5212e" +dependencies = [ + "crossterm", + "strum 0.23.0", + "strum_macros 0.23.1", + "unicode-width", +] + [[package]] name = "console" version = "0.15.0" @@ -429,6 +441,7 @@ dependencies = [ "binread", "bytemuck", "clap", + "comfy-table", "crossterm", "csv", "dialoguer", @@ -445,8 +458,8 @@ dependencies = [ "serialport", "sha2", "slip-codec", - "strum", - "strum_macros", + "strum 0.24.0", + "strum_macros 0.24.0", "thiserror", "toml", "xmas-elf", @@ -534,6 +547,15 @@ version = "0.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ab5ef0d4909ef3724cc8cce6ccc8572c5c817592e9285f5464f8e86f8bd3726e" +[[package]] +name = "heck" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d621efb26863f0e9924c6ac577e8275e5e6b77455db64ffa6c65c904e9e132c" +dependencies = [ + "unicode-segmentation", +] + [[package]] name = "heck" version = "0.4.0" @@ -1189,19 +1211,38 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" +[[package]] +name = "strum" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cae14b91c7d11c9a851d3fbc80a963198998c2a64eec840477fa92d8ce9b70bb" + [[package]] name = "strum" version = "0.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e96acfc1b70604b8b2f1ffa4c57e59176c7dbb05d556c71ecd2f5498a1dee7f8" +[[package]] +name = "strum_macros" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5bb0dc7ee9c15cea6199cde9a127fa16a4c5819af85395457ad72d68edc85a38" +dependencies = [ + "heck 0.3.3", + "proc-macro2", + "quote", + "rustversion", + "syn", +] + [[package]] name = "strum_macros" version = "0.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6878079b17446e4d3eba6192bb0a2950d5b14f0ed8424b852310e5a94345d0ef" dependencies = [ - "heck", + "heck 0.4.0", "proc-macro2", "quote", "rustversion", @@ -1344,6 +1385,12 @@ dependencies = [ "regex", ] +[[package]] +name = "unicode-segmentation" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e8820f5d777f6224dc4be3632222971ac30164d4a258d595640799554ebfd99" + [[package]] name = "unicode-width" version = "0.1.9" diff --git a/espflash/Cargo.toml b/espflash/Cargo.toml index 80a2203..3674952 100644 --- a/espflash/Cargo.toml +++ b/espflash/Cargo.toml @@ -6,7 +6,7 @@ authors = [ "Jesse Braham ", ] edition = "2021" -rust-version = "1.56" +rust-version = "1.58" description = "A command-line tool for flashing Espressif devices over serial" repository = "https://github.com/esp-rs/espflash" license = "GPL-2.0" @@ -35,6 +35,7 @@ path = "src/main.rs" binread = "2.2" bytemuck = { version = "1.9", features = ["derive"] } clap = { version = "3.1", features = ["derive"] } +comfy-table = "5" crossterm = "0.23" csv = "1.1" dialoguer = "0.10" diff --git a/espflash/src/error.rs b/espflash/src/error.rs index c367087..7d3aa9a 100644 --- a/espflash/src/error.rs +++ b/espflash/src/error.rs @@ -346,6 +346,18 @@ pub enum PartitionTableError { #[error(transparent)] #[diagnostic(transparent)] UnalignedPartitionError(#[from] UnalignedPartitionError), + #[error(transparent)] + #[diagnostic(transparent)] + LengthNotMultipleOf32(#[from] LengthNotMultipleOf32), + #[error(transparent)] + #[diagnostic(transparent)] + InvalidChecksum(#[from] InvalidChecksum), + #[error(transparent)] + #[diagnostic(transparent)] + NoEndMarker(#[from] NoEndMarker), + #[error(transparent)] + #[diagnostic(transparent)] + InvalidPartitionTable(#[from] InvalidPartitionTable), } #[derive(Debug, Error, Diagnostic)] @@ -535,6 +547,26 @@ impl UnalignedPartitionError { } } +#[derive(Debug, Error, Diagnostic)] +#[error("Partition table length not a multiple of 32")] +#[diagnostic(code(espflash::partition_table::invalid_length))] +pub struct LengthNotMultipleOf32; + +#[derive(Debug, Error, Diagnostic)] +#[error("Checksum invalid")] +#[diagnostic(code(espflash::partition_table::invalid_checksum))] +pub struct InvalidChecksum; + +#[derive(Debug, Error, Diagnostic)] +#[error("No end marker found")] +#[diagnostic(code(espflash::partition_table::no_end_marker))] +pub struct NoEndMarker; + +#[derive(Debug, Error, Diagnostic)] +#[error("Invalid partition table")] +#[diagnostic(code(espflash::partition_table::invalid_partition_table))] +pub struct InvalidPartitionTable; + #[derive(Debug, Error)] #[error("{0}")] pub struct ElfError(&'static str); diff --git a/espflash/src/lib.rs b/espflash/src/lib.rs index ed72880..ed19a37 100644 --- a/espflash/src/lib.rs +++ b/espflash/src/lib.rs @@ -1,7 +1,7 @@ pub use chip::Chip; pub use cli::config::Config; pub use elf::{FlashFrequency, FlashMode}; -pub use error::Error; +pub use error::{Error, InvalidPartitionTable}; pub use flasher::{FlashSize, Flasher}; pub use image_format::ImageFormatId; pub use partition_table::PartitionTable; diff --git a/espflash/src/main.rs b/espflash/src/main.rs index 0680f14..b6c4f96 100644 --- a/espflash/src/main.rs +++ b/espflash/src/main.rs @@ -1,4 +1,4 @@ -use std::{fs, mem::swap, path::PathBuf, str::FromStr}; +use std::{fs, io::Write, mem::swap, path::PathBuf, str::FromStr}; use clap::{IntoApp, Parser}; use espflash::{ @@ -6,7 +6,7 @@ use espflash::{ board_info, connect, flash_elf_image, monitor::monitor, save_elf_as_image, ConnectOpts, FlashConfigOpts, FlashOpts, }, - Chip, Config, ImageFormatId, + Chip, Config, ImageFormatId, InvalidPartitionTable, PartitionTable, }; use miette::{IntoDiagnostic, Result, WrapErr}; @@ -34,6 +34,8 @@ pub enum SubCommand { BoardInfo(ConnectOpts), /// Save the image to disk instead of flashing to device SaveImage(SaveImageOpts), + /// Operations for partitions tables + PartitionTable(PartitionTableOpts), } #[derive(Parser)] @@ -60,13 +62,34 @@ pub struct SaveImageOpts { pub partition_table: Option, } +#[derive(Parser)] +pub struct PartitionTableOpts { + /// Convert CSV parition table to binary representation + #[clap(long, required_unless_present_any = ["info", "to-csv"])] + to_binary: bool, + /// Convert binary partition table to CSV representation + #[clap(long, required_unless_present_any = ["info", "to-binary"])] + to_csv: bool, + /// Show information on partition table + #[clap(short, long, required_unless_present_any = ["to-binary", "to-csv"])] + info: bool, + /// Input partition table + partition_table: PathBuf, + /// Optional output file name, if unset will output to stdout + #[clap(short, long)] + output: Option, +} + fn main() -> Result<()> { miette::set_panic_hook(); let mut opts = Opts::parse(); let config = Config::load()?; - if !matches!(opts.subcommand, Some(SubCommand::BoardInfo(..))) { + if !matches!( + opts.subcommand, + Some(SubCommand::BoardInfo(..) | SubCommand::PartitionTable(..)), + ) { // If neither the IMAGE nor SERIAL arguments have been provided, print the // help message and exit. if opts.image.is_none() && opts.connect_opts.serial.is_none() { @@ -89,6 +112,7 @@ fn main() -> Result<()> { match subcommand { BoardInfo(opts) => board_info(opts, config), SaveImage(opts) => save_image(opts), + PartitionTable(opts) => partition_table(opts), } } else { flash(opts, config) @@ -167,3 +191,50 @@ fn save_image(opts: SaveImageOpts) -> Result<()> { Ok(()) } + +fn partition_table(opts: PartitionTableOpts) -> Result<()> { + if opts.to_binary { + let input = fs::read(&opts.partition_table).into_diagnostic()?; + let part_table = PartitionTable::try_from_str(String::from_utf8(input).into_diagnostic()?) + .into_diagnostic()?; + + // Use either stdout or a file if provided for the output. + let mut writer: Box = if let Some(output) = opts.output { + Box::new(fs::File::create(output).into_diagnostic()?) + } else { + Box::new(std::io::stdout()) + }; + part_table.save_bin(&mut writer).into_diagnostic()?; + } else if opts.to_csv { + let input = fs::read(&opts.partition_table).into_diagnostic()?; + let part_table = PartitionTable::try_from_bytes(input).into_diagnostic()?; + + // Use either stdout or a file if provided for the output. + let mut writer: Box = if let Some(output) = opts.output { + Box::new(fs::File::create(output).into_diagnostic()?) + } else { + Box::new(std::io::stdout()) + }; + part_table.save_csv(&mut writer).into_diagnostic()?; + } else if opts.info { + let input = fs::read(&opts.partition_table).into_diagnostic()?; + + // Try getting the partition table from either the csv or the binary representation and + // fail otherwise. + let part_table = if let Ok(part_table) = + PartitionTable::try_from_bytes(input.clone()).into_diagnostic() + { + part_table + } else if let Ok(part_table) = + PartitionTable::try_from_str(String::from_utf8(input).into_diagnostic()?) + { + part_table + } else { + return Err((InvalidPartitionTable {}).into()); + }; + + part_table.pretty_print(); + } + + Ok(()) +} diff --git a/espflash/src/partition_table.rs b/espflash/src/partition_table.rs index 409ea09..2579f1a 100644 --- a/espflash/src/partition_table.rs +++ b/espflash/src/partition_table.rs @@ -1,10 +1,12 @@ use std::{ cmp::{max, min}, fmt::{Display, Formatter, Write as _}, - io::Write, + io::{Cursor, Write}, ops::Rem, }; +use binread::{BinRead, BinReaderExt}; +use comfy_table::{modifiers, presets::UTF8_FULL, Attribute, Cell, Color, Table}; use md5::{Context, Digest}; use regex::Regex; use serde::{Deserialize, Deserializer, Serialize}; @@ -12,17 +14,24 @@ use strum::IntoEnumIterator; use strum_macros::EnumIter; use crate::error::{ - CSVError, DuplicatePartitionsError, InvalidSubTypeError, NoAppError, - OverlappingPartitionsError, PartitionTableError, UnalignedPartitionError, + CSVError, DuplicatePartitionsError, InvalidChecksum, InvalidSubTypeError, + LengthNotMultipleOf32, NoAppError, NoEndMarker, OverlappingPartitionsError, + PartitionTableError, UnalignedPartitionError, }; const MAX_PARTITION_LENGTH: usize = 0xC00; const PARTITION_TABLE_SIZE: usize = 0x1000; const PARTITION_SIZE: usize = 32; const PARTITION_ALIGNMENT: u32 = 0x10000; +const MAGIC_BYTES: &[u8] = &[0xAA, 0x50]; +const MD5_PART_MAGIC_BYTES: &[u8] = &[ + 0xEB, 0xEB, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, +]; +const END_MARKER: [u8; 32] = [0xFF; 32]; -#[derive(Copy, Clone, Debug, Deserialize, Serialize, PartialEq)] +#[derive(Copy, Clone, Debug, Deserialize, Serialize, PartialEq, BinRead)] #[repr(u8)] +#[br(little, repr = u8)] #[serde(rename_all = "lowercase")] pub enum Type { App = 0x00, @@ -53,8 +62,9 @@ impl Display for Type { } } -#[derive(Copy, Clone, Debug, Deserialize, Serialize, PartialEq)] +#[derive(Copy, Clone, Debug, Deserialize, Serialize, PartialEq, BinRead)] #[repr(u8)] +#[br(little, repr = u8)] pub enum AppType { #[serde(rename = "factory")] Factory = 0x00, @@ -94,8 +104,9 @@ pub enum AppType { Test = 0x20, } -#[derive(Copy, Clone, Debug, Deserialize, EnumIter, Serialize, PartialEq)] +#[derive(Copy, Clone, Debug, Deserialize, EnumIter, Serialize, PartialEq, BinRead)] #[repr(u8)] +#[br(little, repr = u8)] #[serde(rename_all = "lowercase")] pub enum DataType { Ota = 0x00, @@ -116,7 +127,7 @@ impl DataType { } } -#[derive(Debug, Deserialize, PartialEq, Copy, Clone)] +#[derive(Debug, Serialize, Deserialize, PartialEq, Copy, Clone, BinRead)] #[serde(untagged)] pub enum SubType { App(AppType), @@ -151,7 +162,9 @@ impl SubType { } } -#[derive(Copy, Clone, Debug, Deserialize, Serialize, PartialEq)] +#[derive(Copy, Clone, Debug, Deserialize, Serialize, PartialEq, BinRead)] +#[repr(u8)] +#[br(little, repr = u8)] #[serde(rename_all = "lowercase")] pub enum Flags { Encrypted = 0x1, @@ -163,7 +176,7 @@ impl Flags { } } -#[derive(Debug)] +#[derive(Debug, PartialEq, Serialize)] pub struct PartitionTable { partitions: Vec, } @@ -206,8 +219,9 @@ impl PartitionTable { } } - /// Attempt to parse a partition table from the given string. For more - /// information on the partition table CSV format see: + /// Attempt to parse a CSV partition table from the given string. + /// + /// For more information on the partition table format see: /// https://docs.espressif.com/projects/esp-idf/en/latest/esp32/api-guides/partition-tables.html pub fn try_from_str(data: S) -> Result where @@ -245,28 +259,65 @@ impl PartitionTable { Ok(table) } + /// Attempt to parse a binary partition table from the given bytes. + /// + /// For more information on the partition table format see: + /// https://docs.espressif.com/projects/esp-idf/en/latest/esp32/api-guides/partition-tables.html + pub fn try_from_bytes(data: S) -> Result + where + S: Into>, + { + let data = data.into(); + if data.len() % 32 != 0 { + return Err(PartitionTableError::LengthNotMultipleOf32( + LengthNotMultipleOf32 {}, + )); + } + let mut md5 = Context::new(); + + let mut partitions = vec![]; + for line in data.chunks_exact(PARTITION_SIZE) { + if line.starts_with(MD5_PART_MAGIC_BYTES) { + // The first 16 bytes are just the marker. The next 16 bytes is the actual md5 + // string. + let digest_in_file = &line[16..32]; + let digest_computed = *md5.clone().compute(); + if digest_computed != digest_in_file { + return Err(PartitionTableError::InvalidChecksum(InvalidChecksum {})); + } + } else if line == END_MARKER { + let table = Self { partitions }; + return Ok(table); + } else { + let mut reader = Cursor::new(line); + let part: Partition = reader.read_le().unwrap(); + partitions.push(part); + md5.consume(line); + } + } + Err(PartitionTableError::NoEndMarker(NoEndMarker {})) + } + pub fn to_bytes(&self) -> Vec { let mut result = Vec::with_capacity(PARTITION_TABLE_SIZE); - self.save(&mut result).unwrap(); + self.save_bin(&mut result).unwrap(); result } - pub fn save(&self, writer: &mut W) -> std::io::Result<()> + /// Write binary form of partition table into `writer`. + pub fn save_bin(&self, writer: &mut W) -> std::io::Result<()> where W: Write, { let mut hasher = HashWriter::new(writer); for partition in &self.partitions { - partition.save(&mut hasher)?; + partition.save_bin(&mut hasher)?; } let (writer, hash) = hasher.compute(); - writer.write_all(&[ - 0xEB, 0xEB, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, - 0xFF, 0xFF, - ])?; + writer.write_all(MD5_PART_MAGIC_BYTES)?; writer.write_all(&hash.0)?; let written = self.partitions.len() * PARTITION_SIZE + 32; @@ -277,6 +328,21 @@ impl PartitionTable { Ok(()) } + /// Write CSV form of partition table into `writer`. + pub fn save_csv(&self, writer: &mut W) -> std::io::Result<()> + where + W: Write, + { + writeln!(writer, "# ESP-IDF Partition Table")?; + writeln!(writer, "# Name, Type, SubType, Offset, Size, Flags")?; + let mut csv = csv::Writer::from_writer(writer); + for partition in &self.partitions { + partition.save_csv(&mut csv)?; + } + + Ok(()) + } + pub fn find(&self, name: &str) -> Option<&Partition> { self.partitions.iter().find(|&p| p.name == name) } @@ -345,6 +411,50 @@ impl PartitionTable { Ok(()) } + + pub fn pretty_print(&self) { + let mut table = Table::new(); + table + .load_preset(UTF8_FULL) + .apply_modifier(modifiers::UTF8_ROUND_CORNERS) + .set_header(vec![ + Cell::new("Name") + .fg(Color::Green) + .add_attribute(Attribute::Bold), + Cell::new("Type") + .fg(Color::Cyan) + .add_attribute(Attribute::Bold), + Cell::new("SubType") + .fg(Color::Magenta) + .add_attribute(Attribute::Bold), + Cell::new("Offset") + .fg(Color::Red) + .add_attribute(Attribute::Bold), + Cell::new("Size") + .fg(Color::Yellow) + .add_attribute(Attribute::Bold), + Cell::new("Flags") + .fg(Color::DarkCyan) + .add_attribute(Attribute::Bold), + ]); + for part in &self.partitions { + table.add_row(vec![ + Cell::new(&part.name).fg(Color::Green), + Cell::new(&part.ty.to_string()).fg(Color::Cyan), + Cell::new(&part.sub_type.to_string()).fg(Color::Magenta), + Cell::new(&format!("{:#x}", part.offset)).fg(Color::Red), + Cell::new(&format!("{:#x} ({}KiB)", part.size, part.size / 1024)).fg(Color::Yellow), + Cell::new( + &part + .flags + .map(|x| format!("{:#x}", x.as_u32())) + .unwrap_or_default(), + ) + .fg(Color::DarkCyan), + ]); + } + println!("{table}"); + } } #[derive(Debug, Deserialize)] @@ -383,17 +493,33 @@ impl DeserializedPartition { } } -#[derive(Debug)] +#[derive(Debug, BinRead, Serialize)] +#[br(magic = b"\xAA\x50", assert(!name.is_empty()))] pub struct Partition { - name: String, ty: Type, sub_type: SubType, offset: u32, size: u32, + #[br(count = 16)] + #[br(map = |s: Vec| String::from_utf8_lossy(&s).trim_matches(char::from(0)).to_string())] + name: String, + #[br(try)] flags: Option, + #[br(ignore)] line: Option, } +impl PartialEq for Partition { + fn eq(&self, other: &Self) -> bool { + self.ty == other.ty + && self.sub_type == other.sub_type + && self.offset == other.offset + && self.size == other.size + && self.name == other.name + && self.flags == other.flags + } +} + impl Partition { pub fn new( name: String, @@ -416,11 +542,11 @@ impl Partition { } } - pub fn save(&self, writer: &mut W) -> std::io::Result<()> + pub fn save_bin(&self, writer: &mut W) -> std::io::Result<()> where W: Write, { - writer.write_all(&[0xAA, 0x50])?; + writer.write_all(MAGIC_BYTES)?; writer.write_all(&[self.ty as u8, self.sub_type.as_u8()])?; writer.write_all(&self.offset.to_le_bytes())?; writer.write_all(&self.size.to_le_bytes())?; @@ -440,6 +566,24 @@ impl Partition { Ok(()) } + pub fn save_csv(&self, csv: &mut csv::Writer) -> std::io::Result<()> + where + W: Write, + { + csv.write_record(&[ + &self.name, + &self.ty.to_string(), + &self.sub_type.to_string(), + &format!("{:#x}", self.offset), + &format!("{:#x}", self.size), + &self + .flags + .map(|x| format!("{:#x}", x.as_u32())) + .unwrap_or_default(), + ])?; + Ok(()) + } + pub fn offset(&self) -> u32 { self.offset } @@ -692,6 +836,18 @@ phy_init, data, phy, 0xf000, 0x1000, .expect_err("Failed to reject partition table without factory or ota partition"); } + #[test] + fn test_from_bytes() { + use std::fs::{read, read_to_string}; + let binary_table = read("./tests/data/partitions.bin").unwrap(); + let binary_parsed = PartitionTable::try_from_bytes(binary_table).unwrap(); + + let csv_table = read_to_string("./tests/data/partitions.csv").unwrap(); + let csv_parsed = PartitionTable::try_from_str(csv_table).unwrap(); + + assert_eq!(binary_parsed, csv_parsed); + } + #[test] fn blank_offsets_are_filled_in() { let pt2 = PartitionTable::try_from_str(PTABLE_2)