Add support for flashing S2 via CDC UART (#228)

* Correct reset_after_flash behavior

* Add support for flashing S2 via CDC UART

This change makes it possible to flash an ESP32-S2 via CDC UART.

Author: mertzt89

espflash/src/chip/esp32/esp32s2.rs

@@ -6,11 +6,14 @@ use crate::{
     connection::Connection,
     elf::{FirmwareImage, FlashFrequency, FlashMode},
     error::UnsupportedImageFormatError,
-    flasher::FlashSize,
+    flash_target::MAX_RAM_BLOCK_SIZE,
+    flasher::{FlashSize, FLASH_WRITE_SIZE},
     image_format::{Esp32BootloaderFormat, ImageFormat, ImageFormatId},
     Chip, Error, PartitionTable,
 };
 
+const MAX_USB_BLOCK_SIZE: usize = 0x800;
+
 pub struct Esp32s2;
 
 pub const PARAMS: Esp32Params = Esp32Params::new(
@@ -79,6 +82,22 @@ impl ChipType for Esp32s2 {
         Ok(40)
     }
 
+    fn flash_write_size(&self, connection: &mut Connection) -> Result<usize, Error> {
+        Ok(if self.connection_is_usb_otg(connection)? {
+            MAX_USB_BLOCK_SIZE
+        } else {
+            FLASH_WRITE_SIZE
+        })
+    }
+
+    fn max_ram_block_size(&self, connection: &mut Connection) -> Result<usize, Error> {
+        Ok(if self.connection_is_usb_otg(connection)? {
+            MAX_USB_BLOCK_SIZE
+        } else {
+            MAX_RAM_BLOCK_SIZE
+        })
+    }
+
     fn get_flash_segments<'a>(
         image: &'a dyn FirmwareImage<'a>,
         bootloader: Option<Vec<u8>>,
@@ -110,6 +129,13 @@ impl ReadEFuse for Esp32s2 {
 }
 
 impl Esp32s2 {
+    fn connection_is_usb_otg(&self, connection: &mut Connection) -> Result<bool, Error> {
+        const UARTDEV_BUF_NO: u32 = 0x3FFFFD14; // Address which indicates OTG in use
+        const UARTDEV_BUF_NO_USB_OTG: u32 = 2; // Value of UARTDEV_BUF_NO when OTG is in use
+
+        Ok(connection.read_reg(UARTDEV_BUF_NO)? == UARTDEV_BUF_NO_USB_OTG)
+    }
+
     fn get_flash_version(&self, connection: &mut Connection) -> Result<u32, Error> {
         let blk1_word3 = self.read_efuse(connection, 8)?;
         let flash_version = (blk1_word3 >> 21) & 0xf;

espflash/src/chip/mod.rs

@@ -11,8 +11,8 @@ use crate::{
     connection::Connection,
     elf::{FirmwareImage, FlashFrequency, FlashMode},
     error::ChipDetectError,
-    flash_target::{Esp32Target, Esp8266Target, FlashTarget, RamTarget},
-    flasher::{FlashSize, SpiAttachParams},
+    flash_target::{Esp32Target, Esp8266Target, FlashTarget, RamTarget, MAX_RAM_BLOCK_SIZE},
+    flasher::{FlashSize, SpiAttachParams, FLASH_WRITE_SIZE},
     image_format::{ImageFormat, ImageFormatId},
     Error, PartitionTable,
 };
@@ -91,6 +91,14 @@ pub trait ChipType: ReadEFuse {
     fn supports_target(target: &str) -> bool {
         Self::SUPPORTED_TARGETS.contains(&target)
     }
+
+    fn flash_write_size(&self, _connection: &mut Connection) -> Result<usize, Error> {
+        Ok(FLASH_WRITE_SIZE)
+    }
+
+    fn max_ram_block_size(&self, _connection: &mut Connection) -> Result<usize, Error> {
+        Ok(MAX_RAM_BLOCK_SIZE)
+    }
 }
 
 pub trait ReadEFuse {
@@ -295,8 +303,12 @@ impl Chip {
         }
     }
 
-    pub fn ram_target(&self, entry: Option<u32>) -> Box<dyn FlashTarget> {
-        Box::new(RamTarget::new(entry))
+    pub fn ram_target(
+        &self,
+        entry: Option<u32>,
+        max_ram_block_size: usize,
+    ) -> Box<dyn FlashTarget> {
+        Box::new(RamTarget::new(entry, max_ram_block_size))
     }
 
     pub fn flash_target(
@@ -408,6 +420,28 @@ impl Chip {
             Chip::Esp8266 => Esp8266::flash_frequency_encodings(),
         }
     }
+
+    pub fn flash_write_size(&self, connection: &mut Connection) -> Result<usize, Error> {
+        match self {
+            Chip::Esp32 => Esp32.flash_write_size(connection),
+            Chip::Esp32c2 => Esp32c2.flash_write_size(connection),
+            Chip::Esp32c3 => Esp32c3.flash_write_size(connection),
+            Chip::Esp32s2 => Esp32s2.flash_write_size(connection),
+            Chip::Esp32s3 => Esp32s3.flash_write_size(connection),
+            Chip::Esp8266 => Esp8266.flash_write_size(connection),
+        }
+    }
+
+    pub fn max_ram_block_size(&self, connection: &mut Connection) -> Result<usize, Error> {
+        match self {
+            Chip::Esp32 => Esp32.max_ram_block_size(connection),
+            Chip::Esp32c2 => Esp32c2.max_ram_block_size(connection),
+            Chip::Esp32c3 => Esp32c3.max_ram_block_size(connection),
+            Chip::Esp32s2 => Esp32s2.max_ram_block_size(connection),
+            Chip::Esp32s3 => Esp32s3.max_ram_block_size(connection),
+            Chip::Esp8266 => Esp8266.max_ram_block_size(connection),
+        }
+    }
 }
 
 pub(crate) fn bytes_to_mac_addr(bytes: &[u8]) -> String {

espflash/src/connection.rs

@@ -290,7 +290,6 @@ pub fn reset_after_flash(serial: &mut dyn SerialPort, pid: u16) -> Result<(), se
 
         serial.write_request_to_send(false)?;
     } else {
-        serial.write_data_terminal_ready(false)?;
         serial.write_request_to_send(true)?;
 
         sleep(Duration::from_millis(100));

espflash/src/flash_target/esp32.rs

@@ -3,7 +3,7 @@ use crate::connection::{Connection, USB_SERIAL_JTAG_PID};
 use crate::elf::RomSegment;
 use crate::error::Error;
 use crate::flash_target::FlashTarget;
-use crate::flasher::{SpiAttachParams, FLASH_SECTOR_SIZE, FLASH_WRITE_SIZE};
+use crate::flasher::{SpiAttachParams, FLASH_SECTOR_SIZE};
 use crate::Chip;
 use flate2::write::{ZlibDecoder, ZlibEncoder};
 use flate2::Compression;
@@ -94,7 +94,8 @@ impl FlashTarget for Esp32Target {
         let mut encoder = ZlibEncoder::new(Vec::new(), Compression::best());
         encoder.write_all(&segment.data)?;
         let compressed = encoder.finish()?;
-        let block_count = (compressed.len() + FLASH_WRITE_SIZE - 1) / FLASH_WRITE_SIZE;
+        let flash_write_size = self.chip.flash_write_size(connection)?;
+        let block_count = (compressed.len() + flash_write_size - 1) / flash_write_size;
         let erase_count = (segment.data.len() + FLASH_SECTOR_SIZE - 1) / FLASH_SECTOR_SIZE;
 
         // round up to sector size
@@ -106,15 +107,15 @@ impl FlashTarget for Esp32Target {
                 connection.command(Command::FlashDeflateBegin {
                     size: erase_size,
                     blocks: block_count as u32,
-                    block_size: FLASH_WRITE_SIZE as u32,
+                    block_size: flash_write_size as u32,
                     offset: addr,
                     supports_encryption: self.chip != Chip::Esp32 && !self.use_stub,
                 })?;
                 Ok(())
             },
         )?;
 
-        let chunks = compressed.chunks(FLASH_WRITE_SIZE);
+        let chunks = compressed.chunks(flash_write_size);
 
         let (_, chunk_size) = chunks.size_hint();
         let chunk_size = chunk_size.unwrap_or(0) as u64;

espflash/src/flash_target/mod.rs

@@ -11,6 +11,8 @@ pub use esp32::Esp32Target;
 pub use esp8266::Esp8266Target;
 pub use ram::RamTarget;
 
+pub(crate) use ram::MAX_RAM_BLOCK_SIZE;
+
 pub trait FlashTarget {
     fn begin(&mut self, connection: &mut Connection) -> Result<(), Error>;
     fn write_segment(

espflash/src/flash_target/ram.rs

@@ -5,6 +5,8 @@ use crate::error::Error;
 use crate::flash_target::FlashTarget;
 use bytemuck::{Pod, Zeroable};
 
+pub(crate) const MAX_RAM_BLOCK_SIZE: usize = 0x1800;
+
 #[derive(Zeroable, Pod, Copy, Clone)]
 #[repr(C)]
 struct EntryParams {
@@ -14,17 +16,18 @@ struct EntryParams {
 
 pub struct RamTarget {
     entry: Option<u32>,
+    block_size: usize,
 }
 
 impl RamTarget {
-    pub fn new(entry: Option<u32>) -> Self {
-        RamTarget { entry }
+    pub fn new(entry: Option<u32>, block_size: usize) -> Self {
+        RamTarget { entry, block_size }
     }
 }
 
 impl Default for RamTarget {
     fn default() -> Self {
-        Self::new(None)
+        Self::new(None, MAX_RAM_BLOCK_SIZE)
     }
 }
 
@@ -38,21 +41,18 @@ impl FlashTarget for RamTarget {
         connection: &mut Connection,
         segment: RomSegment,
     ) -> Result<(), Error> {
-        const MAX_RAM_BLOCK_SIZE: usize = 0x1800;
-
         let padding = 4 - segment.data.len() % 4;
-        let block_count =
-            (segment.data.len() + padding + MAX_RAM_BLOCK_SIZE - 1) / MAX_RAM_BLOCK_SIZE;
+        let block_count = (segment.data.len() + padding + self.block_size - 1) / self.block_size;
 
         connection.command(Command::MemBegin {
             size: segment.data.len() as u32,
             blocks: block_count as u32,
-            block_size: MAX_RAM_BLOCK_SIZE as u32,
+            block_size: self.block_size as u32,
             offset: segment.addr,
             supports_encryption: false,
         })?;
 
-        for (i, block) in segment.data.chunks(MAX_RAM_BLOCK_SIZE).enumerate() {
+        for (i, block) in segment.data.chunks(self.block_size).enumerate() {
             connection.command(Command::MemData {
                 sequence: i as u32,
                 pad_to: 4,

espflash/src/flasher.rs

@@ -257,7 +257,10 @@ impl Flasher {
         // Load flash stub
         let stub = FlashStub::get(self.chip);
 
-        let mut ram_target = self.chip.ram_target(Some(stub.entry()));
+        let mut ram_target = self.chip.ram_target(
+            Some(stub.entry()),
+            self.chip.max_ram_block_size(&mut self.connection)?,
+        );
         ram_target.begin(&mut self.connection).flashing()?;
 
         let (text_addr, text) = stub.text();
@@ -513,7 +516,10 @@ impl Flasher {
     pub fn load_elf_to_ram(&mut self, elf_data: &[u8]) -> Result<(), Error> {
         let image = ElfFirmwareImage::try_from(elf_data)?;
 
-        let mut target = self.chip.ram_target(Some(image.entry()));
+        let mut target = self.chip.ram_target(
+            Some(image.entry()),
+            self.chip.max_ram_block_size(&mut self.connection)?,
+        );
         target.begin(&mut self.connection).flashing()?;
 
         if image.rom_segments(self.chip).next().is_some() {