Merge pull request #42 from esp-rs/flash-targets

Allow for per-chip flash logic and implement compression for 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/chip/mod.rs

@@ -9,6 +9,8 @@ use crate::{
 
 use std::{io::Write, str::FromStr};
 
+use crate::flash_target::{Esp32Target, Esp8266Target, FlashTarget, RamTarget};
+use crate::flasher::SpiAttachParams;
 pub use esp32::Esp32;
 pub use esp32c3::Esp32c3;
 pub use esp8266::Esp8266;
@@ -139,6 +141,17 @@ impl Chip {
             Chip::Esp8266 => Esp8266::SPI_REGISTERS,
         }
     }
+
+    pub fn ram_target(&self) -> Box<dyn FlashTarget> {
+        Box::new(RamTarget::new())
+    }
+
+    pub fn flash_target(&self, spi_params: SpiAttachParams) -> Box<dyn FlashTarget> {
+        match self {
+            Chip::Esp8266 => Box::new(Esp8266Target::new()),
+            _ => Box::new(Esp32Target::new(*self, spi_params)),
+        }
+    }
 }
 
 impl FromStr for Chip {

espflash/src/flash_target/esp32.rs

@@ -0,0 +1,109 @@
+use crate::connection::Connection;
+use crate::elf::{FirmwareImage, RomSegment};
+use crate::error::Error;
+use crate::flash_target::{begin_command, block_command_with_timeout, FlashTarget};
+use crate::flasher::{Command, SpiAttachParams, FLASH_SECTOR_SIZE, FLASH_WRITE_SIZE};
+use crate::Chip;
+use flate2::write::{ZlibDecoder, ZlibEncoder};
+use flate2::Compression;
+use indicatif::{ProgressBar, ProgressStyle};
+use std::io::Write;
+
+pub struct Esp32Target {
+    chip: Chip,
+    spi_attach_params: SpiAttachParams,
+}
+
+impl Esp32Target {
+    pub fn new(chip: Chip, spi_attach_params: SpiAttachParams) -> Self {
+        Esp32Target {
+            chip,
+            spi_attach_params,
+        }
+    }
+}
+
+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)
+        })?;
+        Ok(())
+    }
+
+    fn write_segment(
+        &mut self,
+        connection: &mut Connection,
+        segment: RomSegment,
+    ) -> Result<(), Error> {
+        let addr = segment.addr;
+        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;
+
+        // round up to sector size
+        let erase_size = (erase_count * FLASH_SECTOR_SIZE) as u32;
+
+        begin_command(
+            connection,
+            Command::FlashDeflateBegin,
+            erase_size,
+            block_count as u32,
+            FLASH_WRITE_SIZE as u32,
+            addr,
+            self.chip != Chip::Esp32,
+        )?;
+
+        let chunks = compressed.chunks(FLASH_WRITE_SIZE);
+
+        let (_, chunk_size) = chunks.size_hint();
+        let chunk_size = chunk_size.unwrap_or(0) as u64;
+        let pb_chunk = ProgressBar::new(chunk_size);
+        pb_chunk.set_style(
+            ProgressStyle::default_bar()
+                .template("[{elapsed_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}")
+                .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));
+            block_command_with_timeout(
+                connection,
+                Command::FlashDeflateData,
+                block,
+                0,
+                0xff,
+                i as u32,
+                Command::FlashDeflateData.timeout_for_size(size as u32),
+            )?;
+            pb_chunk.inc(1);
+        }
+
+        pb_chunk.finish_with_message(format!("segment 0x{:X}", addr));
+
+        Ok(())
+    }
+
+    fn finish(&mut self, connection: &mut Connection, reboot: bool) -> Result<(), Error> {
+        connection.with_timeout(Command::FlashDeflateEnd.timeout(), |connection| {
+            connection.write_command(Command::FlashDeflateEnd as u8, &[1][..], 0)
+        })?;
+        if reboot {
+            connection.reset()
+        } else {
+            Ok(())
+        }
+    }
+}

espflash/src/flash_target/esp8266.rs

@@ -0,0 +1,89 @@
+use crate::connection::Connection;
+use crate::elf::{FirmwareImage, RomSegment};
+use crate::error::Error;
+use crate::flash_target::{begin_command, block_command, FlashTarget};
+use crate::flasher::{get_erase_size, Command, FLASH_WRITE_SIZE};
+use indicatif::{ProgressBar, ProgressStyle};
+
+pub struct Esp8266Target;
+
+impl Esp8266Target {
+    pub fn new() -> Self {
+        Esp8266Target
+    }
+}
+
+impl FlashTarget for Esp8266Target {
+    fn begin(&mut self, connection: &mut Connection, _image: &FirmwareImage) -> Result<(), Error> {
+        begin_command(
+            connection,
+            Command::FlashBegin,
+            0,
+            0,
+            FLASH_WRITE_SIZE as u32,
+            0,
+            false,
+        )
+    }
+
+    fn write_segment(
+        &mut self,
+        connection: &mut Connection,
+        segment: RomSegment,
+    ) -> Result<(), Error> {
+        let addr = segment.addr;
+        let block_count = (segment.data.len() + FLASH_WRITE_SIZE - 1) / FLASH_WRITE_SIZE;
+
+        let erase_size = get_erase_size(addr as usize, segment.data.len()) as u32;
+
+        begin_command(
+            connection,
+            Command::FlashBegin,
+            erase_size,
+            block_count as u32,
+            FLASH_WRITE_SIZE as u32,
+            addr,
+            false,
+        )?;
+
+        let chunks = segment.data.chunks(FLASH_WRITE_SIZE);
+
+        let (_, chunk_size) = chunks.size_hint();
+        let chunk_size = chunk_size.unwrap_or(0) as u64;
+        let pb_chunk = ProgressBar::new(chunk_size);
+        pb_chunk.set_style(
+            ProgressStyle::default_bar()
+                .template("[{elapsed_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}")
+                .progress_chars("#>-"),
+        );
+
+        for (i, block) in chunks.enumerate() {
+            pb_chunk.set_message(format!("segment 0x{:X} writing chunks", addr));
+            let block_padding = FLASH_WRITE_SIZE - block.len();
+            block_command(
+                connection,
+                Command::FlashData,
+                block,
+                block_padding,
+                0xff,
+                i as u32,
+            )?;
+            pb_chunk.inc(1);
+        }
+
+        pb_chunk.finish_with_message(format!("segment 0x{:X}", addr));
+
+        Ok(())
+    }
+
+    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)
+        })?;
+        if reboot {
+            connection.reset()
+        } else {
+            Ok(())
+        }
+    }
+}

espflash/src/flash_target/mod.rs

@@ -0,0 +1,135 @@
+mod esp32;
+mod esp8266;
+mod ram;
+
+use crate::connection::Connection;
+use crate::elf::{FirmwareImage, RomSegment};
+use crate::error::Error;
+use crate::flasher::{checksum, Command, Encoder, CHECKSUM_INIT, FLASH_WRITE_SIZE};
+use bytemuck::{bytes_of, Pod, Zeroable};
+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>;
+    fn write_segment(
+        &mut self,
+        connection: &mut Connection,
+        segment: RomSegment,
+    ) -> Result<(), Error>;
+    fn finish(&mut self, connection: &mut Connection, reboot: bool) -> Result<(), Error>;
+}
+
+#[derive(Zeroable, Pod, Copy, Clone, Debug)]
+#[repr(C)]
+struct BeginParams {
+    size: u32,
+    blocks: u32,
+    block_size: u32,
+    offset: u32,
+    encrypted: u32,
+}
+
+fn begin_command(
+    connection: &mut Connection,
+    command: Command,
+    size: u32,
+    blocks: u32,
+    block_size: u32,
+    offset: u32,
+    supports_encrypted: bool,
+) -> Result<(), Error> {
+    let params = BeginParams {
+        size,
+        blocks,
+        block_size,
+        offset,
+        encrypted: 0,
+    };
+
+    let bytes = bytes_of(&params);
+    let data = if !supports_encrypted {
+        // The ESP32 and ESP8266 do not take the `encrypted` field, so truncate the last
+        // 4 bytes of the slice where it resides.
+        let end = bytes.len() - 4;
+        &bytes[0..end]
+    } else {
+        bytes
+    };
+
+    connection.with_timeout(command.timeout_for_size(size), |connection| {
+        connection.command(command as u8, data, 0)?;
+        Ok(())
+    })
+}
+
+#[derive(Zeroable, Pod, Copy, Clone, Debug)]
+#[repr(C)]
+struct BlockParams {
+    size: u32,
+    sequence: u32,
+    dummy1: u32,
+    dummy2: u32,
+}
+
+fn block_command(
+    connection: &mut Connection,
+    command: Command,
+    data: &[u8],
+    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,
+        sequence,
+        dummy1: 0,
+        dummy2: 0,
+    };
+
+    let length = size_of::<BlockParams>() + data.len() + padding;
+
+    let mut check = checksum(data, CHECKSUM_INIT);
+
+    for _ in 0..padding {
+        check = checksum(&[padding_byte], check);
+    }
+
+    connection.with_timeout(timout, |connection| {
+        connection.command(
+            command as u8,
+            (length as u16, |encoder: &mut Encoder| {
+                encoder.write(bytes_of(&params))?;
+                encoder.write(data)?;
+                let padding = &[padding_byte; FLASH_WRITE_SIZE][0..padding];
+                encoder.write(padding)?;
+                Ok(())
+            }),
+            check as u32,
+        )?;
+        Ok(())
+    })
+}