implement compression for flashing to esp32

Author: icewind1991

espflash/Cargo.toml

@@ -32,6 +32,7 @@ strum = "0.21.0"
 strum_macros = "0.21.1"
 csv = "1.1.6"
 regex = "1.5.4"
+flate2 = "1"
 
 [dev-dependencies]
 pretty_assertions = "0.7.1"

espflash/src/flasher.rs

@@ -12,7 +12,7 @@ use crate::{
 
 pub(crate) type Encoder<'a> = SlipEncoder<'a, Box<dyn SerialPort>>;
 
-const FLASH_SECTOR_SIZE: usize = 0x1000;
+pub(crate) const FLASH_SECTOR_SIZE: usize = 0x1000;
 const FLASH_BLOCK_SIZE: usize = 0x100;
 const FLASH_SECTORS_PER_BLOCK: usize = FLASH_SECTOR_SIZE / FLASH_BLOCK_SIZE;
 pub(crate) const FLASH_WRITE_SIZE: usize = 0x400;
@@ -42,6 +42,10 @@ pub(crate) enum Command {
     SpiSetParams = 0x0B,
     SpiAttach = 0x0D,
     ChangeBaud = 0x0F,
+    FlashDeflateBegin = 0x10,
+    FlashDeflateData = 0x11,
+    FlashDeflateEnd = 0x12,
+    FlashMd5 = 0x13,
 }
 
 impl Command {
@@ -63,7 +67,9 @@ impl Command {
         }
         match self {
             Command::FlashBegin => calc_timeout(ERASE_REGION_TIMEOUT_PER_MB, size),
-            Command::FlashData => calc_timeout(ERASE_WRITE_TIMEOUT_PER_MB, size),
+            Command::FlashData | Command::FlashDeflateData => {
+                calc_timeout(ERASE_WRITE_TIMEOUT_PER_MB, size)
+            }
             _ => self.timeout(),
         }
     }

espflash/src/flashtarget/esp32.rs

@@ -1,10 +1,13 @@
 use crate::connection::Connection;
 use crate::elf::{FirmwareImage, RomSegment};
 use crate::error::Error;
-use crate::flasher::{get_erase_size, Command, SpiAttachParams, FLASH_WRITE_SIZE};
-use crate::flashtarget::{begin_command, block_command, FlashTarget};
+use crate::flasher::{Command, SpiAttachParams, FLASH_SECTOR_SIZE, FLASH_WRITE_SIZE};
+use crate::flashtarget::{begin_command, block_command_with_timeout, FlashTarget};
 use crate::Chip;
+use flate2::write::{ZlibDecoder, ZlibEncoder};
+use flate2::Compression;
 use indicatif::{ProgressBar, ProgressStyle};
+use std::io::Write;
 
 pub struct Esp32Target {
     chip: Chip,
@@ -35,21 +38,26 @@ impl FlashTarget for Esp32Target {
         segment: RomSegment,
     ) -> Result<(), Error> {
         let addr = segment.addr;
-        let block_count = (segment.data.len() + FLASH_WRITE_SIZE - 1) / FLASH_WRITE_SIZE;
+        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 erase_count = (segment.data.len() + FLASH_SECTOR_SIZE - 1) / FLASH_SECTOR_SIZE;
 
-        let erase_size = segment.data.len() as u32;
+        // round up to sector size
+        let erase_size = (erase_count * FLASH_SECTOR_SIZE) as u32;
 
         begin_command(
             connection,
-            Command::FlashBegin,
+            Command::FlashDeflateBegin,
             erase_size,
             block_count as u32,
             FLASH_WRITE_SIZE as u32,
             addr,
             self.chip != Chip::Esp32,
         )?;
 
-        let chunks = segment.data.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;
@@ -60,16 +68,25 @@ impl FlashTarget for Esp32Target {
                 .progress_chars("#>-"),
         );
 
+        // 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() {
+            decoder.write_all(block)?;
+            decoder.flush()?;
+            let size = decoder.get_ref().len() - decoded_size;
+            decoded_size = decoder.get_ref().len();
+
             pb_chunk.set_message(format!("segment 0x{:X} writing chunks", addr));
-            let block_padding = FLASH_WRITE_SIZE - block.len();
-            block_command(
+            block_command_with_timeout(
                 connection,
-                Command::FlashData,
+                Command::FlashDeflateData,
                 block,
-                block_padding,
+                0,
                 0xff,
                 i as u32,
+                Command::FlashDeflateData.timeout_for_size(size as u32),
             )?;
             pb_chunk.inc(1);
         }
@@ -80,8 +97,8 @@ impl FlashTarget for Esp32Target {
     }
 
     fn finish(&mut self, connection: &mut Connection, reboot: bool) -> Result<(), Error> {
-        connection.with_timeout(Command::FlashEnd.timeout(), |connection| {
-            connection.write_command(Command::FlashEnd as u8, &[1][..], 0)
+        connection.with_timeout(Command::FlashDeflateEnd.timeout(), |connection| {
+            connection.write_command(Command::FlashDeflateEnd as u8, &[1][..], 0)
         })?;
         if reboot {
             connection.reset()

espflash/src/flashtarget/mod.rs

@@ -11,6 +11,7 @@ pub use esp32::Esp32Target;
 pub use esp8266::Esp8266Target;
 pub use ram::RamTarget;
 use std::mem::size_of;
+use std::time::Duration;
 
 pub trait FlashTarget {
     fn begin(&mut self, connection: &mut Connection, image: &FirmwareImage) -> Result<(), Error>;
@@ -81,6 +82,26 @@ fn block_command(
     padding: usize,
     padding_byte: u8,
     sequence: u32,
+) -> Result<(), Error> {
+    block_command_with_timeout(
+        connection,
+        command,
+        data,
+        padding,
+        padding_byte,
+        sequence,
+        command.timeout_for_size(data.len() as u32),
+    )
+}
+
+fn block_command_with_timeout(
+    connection: &mut Connection,
+    command: Command,
+    data: &[u8],
+    padding: usize,
+    padding_byte: u8,
+    sequence: u32,
+    timout: Duration,
 ) -> Result<(), Error> {
     let params = BlockParams {
         size: (data.len() + padding) as u32,
@@ -97,7 +118,7 @@ fn block_command(
         check = checksum(&[padding_byte], check);
     }
 
-    connection.with_timeout(command.timeout_for_size(data.len() as u32), |connection| {
+    connection.with_timeout(timout, |connection| {
         connection.command(
             command as u8,
             (length as u16, |encoder: &mut Encoder| {