Merge pull request #54 from esp-rs/deflate-begin-timeout

fix erease timeout not being used for compressed flashing

Author: icewind1991

espflash/src/connection.rs

@@ -3,7 +3,8 @@ use std::thread::sleep;
 use std::time::Duration;
 
 use crate::encoder::SlipEncoder;
-use crate::error::{ConnectionError, Error, RomError};
+use crate::error::{ConnectionError, Error, ResultExt, RomError};
+use crate::flasher::Command;
 use binread::io::Cursor;
 use binread::{BinRead, BinReaderExt};
 use serial::{BaudRate, SerialPort, SerialPortSettings, SystemPort};
@@ -114,14 +115,15 @@ impl Connection {
 
     pub fn command<Data: LazyBytes<SystemPort>>(
         &mut self,
-        command: u8,
+        command: Command,
         data: Data,
         check: u32,
     ) -> Result<u32, Error> {
-        self.write_command(command, data, check)?;
+        self.write_command(command as u8, data, check)
+            .for_command(command)?;
 
         for _ in 0..100 {
-            match self.read_response()? {
+            match self.read_response().for_command(command)? {
                 Some(response) if response.return_op == command as u8 => {
                     if response.status == 1 {
                         let _error = self.flush();

espflash/src/error.rs

@@ -1,6 +1,8 @@
+use crate::flasher::Command;
 use csv::Position;
 use miette::{Diagnostic, SourceOffset, SourceSpan};
 use slip_codec::Error as SlipError;
+use std::fmt::{Display, Formatter};
 use std::io;
 use thiserror::Error;
 
@@ -63,9 +65,9 @@ pub enum ConnectionError {
         help("Ensure that the device is connected and your host recognizes the serial adapter")
     )]
     DeviceNotFound,
-    #[error("Timeout while running command")]
+    #[error("Timeout while running {0}command")]
     #[diagnostic(code(espflash::timeout))]
-    Timeout,
+    Timeout(TimedOutCommand),
     #[error("Received packet has invalid SLIP framing")]
     #[diagnostic(
         code(espflash::slip_framing),
@@ -80,6 +82,26 @@ pub enum ConnectionError {
     OverSizedPacket,
 }
 
+#[derive(Debug, Default, Clone)]
+pub struct TimedOutCommand {
+    command: Option<Command>,
+}
+
+impl From<Command> for TimedOutCommand {
+    fn from(c: Command) -> Self {
+        TimedOutCommand { command: Some(c) }
+    }
+}
+
+impl Display for TimedOutCommand {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match &self.command {
+            Some(command) => write!(f, "{} ", command),
+            None => Ok(()),
+        }
+    }
+}
+
 impl From<serial::Error> for ConnectionError {
     fn from(err: serial::Error) -> Self {
         match err.kind() {
@@ -110,7 +132,7 @@ impl From<io::Error> for Error {
 
 fn from_error_kind<E: Into<serial::Error>>(kind: io::ErrorKind, err: E) -> ConnectionError {
     match kind {
-        io::ErrorKind::TimedOut => ConnectionError::Timeout,
+        io::ErrorKind::TimedOut => ConnectionError::Timeout(TimedOutCommand::default()),
         io::ErrorKind::NotFound => ConnectionError::DeviceNotFound,
         _ => ConnectionError::Serial(err.into()),
     }
@@ -197,6 +219,8 @@ impl From<u8> for RomError {
 pub(crate) trait ResultExt {
     /// mark an error as having occurred during the flashing stage
     fn flashing(self) -> Self;
+    /// mark the command from which this error originates
+    fn for_command(self, command: Command) -> Self;
 }
 
 impl<T> ResultExt for Result<T, Error> {
@@ -206,6 +230,18 @@ impl<T> ResultExt for Result<T, Error> {
             res => res,
         }
     }
+
+    fn for_command(self, command: Command) -> Self {
+        match self {
+            Err(Error::Connection(ConnectionError::Timeout(_))) => {
+                Err(Error::Connection(ConnectionError::Timeout(command.into())))
+            }
+            Err(Error::Flashing(ConnectionError::Timeout(_))) => {
+                Err(Error::Flashing(ConnectionError::Timeout(command.into())))
+            }
+            res => res,
+        }
+    }
 }
 
 #[derive(Debug, Error, Diagnostic)]

espflash/src/flash_target/esp32.rs

@@ -27,7 +27,7 @@ impl FlashTarget for Esp32Target {
     fn begin(&mut self, connection: &mut Connection, _image: &FirmwareImage) -> Result<(), Error> {
         let spi_params = self.spi_attach_params.encode();
         connection.with_timeout(Command::SpiAttach.timeout(), |connection| {
-            connection.command(Command::SpiAttach as u8, spi_params.as_slice(), 0)
+            connection.command(Command::SpiAttach, spi_params.as_slice(), 0)
         })?;
         Ok(())
     }

espflash/src/flash_target/mod.rs

@@ -61,7 +61,7 @@ fn begin_command(
     };
 
     connection.with_timeout(command.timeout_for_size(size), |connection| {
-        connection.command(command as u8, data, 0)?;
+        connection.command(command, data, 0)?;
         Ok(())
     })
 }
@@ -120,7 +120,7 @@ fn block_command_with_timeout(
 
     connection.with_timeout(timout, |connection| {
         connection.command(
-            command as u8,
+            command,
             (length as u16, |encoder: &mut Encoder| {
                 encoder.write(bytes_of(&params))?;
                 encoder.write(data)?;

espflash/src/flasher.rs

@@ -28,10 +28,12 @@ const ERASE_WRITE_TIMEOUT_PER_MB: Duration = Duration::from_secs(40);
 const MEM_END_TIMEOUT: Duration = Duration::from_millis(50);
 const SYNC_TIMEOUT: Duration = Duration::from_millis(100);
 
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, Debug, Display)]
 #[allow(dead_code)]
 #[repr(u8)]
-pub(crate) enum Command {
+#[non_exhaustive]
+pub enum Command {
+    Unknown = 0,
     FlashBegin = 0x02,
     FlashData = 0x03,
     FlashEnd = 0x04,
@@ -48,6 +50,7 @@ pub(crate) enum Command {
     FlashDeflateData = 0x11,
     FlashDeflateEnd = 0x12,
     FlashMd5 = 0x13,
+    FlashDetect = 0x9f,
 }
 
 impl Command {
@@ -68,7 +71,9 @@ impl Command {
             )
         }
         match self {
-            Command::FlashBegin => calc_timeout(ERASE_REGION_TIMEOUT_PER_MB, size),
+            Command::FlashBegin | Command::FlashDeflateBegin => {
+                calc_timeout(ERASE_REGION_TIMEOUT_PER_MB, size)
+            }
             Command::FlashData | Command::FlashDeflateData => {
                 calc_timeout(ERASE_WRITE_TIMEOUT_PER_MB, size)
             }
@@ -258,7 +263,7 @@ impl Flasher {
     }
 
     fn flash_detect(&mut self) -> Result<bool, Error> {
-        let flash_id = self.spi_command(0x9f, &[], 24)?;
+        let flash_id = self.spi_command(Command::FlashDetect, &[], 24)?;
         let size_id = flash_id >> 16;
 
         self.flash_size = FlashSize::from(size_id as u8)?;
@@ -340,7 +345,7 @@ impl Flasher {
 
         self.connection
             .with_timeout(command.timeout_for_size(size), |connection| {
-                connection.command(command as u8, data, 0)?;
+                connection.command(command, data, 0)?;
                 Ok(())
             })
     }
@@ -354,14 +359,14 @@ impl Flasher {
                 let spi_params = spi_attach_params.encode();
                 self.connection
                     .with_timeout(Command::SpiAttach.timeout(), |connection| {
-                        connection.command(Command::SpiAttach as u8, spi_params.as_slice(), 0)
+                        connection.command(Command::SpiAttach, spi_params.as_slice(), 0)
                     })?;
             }
         }
         Ok(())
     }
 
-    fn spi_command(&mut self, command: u8, data: &[u8], read_bits: u32) -> Result<u32, Error> {
+    fn spi_command(&mut self, command: Command, data: &[u8], read_bits: u32) -> Result<u32, Error> {
         assert!(read_bits < 32);
         assert!(data.len() < 64);
 
@@ -421,7 +426,7 @@ impl Flasher {
             }
             i += 1;
             if i > 10 {
-                return Err(Error::Connection(ConnectionError::Timeout));
+                return Err(Error::Connection(ConnectionError::Timeout(command.into())));
             }
         }
 
@@ -435,7 +440,7 @@ impl Flasher {
     fn read_reg(&mut self, reg: u32) -> Result<u32, Error> {
         self.connection
             .with_timeout(Command::ReadReg.timeout(), |connection| {
-                connection.command(Command::ReadReg as u8, &reg.to_le_bytes()[..], 0)
+                connection.command(Command::ReadReg, &reg.to_le_bytes()[..], 0)
             })
     }
 
@@ -448,7 +453,7 @@ impl Flasher {
         };
         self.connection
             .with_timeout(Command::WriteReg.timeout(), |connection| {
-                connection.command(Command::WriteReg as u8, bytes_of(&params), 0)
+                connection.command(Command::WriteReg, bytes_of(&params), 0)
             })?;
         Ok(())
     }
@@ -524,11 +529,7 @@ impl Flasher {
 
         self.connection
             .with_timeout(Command::ChangeBaud.timeout(), |connection| {
-                connection.command(
-                    Command::ChangeBaud as u8,
-                    &[new_speed, old_speed].concat()[..],
-                    0,
-                )
+                connection.command(Command::ChangeBaud, &[new_speed, old_speed].concat()[..], 0)
             })?;
         self.connection.set_baud(speed)?;
         std::thread::sleep(Duration::from_secs_f32(0.05));