Fix SpiAttach when using stub

Author: mertzt89

espflash/src/chip/mod.rs

@@ -299,10 +299,14 @@ impl Chip {
         Box::new(RamTarget::new(entry))
     }
 
-    pub fn flash_target(&self, spi_params: SpiAttachParams) -> Box<dyn FlashTarget> {
+    pub fn flash_target(
+        &self,
+        spi_params: SpiAttachParams,
+        use_stub: bool,
+    ) -> Box<dyn FlashTarget> {
         match self {
             Chip::Esp8266 => Box::new(Esp8266Target::new()),
-            _ => Box::new(Esp32Target::new(*self, spi_params)),
+            _ => Box::new(Esp32Target::new(*self, spi_params, use_stub)),
         }
     }
 

espflash/src/command.rs

@@ -112,6 +112,9 @@ pub enum Command<'a> {
     SpiAttach {
         spi_params: SpiAttachParams,
     },
+    SpiAttachStub {
+        spi_params: SpiAttachParams,
+    },
     ChangeBaud {
         /// New baud rate
         new_baud: u32,
@@ -150,6 +153,7 @@ impl<'a> Command<'a> {
             Command::WriteReg { .. } => CommandType::WriteReg,
             Command::ReadReg { .. } => CommandType::ReadReg,
             Command::SpiAttach { .. } => CommandType::SpiAttach,
+            Command::SpiAttachStub { .. } => CommandType::SpiAttach,
             Command::ChangeBaud { .. } => CommandType::ChangeBaud,
             Command::FlashDeflateBegin { .. } => CommandType::FlashDeflateBegin,
             Command::FlashDeflateData { .. } => CommandType::FlashDeflateData,
@@ -268,9 +272,15 @@ impl<'a> Command<'a> {
                 write_basic(writer, &address.to_le_bytes(), 0)?;
             }
             Command::SpiAttach { spi_params } => {
-                write_basic(writer, &spi_params.encode(), 0)?;
+                write_basic(writer, &spi_params.encode(false), 0)?;
             }
-            Command::ChangeBaud { new_baud, prior_baud } => {
+            Command::SpiAttachStub { spi_params } => {
+                write_basic(writer, &spi_params.encode(true), 0)?;
+            }
+            Command::ChangeBaud {
+                new_baud,
+                prior_baud,
+            } => {
                 // length
                 writer.write_all(&(8u16.to_le_bytes()))?;
                 // checksum

espflash/src/flash_target/esp32.rs

@@ -13,22 +13,30 @@ use std::io::Write;
 pub struct Esp32Target {
     chip: Chip,
     spi_attach_params: SpiAttachParams,
+    use_stub: bool,
 }
 
 impl Esp32Target {
-    pub fn new(chip: Chip, spi_attach_params: SpiAttachParams) -> Self {
+    pub fn new(chip: Chip, spi_attach_params: SpiAttachParams, use_stub: bool) -> Self {
         Esp32Target {
             chip,
             spi_attach_params,
+            use_stub,
         }
     }
 }
 
 impl FlashTarget for Esp32Target {
     fn begin(&mut self, connection: &mut Connection) -> Result<(), Error> {
         connection.with_timeout(CommandType::SpiAttach.timeout(), |connection| {
-            connection.command(Command::SpiAttach {
-                spi_params: self.spi_attach_params,
+            connection.command(if self.use_stub {
+                Command::SpiAttachStub {
+                    spi_params: self.spi_attach_params,
+                }
+            } else {
+                Command::SpiAttach {
+                    spi_params: self.spi_attach_params,
+                }
             })
         })?;
 

espflash/src/flasher.rs

@@ -138,17 +138,20 @@ impl SpiAttachParams {
         }
     }
 
-    pub fn encode(self) -> Vec<u8> {
+    pub fn encode(self, stub: bool) -> Vec<u8> {
         let packed = ((self.hd as u32) << 24)
             | ((self.cs as u32) << 18)
             | ((self.d as u32) << 12)
             | ((self.q as u32) << 6)
             | (self.clk as u32);
-        if packed == 0 {
-            vec![0; 5]
-        } else {
-            packed.to_le_bytes().to_vec()
+
+        let mut encoded: Vec<u8> = packed.to_le_bytes().to_vec();
+
+        if !stub {
+            encoded.append(&mut vec![0u8; 4]);
         }
+
+        encoded
     }
 }
 
@@ -369,7 +372,11 @@ impl Flasher {
             _ => {
                 self.connection
                     .with_timeout(CommandType::SpiAttach.timeout(), |connection| {
-                        connection.command(Command::SpiAttach { spi_params })
+                        connection.command(if self.use_stub {
+                            Command::SpiAttachStub { spi_params }
+                        } else {
+                            Command::SpiAttach { spi_params }
+                        })
                     })?;
             }
         }
@@ -537,7 +544,7 @@ impl Flasher {
     ) -> Result<(), Error> {
         let image = ElfFirmwareImage::try_from(elf_data)?;
 
-        let mut target = self.chip.flash_target(self.spi_params);
+        let mut target = self.chip.flash_target(self.spi_params, self.use_stub);
         target.begin(&mut self.connection).flashing()?;
 
         let flash_image = self.chip.get_flash_image(
@@ -564,7 +571,7 @@ 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);
+        let mut target = self.chip.flash_target(self.spi_params, self.use_stub);
         target.begin(&mut self.connection).flashing()?;
         let segment = RomSegment {
             addr,
@@ -606,7 +613,10 @@ impl Flasher {
 
         self.connection
             .with_timeout(CommandType::ChangeBaud.timeout(), |connection| {
-                connection.command(Command::ChangeBaud { new_baud: speed, prior_baud })
+                connection.command(Command::ChangeBaud {
+                    new_baud: speed,
+                    prior_baud,
+                })
             })?;
         self.connection.set_baud(speed)?;
         std::thread::sleep(Duration::from_secs_f32(0.05));