Add before and after reset arguments (#561)

* feat: Initial implementation of reset args

* feat: Add after reset methods

* feat: Allow monitoring without reseting

* style: Fix clippy warnings

* feat: Use reset_after when possible to take into account `--after` flag

* feat: Add boot_mode and download mode checks

* docs: Update changelog

* chore: Cleanup comments

* style: Use info! instead of println!

Author: SergioGasquez

CHANGELOG.md

@@ -22,6 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add `serialport` feature. (#535)
 - Add support for 26 MHz bootloader for ESP32 and ESP32-C2 (#553)
 - Add CI check to verify that CHANGELOG is updated (#560)
+- Add `--before` and `--after` reset arguments (#561)
 
 ### Fixed
 

cargo-espflash/src/main.rs

@@ -252,8 +252,11 @@ pub fn erase_parts(args: ErasePartsArgs, config: &Config) -> Result<()> {
     };
 
     info!("Erasing the following partitions: {:?}", args.erase_parts);
+    let chip: Chip = flash.chip();
     erase_partitions(&mut flash, partition_table, Some(args.erase_parts), None)?;
-    flash.connection().reset()?;
+    flash
+        .connection()
+        .reset_after(!args.connect_args.no_stub, chip)?;
 
     Ok(())
 }

espflash/src/bin/espflash.rs

@@ -195,8 +195,11 @@ pub fn erase_parts(args: ErasePartsArgs, config: &Config) -> Result<()> {
     };
 
     info!("Erasing the following partitions: {:?}", args.erase_parts);
+    let chip = flash.chip();
     erase_partitions(&mut flash, partition_table, Some(args.erase_parts), None)?;
-    flash.connection().reset()?;
+    flash
+        .connection()
+        .reset_after(!args.connect_args.no_stub, chip)?;
 
     Ok(())
 }

espflash/src/cli/mod.rs

@@ -18,7 +18,7 @@ use clap_complete::Shell;
 use comfy_table::{modifiers, presets::UTF8_FULL, Attribute, Cell, Color, Table};
 use esp_idf_part::{DataType, Partition, PartitionTable};
 use indicatif::{style::ProgressStyle, HumanCount, ProgressBar};
-use log::{debug, info};
+use log::{debug, info, warn};
 use miette::{IntoDiagnostic, Result, WrapErr};
 use serialport::{SerialPortType, UsbPortInfo};
 
@@ -28,6 +28,7 @@ use self::{
     serial::get_serial_port_info,
 };
 use crate::{
+    connection::reset::{ResetAfterOperation, ResetBeforeOperation},
     elf::ElfFirmwareImage,
     error::{Error, MissingPartition, MissingPartitionTable},
     flasher::{
@@ -48,9 +49,15 @@ mod serial;
 #[derive(Debug, Args)]
 #[non_exhaustive]
 pub struct ConnectArgs {
+    /// Reset operation to perform after connecting to the target
+    #[arg(short = 'a', long, default_value = "hard-reset")]
+    pub after: ResetAfterOperation,
     /// Baud rate at which to communicate with target device
-    #[arg(short = 'b', long, env = "ESPFLASH_BAUD")]
+    #[arg(short = 'B', long, env = "ESPFLASH_BAUD")]
     pub baud: Option<u32>,
+    /// Reset operation to perform before connecting to the target
+    #[arg(short = 'b', long, default_value = "default-reset")]
+    pub before: ResetBeforeOperation,
     /// Target device
     #[arg(short = 'c', long)]
     pub chip: Option<Chip>,
@@ -263,6 +270,15 @@ pub fn connect(
     no_verify: bool,
     no_skip: bool,
 ) -> Result<Flasher> {
+    if args.before == ResetBeforeOperation::NoReset
+        || args.before == ResetBeforeOperation::NoResetNoSync
+    {
+        warn!(
+            "Pre-connection option '{:#?}' was selected. Connection may fail if the chip is not in bootloader or flasher stub mode.",
+            args.before
+        );
+    }
+
     let port_info = get_serial_port_info(args, config)?;
 
     // Attempt to open the serial port and set its initial baud rate.
@@ -305,6 +321,8 @@ pub fn connect(
         !no_verify,
         !no_skip,
         args.chip,
+        args.after,
+        args.before,
     )?)
 }
 
@@ -547,8 +565,10 @@ pub fn erase_flash(args: EraseFlashArgs, config: &Config) -> Result<()> {
 
     info!("Erasing Flash...");
     flash.erase_flash()?;
-
-    flash.connection().reset()?;
+    let chip = flash.chip();
+    flash
+        .connection()
+        .reset_after(!args.connect_args.no_stub, chip)?;
 
     Ok(())
 }
@@ -565,7 +585,10 @@ pub fn erase_region(args: EraseRegionArgs, config: &Config) -> Result<()> {
         args.addr, args.size
     );
     flash.erase_region(args.addr, args.size)?;
-    flash.connection().reset()?;
+    let chip = flash.chip();
+    flash
+        .connection()
+        .reset_after(!args.connect_args.no_stub, chip)?;
 
     Ok(())
 }

espflash/src/command.rs

@@ -42,6 +42,7 @@ pub enum CommandType {
     // Some commands supported by stub only
     EraseFlash = 0xd0,
     EraseRegion = 0xd1,
+    RunUserCode = 0xd3,
 }
 
 impl CommandType {
@@ -165,6 +166,7 @@ pub enum Command<'a> {
         offset: u32,
         size: u32,
     },
+    RunUserCode,
 }
 
 impl<'a> Command<'a> {
@@ -191,6 +193,7 @@ impl<'a> Command<'a> {
             Command::EraseFlash { .. } => CommandType::EraseFlash,
             Command::EraseRegion { .. } => CommandType::EraseRegion,
             Command::FlashMd5 { .. } => CommandType::FlashMd5,
+            Command::RunUserCode { .. } => CommandType::RunUserCode,
         }
     }
 
@@ -379,6 +382,9 @@ impl<'a> Command<'a> {
                 writer.write_all(&(0u32.to_le_bytes()))?;
                 writer.write_all(&(0u32.to_le_bytes()))?;
             }
+            Command::RunUserCode => {
+                write_basic(writer, &[], 0)?;
+            }
         };
         Ok(())
     }

espflash/src/connection/mod.rs

@@ -11,23 +11,29 @@ use std::{
     time::Duration,
 };
 
-use log::debug;
+use log::{debug, info};
+use regex::Regex;
 use serialport::UsbPortInfo;
 use slip_codec::SlipDecoder;
 
 #[cfg(unix)]
 use self::reset::UnixTightReset;
 use self::{
     encoder::SlipEncoder,
-    reset::{construct_reset_strategy_sequence, ClassicReset, ResetStrategy, UsbJtagSerialReset},
+    reset::{
+        construct_reset_strategy_sequence, ClassicReset, HardReset, ResetAfterOperation,
+        ResetBeforeOperation, ResetStrategy, UsbJtagSerialReset,
+    },
 };
 use crate::{
     command::{Command, CommandType},
+    connection::reset::soft_reset,
     error::{ConnectionError, Error, ResultExt, RomError, RomErrorKind},
     interface::Interface,
+    targets::Chip,
 };
 
-mod reset;
+pub mod reset;
 
 const MAX_CONNECT_ATTEMPTS: usize = 7;
 const MAX_SYNC_ATTEMPTS: usize = 5;
@@ -77,21 +83,34 @@ pub struct Connection {
     serial: Interface,
     port_info: UsbPortInfo,
     decoder: SlipDecoder,
+    after_operation: ResetAfterOperation,
+    before_operation: ResetBeforeOperation,
 }
 
 impl Connection {
-    pub fn new(serial: Interface, port_info: UsbPortInfo) -> Self {
+    pub fn new(
+        serial: Interface,
+        port_info: UsbPortInfo,
+        after_operation: ResetAfterOperation,
+        before_operation: ResetBeforeOperation,
+    ) -> Self {
         Connection {
             serial,
             port_info,
             decoder: SlipDecoder::new(),
+            after_operation,
+            before_operation,
         }
     }
 
     /// Initialize a connection with a device
     pub fn begin(&mut self) -> Result<(), Error> {
         let port_name = self.serial.serial_port().name().unwrap_or_default();
-        let reset_sequence = construct_reset_strategy_sequence(&port_name, self.port_info.pid);
+        let reset_sequence = construct_reset_strategy_sequence(
+            &port_name,
+            self.port_info.pid,
+            self.before_operation,
+        );
 
         for (_, reset_strategy) in zip(0..MAX_CONNECT_ATTEMPTS, reset_sequence.iter().cycle()) {
             match self.connect_attempt(reset_strategy) {
@@ -110,7 +129,47 @@ impl Connection {
     /// Try to connect to a device
     #[allow(clippy::borrowed_box)]
     fn connect_attempt(&mut self, reset_strategy: &Box<dyn ResetStrategy>) -> Result<(), Error> {
-        reset_strategy.reset(&mut self.serial)?;
+        // If we're doing no_sync, we're likely communicating as a pass through
+        // with an intermediate device to the ESP32
+        if self.before_operation == ResetBeforeOperation::NoResetNoSync {
+            return Ok(());
+        }
+        let mut download_mode: bool = false;
+        let mut boot_mode: &str = "";
+        let mut boot_log_detected = false;
+        let mut buff: Vec<u8>;
+        if self.before_operation != ResetBeforeOperation::NoReset {
+            // Reset the chip to bootloader (download mode)
+            reset_strategy.reset(&mut self.serial)?;
+
+            let available_bytes = self.serial.serial_port_mut().bytes_to_read()?;
+            buff = vec![0; available_bytes as usize];
+            let read_bytes = self.serial.serial_port_mut().read(&mut buff)? as u32;
+
+            if read_bytes != available_bytes {
+                return Err(Error::Connection(ConnectionError::ReadMissmatch(
+                    available_bytes,
+                    read_bytes,
+                )));
+            }
+
+            let read_slice = std::str::from_utf8(&buff[..read_bytes as usize]).unwrap();
+
+            let pattern = Regex::new(r"boot:(0x[0-9a-fA-F]+)(.*waiting for download)?").unwrap();
+
+            // Search for the pattern in the read data
+            if let Some(data) = pattern.captures(read_slice) {
+                boot_log_detected = true;
+                // Boot log detected
+                boot_mode = data.get(1).map(|m| m.as_str()).unwrap_or_default();
+                download_mode = data.get(2).is_some();
+
+                // Further processing or printing the results
+                debug!("Boot Mode: {}", boot_mode);
+                debug!("Download Mode: {}", download_mode);
+            };
+        }
+
         for _ in 0..MAX_SYNC_ATTEMPTS {
             self.flush()?;
 
@@ -119,6 +178,16 @@ impl Connection {
             }
         }
 
+        if boot_log_detected {
+            if download_mode {
+                return Err(Error::Connection(ConnectionError::NoSyncReply));
+            } else {
+                return Err(Error::Connection(ConnectionError::WrongBootMode(
+                    boot_mode.to_string(),
+                )));
+            }
+        }
+
         Err(Error::Connection(ConnectionError::ConnectionFailed))
     }
 
@@ -163,6 +232,28 @@ impl Connection {
         Ok(())
     }
 
+    // Reset the device taking into account the reset after argument
+    pub fn reset_after(&mut self, is_stub: bool, chip: Chip) -> Result<(), Error> {
+        match self.after_operation {
+            ResetAfterOperation::HardReset => HardReset.reset(&mut self.serial),
+            ResetAfterOperation::SoftReset => {
+                info!("Soft resetting");
+                soft_reset(self, false, is_stub, chip)?;
+                Ok(())
+            }
+            ResetAfterOperation::NoReset => {
+                info!("Staying in bootloader");
+                soft_reset(self, true, is_stub, chip)?;
+
+                Ok(())
+            }
+            ResetAfterOperation::NoResetNoStub => {
+                info!("Staying in flasher stub");
+                Ok(())
+            }
+        }
+    }
+
     // Reset the device to flash mode
     pub fn reset_to_flash(&mut self, extra_delay: bool) -> Result<(), Error> {
         if self.port_info.pid == USB_SERIAL_JTAG_PID {