From 79351bd1886bd69e8bde77e2b9a019bbcc297cf9 Mon Sep 17 00:00:00 2001 From: Jesse Braham Date: Tue, 1 Nov 2022 09:45:14 -0700 Subject: [PATCH 1/3] Remove progress bars from flash target modules, make `indicatif` optional --- espflash/Cargo.toml | 6 +-- espflash/src/targets/flash_target/esp32.rs | 47 ++++++++------------ espflash/src/targets/flash_target/esp8266.rs | 24 +++------- 3 files changed, 26 insertions(+), 51 deletions(-) diff --git a/espflash/Cargo.toml b/espflash/Cargo.toml index 158bda4c..30653c9b 100644 --- a/espflash/Cargo.toml +++ b/espflash/Cargo.toml @@ -47,7 +47,7 @@ directories-next = { version = "2.0.0", optional = true } esp-idf-part = "0.1.1" env_logger = { version = "0.9.1", optional = true } flate2 = "1.0.24" -indicatif = "0.17.1" +indicatif = { version = "0.17.1", optional = true } lazy_static = { version = "1.4.0", optional = true } log = "0.4.17" miette = { version = "5.3.0", features = ["fancy"] } @@ -69,7 +69,7 @@ xmas-elf = "0.8.0" default = ["cli"] cli = [ "dep:addr2line", "dep:clap", "dep:comfy-table", "dep:crossterm", "dep:dialoguer", - "dep:directories-next", "dep:env_logger", "dep:lazy_static", "dep:parse_int", - "dep:regex", "dep:serde-hex", "dep:update-informer" + "dep:directories-next", "dep:env_logger", "dep:indicatif", "dep:lazy_static", + "dep:parse_int", "dep:regex", "dep:serde-hex", "dep:update-informer" ] raspberry = ["dep:rppal"] diff --git a/espflash/src/targets/flash_target/esp32.rs b/espflash/src/targets/flash_target/esp32.rs index 2970801b..eb7b1413 100644 --- a/espflash/src/targets/flash_target/esp32.rs +++ b/espflash/src/targets/flash_target/esp32.rs @@ -4,7 +4,6 @@ use flate2::{ write::{ZlibDecoder, ZlibEncoder}, Compression, }; -use indicatif::{ProgressBar, ProgressStyle}; use super::FlashTarget; use crate::{ @@ -36,7 +35,7 @@ impl Esp32Target { impl FlashTarget for Esp32Target { fn begin(&mut self, connection: &mut Connection) -> Result<(), Error> { connection.with_timeout(CommandType::SpiAttach.timeout(), |connection| { - connection.command(if self.use_stub { + let command = if self.use_stub { Command::SpiAttachStub { spi_params: self.spi_attach_params, } @@ -44,26 +43,29 @@ impl FlashTarget for Esp32Target { Command::SpiAttach { spi_params: self.spi_attach_params, } - }) + }; + + connection.command(command) })?; - // TODO remove this when we use the stub, the stub should be taking care of - // this. TODO do we also need to disable rtc super wdt? + // TODO: Remove this when we use the stub, the stub should be taking care of + // this. + // TODO: Do we also need to disable RTC super WDT? if connection.get_usb_pid()? == USB_SERIAL_JTAG_PID { match self.chip { Chip::Esp32c3 => { connection.command(Command::WriteReg { - address: 0x600080a8, - value: 0x50D83AA1u32, + address: 0x6000_80a8, + value: 0x50D8_3AA1, mask: None, })?; // WP disable connection.command(Command::WriteReg { - address: 0x60008090, + address: 0x6000_8090, value: 0x0, mask: None, - })?; // turn off RTC WDG + })?; // turn off RTC WDT connection.command(Command::WriteReg { - address: 0x600080a8, + address: 0x6000_80a8, value: 0x0, mask: None, })?; // WP enable @@ -71,14 +73,14 @@ impl FlashTarget for Esp32Target { Chip::Esp32s3 => { connection.command(Command::WriteReg { address: 0x6000_80B0, - value: 0x50D83AA1u32, + value: 0x50D8_3AA1, mask: None, })?; // WP disable connection.command(Command::WriteReg { address: 0x6000_8098, value: 0x0, mask: None, - })?; // turn off RTC WDG + })?; // turn off RTC WDT connection.command(Command::WriteReg { address: 0x6000_80B0, value: 0x0, @@ -127,16 +129,6 @@ impl FlashTarget for Esp32Target { let chunks = compressed.chunks(flash_write_size); - let (_, chunk_size) = chunks.size_hint(); - let chunk_size = chunk_size.unwrap_or(0) as u64; - let pb_chunk = ProgressBar::new(chunk_size); - pb_chunk.set_style( - ProgressStyle::default_bar() - .template("[{elapsed_precise}] [{bar:40}] {pos:>7}/{len:7} {msg}") - .unwrap() - .progress_chars("=> "), - ); - // decode the chunks to see how much data the device will have to save let mut decoder = ZlibDecoder::new(Vec::new()); let mut decoded_size = 0; @@ -147,7 +139,6 @@ impl FlashTarget for Esp32Target { let size = decoder.get_ref().len() - decoded_size; decoded_size = decoder.get_ref().len(); - pb_chunk.set_message(format!("segment 0x{:X} writing chunks", addr)); connection.with_timeout( CommandType::FlashDeflateData.timeout_for_size(size as u32), |connection| { @@ -160,11 +151,8 @@ impl FlashTarget for Esp32Target { Ok(()) }, )?; - pb_chunk.inc(1); } - pb_chunk.finish_with_message(format!("segment 0x{:X}", addr)); - Ok(()) } @@ -172,10 +160,11 @@ impl FlashTarget for Esp32Target { connection.with_timeout(CommandType::FlashDeflateEnd.timeout(), |connection| { connection.command(Command::FlashDeflateEnd { reboot: false }) })?; + if reboot { - connection.reset() - } else { - Ok(()) + connection.reset()?; } + + Ok(()) } } diff --git a/espflash/src/targets/flash_target/esp8266.rs b/espflash/src/targets/flash_target/esp8266.rs index 82e778f3..dd0822b6 100644 --- a/espflash/src/targets/flash_target/esp8266.rs +++ b/espflash/src/targets/flash_target/esp8266.rs @@ -1,5 +1,3 @@ -use indicatif::{ProgressBar, ProgressStyle}; - use super::FlashTarget; use crate::{ command::{Command, CommandType}, @@ -28,6 +26,7 @@ impl FlashTarget for Esp8266Target { offset: 0, supports_encryption: false, })?; + Ok(()) } @@ -56,29 +55,15 @@ impl FlashTarget for Esp8266Target { let chunks = segment.data.chunks(FLASH_WRITE_SIZE); - let (_, chunk_size) = chunks.size_hint(); - let chunk_size = chunk_size.unwrap_or(0) as u64; - let pb_chunk = ProgressBar::new(chunk_size); - pb_chunk.set_style( - ProgressStyle::default_bar() - .template("[{elapsed_precise}] [{bar:40}] {pos:>7}/{len:7} {msg}") - .unwrap() - .progress_chars("=> "), - ); - for (i, block) in chunks.enumerate() { - pb_chunk.set_message(format!("segment 0x{:X} writing chunks", addr)); connection.command(Command::FlashData { sequence: i as u32, pad_to: FLASH_WRITE_SIZE, pad_byte: 0xff, data: block, })?; - pb_chunk.inc(1); } - pb_chunk.finish_with_message(format!("segment 0x{:X}", addr)); - Ok(()) } @@ -86,10 +71,11 @@ impl FlashTarget for Esp8266Target { connection.with_timeout(CommandType::FlashEnd.timeout(), |connection| { connection.write_command(Command::FlashEnd { reboot: false }) })?; + if reboot { - connection.reset() - } else { - Ok(()) + connection.reset()?; } + + Ok(()) } } From d47d994942ed278ccab2065844065af0cbd0545e Mon Sep 17 00:00:00 2001 From: Jesse Braham Date: Fri, 4 Nov 2022 13:25:51 -0700 Subject: [PATCH 2/3] Make flash targets take an optional callback function to update progress --- espflash/src/targets/flash_target/esp32.rs | 6 ++++++ espflash/src/targets/flash_target/esp8266.rs | 6 ++++++ espflash/src/targets/flash_target/mod.rs | 1 + espflash/src/targets/flash_target/ram.rs | 12 +++++++++++- 4 files changed, 24 insertions(+), 1 deletion(-) diff --git a/espflash/src/targets/flash_target/esp32.rs b/espflash/src/targets/flash_target/esp32.rs index eb7b1413..f482179d 100644 --- a/espflash/src/targets/flash_target/esp32.rs +++ b/espflash/src/targets/flash_target/esp32.rs @@ -98,6 +98,7 @@ impl FlashTarget for Esp32Target { &mut self, connection: &mut Connection, segment: RomSegment, + progress_cb: Option>, ) -> Result<(), Error> { let addr = segment.addr; @@ -128,6 +129,7 @@ impl FlashTarget for Esp32Target { )?; let chunks = compressed.chunks(flash_write_size); + let num_chunks = chunks.len(); // decode the chunks to see how much data the device will have to save let mut decoder = ZlibDecoder::new(Vec::new()); @@ -151,6 +153,10 @@ impl FlashTarget for Esp32Target { Ok(()) }, )?; + + if let Some(ref cb) = progress_cb { + cb(i + 1, num_chunks); + } } Ok(()) diff --git a/espflash/src/targets/flash_target/esp8266.rs b/espflash/src/targets/flash_target/esp8266.rs index dd0822b6..507c5ca5 100644 --- a/espflash/src/targets/flash_target/esp8266.rs +++ b/espflash/src/targets/flash_target/esp8266.rs @@ -34,6 +34,7 @@ impl FlashTarget for Esp8266Target { &mut self, connection: &mut Connection, segment: RomSegment, + progress_cb: Option>, ) -> Result<(), Error> { let addr = segment.addr; let block_count = (segment.data.len() + FLASH_WRITE_SIZE - 1) / FLASH_WRITE_SIZE; @@ -54,6 +55,7 @@ impl FlashTarget for Esp8266Target { )?; let chunks = segment.data.chunks(FLASH_WRITE_SIZE); + let num_chunks = chunks.len(); for (i, block) in chunks.enumerate() { connection.command(Command::FlashData { @@ -62,6 +64,10 @@ impl FlashTarget for Esp8266Target { pad_byte: 0xff, data: block, })?; + + if let Some(ref cb) = progress_cb { + cb(i + 1, num_chunks); + } } Ok(()) diff --git a/espflash/src/targets/flash_target/mod.rs b/espflash/src/targets/flash_target/mod.rs index 2d86b095..d0c52fa8 100644 --- a/espflash/src/targets/flash_target/mod.rs +++ b/espflash/src/targets/flash_target/mod.rs @@ -18,6 +18,7 @@ pub trait FlashTarget { &mut self, connection: &mut Connection, segment: RomSegment, + progress_cb: Option>, ) -> Result<(), Error>; /// Complete the flashing operation diff --git a/espflash/src/targets/flash_target/ram.rs b/espflash/src/targets/flash_target/ram.rs index 9c64e60f..a2c612d4 100644 --- a/espflash/src/targets/flash_target/ram.rs +++ b/espflash/src/targets/flash_target/ram.rs @@ -44,6 +44,7 @@ impl FlashTarget for RamTarget { &mut self, connection: &mut Connection, segment: RomSegment, + progress_cb: Option>, ) -> Result<(), Error> { let padding = 4 - segment.data.len() % 4; let block_count = (segment.data.len() + padding + self.block_size - 1) / self.block_size; @@ -56,14 +57,22 @@ impl FlashTarget for RamTarget { supports_encryption: false, })?; - for (i, block) in segment.data.chunks(self.block_size).enumerate() { + let chunks = segment.data.chunks(self.block_size); + let num_chunks = chunks.len(); + + for (i, block) in chunks.enumerate() { connection.command(Command::MemData { sequence: i as u32, pad_to: 4, pad_byte: 0, data: block, })?; + + if let Some(ref cb) = progress_cb { + cb(i + 1, num_chunks); + } } + Ok(()) } @@ -77,6 +86,7 @@ impl FlashTarget for RamTarget { }) })?; } + Ok(()) } } From 1d33a5a212adc4bc8fce9f34bd27392dde3da0b6 Mon Sep 17 00:00:00 2001 From: Jesse Braham Date: Fri, 4 Nov 2022 13:43:07 -0700 Subject: [PATCH 3/3] Add back in our progress bars, but only when the `cli` feature is enabled --- espflash/src/bin/espflash.rs | 12 +++++--- espflash/src/cli/mod.rs | 54 ++++++++++++++++++++++++++++++++- espflash/src/flasher/mod.rs | 59 ++++++++++++++++++++++++++---------- 3 files changed, 104 insertions(+), 21 deletions(-) diff --git a/espflash/src/bin/espflash.rs b/espflash/src/bin/espflash.rs index 5ed05475..16ad37a1 100644 --- a/espflash/src/bin/espflash.rs +++ b/espflash/src/bin/espflash.rs @@ -8,9 +8,10 @@ use std::{ use clap::{Args, Parser, Subcommand}; use espflash::{ cli::{ - 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, + self, board_info, build_progress_bar_callback, clap_enum_variants, config::Config, connect, + erase_partitions, flash_elf_image, monitor::monitor, parse_partition_table, + partition_table, progress_bar, save_elf_as_image, serial_monitor, ConnectArgs, + FlashConfigArgs, PartitionTableArgs, }, image_format::ImageFormatKind, logging::initialize_logger, @@ -206,7 +207,10 @@ fn write_bin(args: WriteBinArgs, config: &Config) -> Result<()> { let mut buffer = Vec::with_capacity(size.try_into().into_diagnostic()?); f.read_to_end(&mut buffer).into_diagnostic()?; - flasher.write_bin_to_flash(args.addr, &buffer)?; + let progress = progress_bar(format!("segment 0x{:X}", args.addr), None); + let progress_cb = Some(build_progress_bar_callback(progress)); + + flasher.write_bin_to_flash(args.addr, &buffer, progress_cb)?; Ok(()) } diff --git a/espflash/src/cli/mod.rs b/espflash/src/cli/mod.rs index dd4bb102..e796d3b2 100644 --- a/espflash/src/cli/mod.rs +++ b/espflash/src/cli/mod.rs @@ -8,6 +8,7 @@ //! [espflash]: https://crates.io/crates/espflash use std::{ + borrow::Cow, collections::HashMap, fs, io::Write, @@ -17,7 +18,7 @@ use std::{ use clap::{builder::ArgPredicate, Args}; use comfy_table::{modifiers, presets::UTF8_FULL, Attribute, Cell, Color, Table}; use esp_idf_part::{DataType, Partition, PartitionTable}; -use indicatif::HumanCount; +use indicatif::{style::ProgressStyle, HumanCount, ProgressBar, ProgressDrawTarget}; use log::{debug, info}; use miette::{IntoDiagnostic, Result, WrapErr}; use serialport::{SerialPortType, UsbPortInfo}; @@ -169,6 +170,57 @@ pub struct SaveImageArgs { pub skip_padding: bool, } +/// Create a new [ProgressBar] with some message and styling applied +pub fn progress_bar(msg: S, len: Option) -> ProgressBar +where + S: Into>, +{ + // If no length was provided, or the provided length is 0, we will initially + // hide the progress bar. This is done to avoid drawing a full-width progress + // bar before we've determined the length. + let draw_target = match len { + Some(len) if len > 0 => ProgressDrawTarget::stderr(), + _ => ProgressDrawTarget::hidden(), + }; + + let progress = ProgressBar::with_draw_target(len, draw_target) + .with_message(msg) + .with_style( + ProgressStyle::default_bar() + .template("[{elapsed_precise}] [{bar:40}] {pos:>7}/{len:7} {msg}") + .unwrap() + .progress_chars("=> "), + ); + + progress +} + +/// Create a callback function for the provided [ProgressBar] +pub fn build_progress_bar_callback(pb: ProgressBar) -> Box { + // This is a bit of an odd callback function, as it handles the entire lifecycle + // of the progress bar. + Box::new(move |current: usize, total: usize| { + // If the length has not yet been set, set it and then change the draw target to + // make the progress bar visible. + match pb.length() { + Some(0) | None => { + pb.set_length(total as u64); + pb.set_draw_target(ProgressDrawTarget::stderr()); + } + _ => {} + } + + // Set the new position of the progress bar. + pb.set_position(current as u64); + + // If we have reached the end, make sure to finish the progress bar to ensure + // proper output on the terminal. + if current == total { + pb.finish(); + } + }) +} + /// Select a serial port and establish a connection with a target device pub fn connect(args: &ConnectArgs, config: &Config) -> Result { let port_info = get_serial_port_info(args, config)?; diff --git a/espflash/src/flasher/mod.rs b/espflash/src/flasher/mod.rs index 88ff2ed1..c668f96d 100644 --- a/espflash/src/flasher/mod.rs +++ b/espflash/src/flasher/mod.rs @@ -394,6 +394,7 @@ impl Flasher { addr: text_addr, data: Cow::Borrowed(&text), }, + None, ) .flashing()?; @@ -407,6 +408,7 @@ impl Flasher { addr: data_addr, data: Cow::Borrowed(&data), }, + None, ) .flashing()?; @@ -637,6 +639,9 @@ impl Flasher { /// Note that this will not touch the flash on the device pub fn load_elf_to_ram(&mut self, elf_data: &[u8]) -> Result<(), Error> { let image = ElfFirmwareImage::try_from(elf_data)?; + if image.rom_segments(self.chip).next().is_some() { + return Err(Error::ElfNotRamLoadable); + } let mut target = self.chip.ram_target( Some(image.entry()), @@ -646,19 +651,21 @@ impl Flasher { ); target.begin(&mut self.connection).flashing()?; - if image.rom_segments(self.chip).next().is_some() { - return Err(Error::ElfNotRamLoadable); - } - for segment in image.ram_segments(self.chip) { + // Only display progress bars when the "cli" feature is enabled. + let progress_cb = if cfg!(feature = "cli") { + use crate::cli::{build_progress_bar_callback, progress_bar}; + + let progress = progress_bar(format!("segment 0x{:X}", segment.addr), None); + let progress_cb = build_progress_bar_callback(progress); + + Some(progress_cb) + } else { + None + }; + target - .write_segment( - &mut self.connection, - RomSegment { - addr: segment.addr, - data: Cow::Borrowed(segment.data()), - }, - ) + .write_segment(&mut self.connection, segment.into(), progress_cb) .flashing()?; } @@ -694,12 +701,25 @@ impl Flasher { flash_freq, )?; + // When the "cli" feature is enabled, display the image size information. #[cfg(feature = "cli")] crate::cli::display_image_size(image.app_size(), image.part_size()); for segment in image.flash_segments() { + // Only display progress bars when the "cli" feature is enabled. + let progress_cb = if cfg!(feature = "cli") { + use crate::cli::{build_progress_bar_callback, progress_bar}; + + let progress = progress_bar(format!("segment 0x{:X}", segment.addr), None); + let progress_cb = build_progress_bar_callback(progress); + + Some(progress_cb) + } else { + None + }; + target - .write_segment(&mut self.connection, segment) + .write_segment(&mut self.connection, segment, progress_cb) .flashing()?; } @@ -709,15 +729,22 @@ impl Flasher { } /// Load an bin image to flash at a specific address - pub fn write_bin_to_flash(&mut self, addr: u32, data: &[u8]) -> Result<(), Error> { - let mut target = self.chip.flash_target(self.spi_params, self.use_stub); - target.begin(&mut self.connection).flashing()?; + pub fn write_bin_to_flash( + &mut self, + addr: u32, + data: &[u8], + progress_cb: Option>, + ) -> Result<(), Error> { let segment = RomSegment { addr, data: Cow::from(data), }; - target.write_segment(&mut self.connection, segment)?; + + let mut target = self.chip.flash_target(self.spi_params, self.use_stub); + target.begin(&mut self.connection).flashing()?; + target.write_segment(&mut self.connection, segment, progress_cb)?; target.finish(&mut self.connection, true).flashing()?; + Ok(()) }