Implement skip and verify (#538)

* Implement skip and verify

* CHANGELOG.md entry

* Dedicated error for "verify failed"

CHANGELOG.md

@@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add --partition-table-offset argument for specifying the partition table offset (#516)
 - Add `Serialize` and `Deserialize` to `FlashFrequency`, `FlashMode` and `FlashSize`. (#528)
 - Add `checksum-md5` command (#536)
+- Add verify and skipping of unchanged flash regions - add `--no-verify` and `--no-skip` (#538)
 
 ### Fixed
 

cargo-espflash/src/main.rs

@@ -244,7 +244,7 @@ pub fn erase_parts(args: ErasePartsArgs, config: &Config) -> Result<()> {
         .as_deref()
         .or(metadata_partition_table.as_deref());
 
-    let mut flash = connect(&args.connect_args, config)?;
+    let mut flash = connect(&args.connect_args, config, false, false)?;
     let partition_table = match partition_table {
         Some(path) => Some(parse_partition_table(path)?),
         None => None,
@@ -261,7 +261,12 @@ fn flash(args: FlashArgs, config: &Config) -> Result<()> {
     let metadata = PackageMetadata::load(&args.build_args.package)?;
     let cargo_config = CargoConfig::load(&metadata.workspace_root, &metadata.package_root);
 
-    let mut flasher = connect(&args.connect_args, config)?;
+    let mut flasher = connect(
+        &args.connect_args,
+        config,
+        args.flash_args.no_verify,
+        args.flash_args.no_skip,
+    )?;
     flasher.verify_minimum_revision(args.flash_args.min_chip_rev)?;
 
     // If the user has provided a flash size via a command-line argument, we'll

espflash/Cargo.toml

@@ -43,6 +43,7 @@ hex = { version = "0.4.3", features = ["serde"], optional = true }
 indicatif = { version = "0.17.7", optional = true }
 lazy_static = { version = "1.4.0", optional = true }
 log = "0.4.20"
+md-5 = "0.10.6"
 miette = { version = "5.10.0", features = ["fancy"] }
 parse_int = { version = "0.6.0", optional = true }
 regex = { version = "1.9.6", optional = true }

espflash/src/bin/espflash.rs

@@ -187,7 +187,7 @@ pub fn erase_parts(args: ErasePartsArgs, config: &Config) -> Result<()> {
         return Err(Error::StubRequiredToEraseFlash).into_diagnostic();
     }
 
-    let mut flash = connect(&args.connect_args, config)?;
+    let mut flash = connect(&args.connect_args, config, false, false)?;
     let partition_table = match args.partition_table {
         Some(path) => Some(parse_partition_table(&path)?),
         None => None,
@@ -201,7 +201,12 @@ pub fn erase_parts(args: ErasePartsArgs, config: &Config) -> Result<()> {
 }
 
 fn flash(args: FlashArgs, config: &Config) -> Result<()> {
-    let mut flasher = connect(&args.connect_args, config)?;
+    let mut flasher = connect(
+        &args.connect_args,
+        config,
+        args.flash_args.no_verify,
+        args.flash_args.no_skip,
+    )?;
     flasher.verify_minimum_revision(args.flash_args.min_chip_rev)?;
 
     // If the user has provided a flash size via a command-line argument, we'll
@@ -325,7 +330,7 @@ fn save_image(args: SaveImageArgs) -> Result<()> {
 }
 
 fn write_bin(args: WriteBinArgs, config: &Config) -> Result<()> {
-    let mut flasher = connect(&args.connect_args, config)?;
+    let mut flasher = connect(&args.connect_args, config, false, false)?;
     print_board_info(&mut flasher)?;
 
     let mut f = File::open(&args.bin_file).into_diagnostic()?;

espflash/src/cli/mod.rs

@@ -163,6 +163,12 @@ pub struct FlashArgs {
     /// Load the application to RAM instead of Flash
     #[arg(long)]
     pub ram: bool,
+    /// Don't verify the flash contents after flashing
+    #[arg(long)]
+    pub no_verify: bool,
+    /// Don't skip flashing of parts with matching checksum
+    #[arg(long)]
+    pub no_skip: bool,
 }
 
 /// Operations for partitions tables
@@ -250,7 +256,12 @@ pub fn parse_u32(input: &str) -> Result<u32, ParseIntError> {
 }
 
 /// Select a serial port and establish a connection with a target device
-pub fn connect(args: &ConnectArgs, config: &Config) -> Result<Flasher> {
+pub fn connect(
+    args: &ConnectArgs,
+    config: &Config,
+    no_verify: bool,
+    no_skip: bool,
+) -> Result<Flasher> {
     let port_info = get_serial_port_info(args, config)?;
 
     // Attempt to open the serial port and set its initial baud rate.
@@ -290,21 +301,23 @@ pub fn connect(args: &ConnectArgs, config: &Config) -> Result<Flasher> {
         port_info,
         args.baud,
         !args.no_stub,
+        !no_verify,
+        !no_skip,
         args.chip,
     )?)
 }
 
 /// Connect to a target device and print information about its chip
 pub fn board_info(args: &ConnectArgs, config: &Config) -> Result<()> {
-    let mut flasher = connect(args, config)?;
+    let mut flasher = connect(args, config, true, true)?;
     print_board_info(&mut flasher)?;
 
     Ok(())
 }
 
 /// Connect to a target device and calculate the checksum of the given region
 pub fn checksum_md5(args: &ChecksumMd5Args, config: &Config) -> Result<()> {
-    let mut flasher = connect(&args.connect_args, config)?;
+    let mut flasher = connect(&args.connect_args, config, true, true)?;
 
     let checksum = flasher.checksum_md5(args.address, args.length)?;
     println!("0x{:x}", checksum);
@@ -369,7 +382,7 @@ pub fn print_board_info(flasher: &mut Flasher) -> Result<()> {
 
 /// Open a serial monitor
 pub fn serial_monitor(args: MonitorArgs, config: &Config) -> Result<()> {
-    let mut flasher = connect(&args.connect_args, config)?;
+    let mut flasher = connect(&args.connect_args, config, true, true)?;
     let pid = flasher.get_usb_pid()?;
 
     let elf = if let Some(elf_path) = args.elf {
@@ -585,7 +598,7 @@ pub fn erase_flash(args: EraseFlashArgs, config: &Config) -> Result<()> {
         return Err(Error::StubRequiredToEraseFlash).into_diagnostic();
     }
 
-    let mut flash = connect(&args.connect_args, config)?;
+    let mut flash = connect(&args.connect_args, config, true, true)?;
 
     info!("Erasing Flash...");
     flash.erase_flash()?;
@@ -600,7 +613,7 @@ pub fn erase_region(args: EraseRegionArgs, config: &Config) -> Result<()> {
         return Err(Error::StubRequiredToEraseFlash).into_diagnostic();
     }
 
-    let mut flash = connect(&args.connect_args, config)?;
+    let mut flash = connect(&args.connect_args, config, true, true)?;
 
     info!(
         "Erasing region at 0x{:08x} ({} bytes)",

espflash/src/error.rs

@@ -180,6 +180,10 @@ pub enum Error {
     #[diagnostic(transparent)]
     Defmt(#[from] DefmtError),
 
+    #[error("Verification of flash content failed")]
+    #[diagnostic(code(espflash::verify_failed))]
+    VerifyFailed,
+
     #[error("Internal Error")]
     InternalError,
 }

espflash/src/flasher/mod.rs

@@ -411,6 +411,10 @@ pub struct Flasher {
     spi_params: SpiAttachParams,
     /// Indicate RAM stub loader is in use
     use_stub: bool,
+    /// Indicate verifying flash contents after flashing
+    verify: bool,
+    /// Indicate skipping of already flashed regions
+    skip: bool,
 }
 
 impl Flasher {
@@ -419,6 +423,8 @@ impl Flasher {
         port_info: UsbPortInfo,
         speed: Option<u32>,
         use_stub: bool,
+        verify: bool,
+        skip: bool,
         chip: Option<Chip>,
     ) -> Result<Self, Error> {
         // Establish a connection to the device using the default baud rate of 115,200
@@ -445,6 +451,8 @@ impl Flasher {
             flash_size: FlashSize::_4Mb,
             spi_params: SpiAttachParams::default(),
             use_stub,
+            verify,
+            skip,
         };
 
         // Load flash stub if enabled
@@ -477,7 +485,9 @@ impl Flasher {
     }
 
     pub fn disable_watchdog(&mut self) -> Result<(), Error> {
-        let mut target = self.chip.flash_target(self.spi_params, self.use_stub);
+        let mut target = self
+            .chip
+            .flash_target(self.spi_params, self.use_stub, false, false);
         target.begin(&mut self.connection).flashing()?;
         Ok(())
     }
@@ -814,7 +824,9 @@ impl Flasher {
     ) -> Result<(), Error> {
         let image = ElfFirmwareImage::try_from(elf_data)?;
 
-        let mut target = self.chip.flash_target(self.spi_params, self.use_stub);
+        let mut target =
+            self.chip
+                .flash_target(self.spi_params, self.use_stub, self.verify, self.skip);
         target.begin(&mut self.connection).flashing()?;
 
         // The ESP8266 does not have readable major/minor revision numbers, so we have
@@ -878,7 +890,9 @@ impl Flasher {
         segments: &[RomSegment],
         mut progress: Option<&mut dyn ProgressCallbacks>,
     ) -> Result<(), Error> {
-        let mut target = self.chip.flash_target(self.spi_params, self.use_stub);
+        let mut target = self
+            .chip
+            .flash_target(self.spi_params, self.use_stub, false, false);
         target.begin(&mut self.connection).flashing()?;
         for segment in segments {
             target.write_segment(&mut self.connection, segment.borrow(), &mut progress)?;

espflash/src/targets/flash_target/esp32.rs

@@ -4,6 +4,8 @@ use flate2::{
     write::{ZlibDecoder, ZlibEncoder},
     Compression,
 };
+use log::debug;
+use md5::{Digest, Md5};
 
 use crate::{
     command::{Command, CommandType},
@@ -20,14 +22,26 @@ pub struct Esp32Target {
     chip: Chip,
     spi_attach_params: SpiAttachParams,
     use_stub: bool,
+    verify: bool,
+    skip: bool,
+    need_deflate_end: bool,
 }
 
 impl Esp32Target {
-    pub fn new(chip: Chip, spi_attach_params: SpiAttachParams, use_stub: bool) -> Self {
+    pub fn new(
+        chip: Chip,
+        spi_attach_params: SpiAttachParams,
+        use_stub: bool,
+        verify: bool,
+        skip: bool,
+    ) -> Self {
         Esp32Target {
             chip,
             spi_attach_params,
             use_stub,
+            verify,
+            skip,
+            need_deflate_end: false,
         }
     }
 }
@@ -120,6 +134,27 @@ impl FlashTarget for Esp32Target {
     ) -> Result<(), Error> {
         let addr = segment.addr;
 
+        let mut md5_hasher = Md5::new();
+        md5_hasher.update(&segment.data);
+        let checksum_md5 = md5_hasher.finalize();
+
+        if self.skip {
+            let flash_checksum_md5: u128 =
+                connection.with_timeout(CommandType::FlashMd5.timeout(), |connection| {
+                    connection
+                        .command(crate::command::Command::FlashMd5 {
+                            offset: addr,
+                            size: segment.data.len() as u32,
+                        })?
+                        .try_into()
+                })?;
+
+            if checksum_md5.as_slice() == flash_checksum_md5.to_be_bytes() {
+                debug!("Skipping segment at address {:x}", addr);
+                return Ok(());
+            }
+        }
+
         let mut encoder = ZlibEncoder::new(Vec::new(), Compression::best());
         encoder.write_all(&segment.data)?;
         let compressed = encoder.finish()?;
@@ -145,6 +180,7 @@ impl FlashTarget for Esp32Target {
                 Ok(())
             },
         )?;
+        self.need_deflate_end = true;
 
         let chunks = compressed.chunks(flash_write_size);
         let num_chunks = chunks.len();
@@ -185,13 +221,31 @@ impl FlashTarget for Esp32Target {
             cb.finish()
         }
 
+        if self.verify {
+            let flash_checksum_md5: u128 =
+                connection.with_timeout(CommandType::FlashMd5.timeout(), |connection| {
+                    connection
+                        .command(crate::command::Command::FlashMd5 {
+                            offset: addr,
+                            size: segment.data.len() as u32,
+                        })?
+                        .try_into()
+                })?;
+
+            if checksum_md5.as_slice() != flash_checksum_md5.to_be_bytes() {
+                return Err(Error::VerifyFailed);
+            }
+        }
+
         Ok(())
     }
 
     fn finish(&mut self, connection: &mut Connection, reboot: bool) -> Result<(), Error> {
-        connection.with_timeout(CommandType::FlashDeflateEnd.timeout(), |connection| {
-            connection.command(Command::FlashDeflateEnd { reboot: false })
-        })?;
+        if self.need_deflate_end {
+            connection.with_timeout(CommandType::FlashDeflateEnd.timeout(), |connection| {
+                connection.command(Command::FlashDeflateEnd { reboot: false })
+            })?;
+        }
 
         if reboot {
             connection.reset()?;

espflash/src/targets/mod.rs

@@ -108,10 +108,12 @@ impl Chip {
         &self,
         spi_params: SpiAttachParams,
         use_stub: bool,
+        verify: bool,
+        skip: bool,
     ) -> Box<dyn FlashTarget> {
         match self {
             Chip::Esp8266 => Box::new(Esp8266Target::new()),
-            _ => Box::new(Esp32Target::new(*self, spi_params, use_stub)),
+            _ => Box::new(Esp32Target::new(*self, spi_params, use_stub, verify, skip)),
         }
     }