Support flashing in Secure Download Mode

* With secure download enabled, the stub is not utilized so
compressed flashing is unavailble. In this situation, use the
uncompressed flashing commands.

* If secure download is enabled, prevent flashing data to addresses
below 0x8000 to prevent overwriting the bootloader and bricking the
device

Author: enody-carter

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - esp32p4rc1 support for esp32p4 chips with revision < 3.0 (#922)
 - Support for binaries with INIT_ARRAY sections, which is needed for esp32p4 support. (#991)
 - Add sha256 calculation to match esptool generated binaries (#991)
+- Support flashing in secure download mode (#990)
 
 ### Changed
 

espflash/src/error.rs

@@ -383,6 +383,10 @@ pub enum Error {
     /// Tried to use an unsupported eFuse coding scheme
     #[error("Tried to use an unsupported eFuse coding scheme: {0}")]
     UnsupportedEfuseCodingScheme(String),
+
+    /// Tried to write to address < 0x8000 with Secure Download enabled
+    #[error("Tried to write to address < 0x8000 with Secure Download enabled")]
+    SecureDownloadBootloaderProtection,
 }
 
 #[cfg(feature = "serialport")]

espflash/src/flasher/mod.rs

@@ -57,6 +57,9 @@ pub(crate) const TRY_SPI_PARAMS: [SpiAttachParams; 2] =
 pub(crate) const FLASH_SECTOR_SIZE: usize = 0x1000;
 pub(crate) const FLASH_WRITE_SIZE: usize = 0x400;
 
+#[cfg(feature = "serialport")]
+pub(crate) const BOOTLOADER_PROTECTION_ADDR: u32 = 0x8000;
+
 /// Supported flash frequencies
 ///
 /// Note that not all frequencies are supported by each target device.
@@ -936,15 +939,41 @@ impl Flasher {
         target.finish(&mut self.connection, true).flashing()
     }
 
+    /// Validate flash arguments when in secure download mode.
+    /// Prevent a user from accidentally flashing over a secure boot enabled
+    /// bootloader and bricking their device.
+    fn validate_secure_download_args(&self, segments: &[Segment<'_>]) -> Result<(), Error> {
+        for segment in segments {
+            if segment.addr < BOOTLOADER_PROTECTION_ADDR {
+                return Err(Error::SecureDownloadBootloaderProtection);
+            }
+        }
+
+        if self.verify || self.skip {
+            warn!(
+                "Secure Download Mode enabled: --verify and --skip options are not available \
+                (flash read operations are restricted)"
+            );
+        }
+        Ok(())
+    }
+
     /// Load an ELF image to flash and execute it
     pub fn load_image_to_flash<'a>(
         &mut self,
         progress: &mut dyn ProgressCallbacks,
         image_format: ImageFormat<'a>,
     ) -> Result<(), Error> {
-        let mut target =
-            self.chip
-                .flash_target(self.spi_params, self.use_stub, self.verify, self.skip);
+        let (mut verify, mut skip) = (self.verify, self.skip);
+
+        if self.connection.secure_download_mode {
+            self.validate_secure_download_args(&image_format.clone().flash_segments())?;
+            (verify, skip) = (false, false);
+        }
+
+        let mut target = self
+            .chip
+            .flash_target(self.spi_params, self.use_stub, verify, skip);
         target.begin(&mut self.connection).flashing()?;
 
         // When the `cli` feature is enabled, display the image size information.
@@ -1005,16 +1034,16 @@ impl Flasher {
         segments: &[Segment<'_>],
         progress: &mut dyn ProgressCallbacks,
     ) -> Result<(), Error> {
+        let (mut verify, mut skip) = (self.verify, self.skip);
+
         if self.connection.secure_download_mode {
-            return Err(Error::UnsupportedFeature {
-                chip: self.chip,
-                feature: "writing binaries in Secure Download Mode currently".into(),
-            });
+            self.validate_secure_download_args(segments)?;
+            (verify, skip) = (false, false);
         }
 
-        let mut target =
-            self.chip
-                .flash_target(self.spi_params, self.use_stub, self.verify, self.skip);
+        let mut target = self
+            .chip
+            .flash_target(self.spi_params, self.use_stub, verify, skip);
 
         target.begin(&mut self.connection).flashing()?;
 

espflash/src/target/flash_target/esp32.rs

@@ -34,7 +34,7 @@ pub struct Esp32Target {
     use_stub: bool,
     verify: bool,
     skip: bool,
-    need_deflate_end: bool,
+    need_flash_end: bool,
 }
 
 impl Esp32Target {
@@ -52,7 +52,7 @@ impl Esp32Target {
             use_stub,
             verify,
             skip,
-            need_deflate_end: false,
+            need_flash_end: false,
         }
     }
 }
@@ -79,7 +79,10 @@ impl FlashTarget for Esp32Target {
         //
         // TODO: the stub doesn't appear to disable the watchdog on ESP32-S3, so we
         //       explicitly disable the watchdog here.
-        if connection.is_using_usb_serial_jtag() {
+        //
+        // NOTE: In Secure Download Mode, WRITE_REG commands are not allowed, so we
+        // must skip the watchdog disable.
+        if connection.is_using_usb_serial_jtag() && !connection.secure_download_mode {
             if let (Some(wdt_wprotect), Some(wdt_config0)) =
                 (self.chip.wdt_wprotect(), self.chip.wdt_config0())
             {
@@ -116,22 +119,15 @@ impl FlashTarget for Esp32Target {
         md5_hasher.update(&segment.data);
         let checksum_md5 = md5_hasher.finalize();
 
-        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::best());
-        encoder.write_all(&segment.data)?;
-        let compressed = encoder.finish()?;
+        // use compression only when stub is loaded.
+        let use_compression = self.use_stub;
 
         let flash_write_size = self.chip.flash_write_size();
-        let block_count = compressed.len().div_ceil(flash_write_size);
         let erase_count = segment.data.len().div_ceil(FLASH_SECTOR_SIZE);
 
         // round up to sector size
         let erase_size = (erase_count * FLASH_SECTOR_SIZE) as u32;
 
-        let chunks = compressed.chunks(flash_write_size);
-        let num_chunks = chunks.len();
-
-        progress.init(addr, num_chunks);
-
         if self.skip {
             let flash_checksum_md5: u128 = connection.with_timeout(
                 CommandType::FlashMd5.timeout_for_size(segment.data.len() as u32),
@@ -153,43 +149,88 @@ impl FlashTarget for Esp32Target {
             }
         }
 
-        connection.with_timeout(
-            CommandType::FlashDeflBegin.timeout_for_size(erase_size),
-            |connection| {
-                connection.command(Command::FlashDeflBegin {
-                    size: segment.data.len() as u32,
-                    blocks: block_count as u32,
-                    block_size: flash_write_size as u32,
-                    offset: addr,
-                    supports_encryption: self.chip != Chip::Esp32 && !self.use_stub,
-                })?;
-                Ok(())
-            },
-        )?;
-        self.need_deflate_end = true;
+        let data = if use_compression {
+            let mut encoder = ZlibEncoder::new(Vec::new(), Compression::best());
+            encoder.write_all(&segment.data)?;
+            encoder.finish()?
+        } else {
+            segment.data.to_vec()
+        };
 
-        // decode the chunks to see how much data the device will have to save
-        let mut decoder = ZlibDecoder::new(Vec::new());
-        let mut decoded_size = 0;
+        let block_count = data.len().div_ceil(flash_write_size);
+        let chunks = data.chunks(flash_write_size);
+        let num_chunks = chunks.len();
 
-        for (i, block) in chunks.enumerate() {
-            decoder.write_all(block)?;
-            decoder.flush()?;
-            let size = decoder.get_ref().len() - decoded_size;
-            decoded_size = decoder.get_ref().len();
+        progress.init(addr, num_chunks);
 
+        if use_compression {
+            connection.with_timeout(
+                CommandType::FlashDeflBegin.timeout_for_size(erase_size),
+                |connection| {
+                    connection.command(Command::FlashDeflBegin {
+                        size: segment.data.len() as u32,
+                        blocks: block_count as u32,
+                        block_size: flash_write_size as u32,
+                        offset: addr,
+                        supports_encryption: self.chip != Chip::Esp32 && !self.use_stub,
+                    })?;
+                    Ok(())
+                },
+            )?;
+        } else {
             connection.with_timeout(
-                CommandType::FlashDeflData.timeout_for_size(size as u32),
+                CommandType::FlashBegin.timeout_for_size(erase_size),
                 |connection| {
-                    connection.command(Command::FlashDeflData {
-                        sequence: i as u32,
-                        pad_to: 0,
-                        pad_byte: 0xff,
-                        data: block,
+                    connection.command(Command::FlashBegin {
+                        size: erase_size,
+                        blocks: block_count as u32,
+                        block_size: flash_write_size as u32,
+                        offset: addr,
+                        supports_encryption: self.chip != Chip::Esp32,
                     })?;
                     Ok(())
                 },
             )?;
+        }
+        self.need_flash_end = true;
+
+        // decode the chunks to see how much data the device will have to save
+        let mut decoder = ZlibDecoder::new(Vec::new());
+        let mut decoded_size = 0;
+
+        for (i, block) in chunks.enumerate() {
+            if use_compression {
+                decoder.write_all(block)?;
+                decoder.flush()?;
+                let size = decoder.get_ref().len() - decoded_size;
+                decoded_size = decoder.get_ref().len();
+
+                connection.with_timeout(
+                    CommandType::FlashDeflData.timeout_for_size(size as u32),
+                    |connection| {
+                        connection.command(Command::FlashDeflData {
+                            sequence: i as u32,
+                            pad_to: 0,
+                            pad_byte: 0xff,
+                            data: block,
+                        })?;
+                        Ok(())
+                    },
+                )?;
+            } else {
+                connection.with_timeout(
+                    CommandType::FlashData.timeout_for_size(block.len() as u32),
+                    |connection| {
+                        connection.command(Command::FlashData {
+                            sequence: i as u32,
+                            pad_to: flash_write_size,
+                            pad_byte: 0xff,
+                            data: block,
+                        })?;
+                        Ok(())
+                    },
+                )?;
+            }
 
             progress.update(i + 1)
         }
@@ -220,10 +261,16 @@ impl FlashTarget for Esp32Target {
     }
 
     fn finish(&mut self, connection: &mut Connection, reboot: bool) -> Result<(), Error> {
-        if self.need_deflate_end {
-            connection.with_timeout(CommandType::FlashDeflEnd.timeout(), |connection| {
-                connection.command(Command::FlashDeflEnd { reboot: false })
-            })?;
+        if self.need_flash_end {
+            if self.use_stub {
+                connection.with_timeout(CommandType::FlashDeflEnd.timeout(), |connection| {
+                    connection.command(Command::FlashDeflEnd { reboot: false })
+                })?;
+            } else {
+                connection.with_timeout(CommandType::FlashEnd.timeout(), |connection| {
+                    connection.command(Command::FlashEnd { reboot: false })
+                })?;
+            }
         }
 
         if reboot {