Refactoring, various changes/improvements

- chip detect using magic value instead of date regs
- handle differences in BeginParams for older and newer chips
- extract functions common to ESP32 devices

Author: jessebraham

espflash/bootloader/bootloader.bin renamed to espflash/bootloader/esp32-bootloader.bin

@@ -1,20 +1,20 @@
-use std::borrow::Cow;
-use std::io::Write;
-use std::iter::once;
-
-use crate::chip::esp32::partition_table::PartitionTable;
-use crate::chip::{Chip, ChipType, EspCommonHeader, SegmentHeader, SpiRegisters, ESP_MAGIC};
-use crate::elf::{update_checksum, CodeSegment, FirmwareImage, RomSegment, ESP_CHECKSUM_MAGIC};
-use crate::flasher::FlashSize;
-use crate::Error;
-use bytemuck::{bytes_of, Pod, Zeroable};
+use bytemuck::bytes_of;
 use sha2::{Digest, Sha256};
 
-mod partition_table;
+use crate::{
+    chip::{
+        encode_flash_size, get_segment_padding, save_flash_segment, save_segment, Chip, ChipType,
+        EspCommonHeader, ExtendedHeader, SegmentHeader, SpiRegisters, ESP_MAGIC, SEG_HEADER_LEN,
+        WP_PIN_DISABLED,
+    },
+    elf::{FirmwareImage, RomSegment, ESP_CHECKSUM_MAGIC},
+    partition_table::PartitionTable,
+    Error,
+};
 
-pub struct Esp32;
+use std::{borrow::Cow, io::Write, iter::once};
 
-const WP_PIN_DISABLED: u8 = 0xEE;
+pub struct Esp32;
 
 const IROM_MAP_START: u32 = 0x400d0000;
 const IROM_MAP_END: u32 = 0x40400000;
@@ -26,22 +26,9 @@ const BOOT_ADDR: u32 = 0x1000;
 const PARTION_ADDR: u32 = 0x8000;
 const APP_ADDR: u32 = 0x10000;
 
-#[derive(Copy, Clone, Zeroable, Pod)]
-#[repr(C)]
-struct ExtendedHeader {
-    wp_pin: u8,
-    clk_q_drv: u8,
-    d_cs_drv: u8,
-    gd_wp_drv: u8,
-    chip_id: u16,
-    min_rev: u8,
-    padding: [u8; 8],
-    append_digest: u8,
-}
-
 impl ChipType for Esp32 {
-    const DATE_REG1_VALUE: u32 = 0x15122500;
-    const DATE_REG2_VALUE: u32 = 0;
+    const CHIP_DETECT_MAGIC_VALUE: u32 = 0x00f01d83;
+
     const SPI_REGISTERS: SpiRegisters = SpiRegisters {
         base: 0x3ff42000,
         usr_offset: 0x1c,
@@ -60,7 +47,7 @@ impl ChipType for Esp32 {
     fn get_flash_segments<'a>(
         image: &'a FirmwareImage,
     ) -> Box<dyn Iterator<Item = Result<RomSegment<'a>, Error>> + 'a> {
-        let bootloader = include_bytes!("../../bootloader/bootloader.bin");
+        let bootloader = include_bytes!("../../bootloader/esp32-bootloader.bin");
 
         let partition_table = PartitionTable::basic(0x10000, 0x3f0000).to_bytes();
 
@@ -167,74 +154,6 @@ impl ChipType for Esp32 {
     }
 }
 
-fn encode_flash_size(size: FlashSize) -> Result<u8, Error> {
-    match size {
-        FlashSize::Flash256Kb => Err(Error::UnsupportedFlash(size as u8)),
-        FlashSize::Flash512Kb => Err(Error::UnsupportedFlash(size as u8)),
-        FlashSize::Flash1Mb => Ok(0x00),
-        FlashSize::Flash2Mb => Ok(0x10),
-        FlashSize::Flash4Mb => Ok(0x20),
-        FlashSize::Flash8Mb => Ok(0x30),
-        FlashSize::Flash16Mb => Ok(0x40),
-        FlashSize::FlashRetry => Err(Error::UnsupportedFlash(size as u8)),
-    }
-}
-
-const IROM_ALIGN: u32 = 65536;
-const SEG_HEADER_LEN: u32 = 8;
-
-/// Actual alignment (in data bytes) required for a segment header: positioned
-/// so that after we write the next 8 byte header, file_offs % IROM_ALIGN ==
-/// segment.addr % IROM_ALIGN
-///
-/// (this is because the segment's vaddr may not be IROM_ALIGNed, more likely is
-/// aligned IROM_ALIGN+0x18 to account for the binary file header
-fn get_segment_padding(offset: usize, segment: &CodeSegment) -> u32 {
-    let align_past = (segment.addr % IROM_ALIGN) - SEG_HEADER_LEN;
-    let pad_len = (IROM_ALIGN - ((offset as u32) % IROM_ALIGN)) + align_past;
-    if pad_len == 0 || pad_len == IROM_ALIGN {
-        0
-    } else if pad_len > SEG_HEADER_LEN {
-        pad_len - SEG_HEADER_LEN
-    } else {
-        pad_len + IROM_ALIGN - SEG_HEADER_LEN
-    }
-}
-
-fn save_flash_segment(
-    data: &mut Vec<u8>,
-    segment: &CodeSegment,
-    checksum: u8,
-) -> Result<u8, Error> {
-    let end_pos = (data.len() + segment.data.len()) as u32 + SEG_HEADER_LEN;
-    let segment_reminder = end_pos % IROM_ALIGN;
-
-    let checksum = save_segment(data, segment, checksum)?;
-
-    if segment_reminder < 0x24 {
-        // Work around a bug in ESP-IDF 2nd stage bootloader, that it didn't map the
-        // last MMU page, if an IROM/DROM segment was < 0x24 bytes over the page
-        // boundary.
-        data.write_all(&[0u8; 0x24][0..(0x24 - segment_reminder as usize)])?;
-    }
-    Ok(checksum)
-}
-
-fn save_segment(data: &mut Vec<u8>, segment: &CodeSegment, checksum: u8) -> Result<u8, Error> {
-    let padding = (4 - segment.data.len() % 4) % 4;
-
-    let header = SegmentHeader {
-        addr: segment.addr,
-        length: (segment.data.len() + padding) as u32,
-    };
-    data.write_all(bytes_of(&header))?;
-    data.write_all(segment.data)?;
-    let padding = &[0u8; 4][0..padding];
-    data.write_all(padding)?;
-
-    Ok(update_checksum(segment.data, checksum))
-}
-
 #[test]
 fn test_esp32_rom() {
     use std::fs::read;

espflash/src/chip/esp32.rs

@@ -1,23 +1,23 @@
-use std::borrow::Cow;
-use std::io::Write;
-use std::iter::once;
-use std::mem::size_of;
+use bytemuck::bytes_of;
 
 use super::{ChipType, EspCommonHeader, SegmentHeader, ESP_MAGIC};
-use crate::chip::{Chip, SpiRegisters};
-use crate::elf::{update_checksum, CodeSegment, FirmwareImage, RomSegment, ESP_CHECKSUM_MAGIC};
-use crate::flasher::FlashSize;
-use crate::Error;
-use bytemuck::bytes_of;
+use crate::{
+    chip::{Chip, SpiRegisters},
+    elf::{update_checksum, CodeSegment, FirmwareImage, RomSegment, ESP_CHECKSUM_MAGIC},
+    flasher::FlashSize,
+    Error,
+};
+
+use std::{borrow::Cow, io::Write, iter::once, mem::size_of};
 
 pub const IROM_MAP_START: u32 = 0x40200000;
 const IROM_MAP_END: u32 = 0x40300000;
 
 pub struct Esp8266;
 
 impl ChipType for Esp8266 {
-    const DATE_REG1_VALUE: u32 = 0x00062000;
-    const DATE_REG2_VALUE: u32 = 0;
+    const CHIP_DETECT_MAGIC_VALUE: u32 = 0xfff0c101;
+
     const SPI_REGISTERS: SpiRegisters = SpiRegisters {
         base: 0x60000200,
         usr_offset: 0x1c,

espflash/src/chip/esp8266.rs

@@ -1,19 +1,26 @@
-use crate::elf::{FirmwareImage, RomSegment};
-use crate::Error;
-use bytemuck::{Pod, Zeroable};
-use std::str::FromStr;
+use bytemuck::{bytes_of, Pod, Zeroable};
+
+use crate::{
+    elf::{update_checksum, CodeSegment, FirmwareImage, RomSegment},
+    flasher::FlashSize,
+    Error,
+};
+
+use std::{io::Write, str::FromStr};
 
 pub use esp32::Esp32;
 pub use esp8266::Esp8266;
 
 mod esp32;
 mod esp8266;
 
-const ESP_MAGIC: u8 = 0xe9;
+const ESP_MAGIC: u8 = 0xE9;
+const WP_PIN_DISABLED: u8 = 0xEE;
 
 pub trait ChipType {
-    const DATE_REG1_VALUE: u32;
-    const DATE_REG2_VALUE: u32;
+    const CHIP_DETECT_MAGIC_VALUE: u32;
+    const CHIP_DETECT_MAGIC_VALUE2: u32 = 0x0; // give default value, as most chips don't only have one
+
     const SPI_REGISTERS: SpiRegisters;
 
     /// Get the firmware segments for writing an image to flash
@@ -64,17 +71,30 @@ impl SpiRegisters {
     }
 }
 
-#[derive(Debug, Copy, Clone)]
+#[derive(Copy, Clone, Zeroable, Pod)]
+#[repr(C)]
+struct ExtendedHeader {
+    wp_pin: u8,
+    clk_q_drv: u8,
+    d_cs_drv: u8,
+    gd_wp_drv: u8,
+    chip_id: u16,
+    min_rev: u8,
+    padding: [u8; 8],
+    append_digest: u8,
+}
+
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub enum Chip {
-    Esp8266,
     Esp32,
+    Esp8266,
 }
 
 impl Chip {
-    pub fn from_regs(value1: u32, value2: u32) -> Option<Self> {
-        match (value1, value2) {
-            (Esp8266::DATE_REG1_VALUE, _) => Some(Chip::Esp8266),
-            (Esp32::DATE_REG1_VALUE, _) => Some(Chip::Esp32),
+    pub fn from_magic(magic: u32) -> Option<Self> {
+        match magic {
+            Esp32::CHIP_DETECT_MAGIC_VALUE => Some(Chip::Esp32),
+            Esp8266::CHIP_DETECT_MAGIC_VALUE => Some(Chip::Esp8266),
             _ => None,
         }
     }
@@ -84,22 +104,22 @@ impl Chip {
         image: &'a FirmwareImage,
     ) -> Box<dyn Iterator<Item = Result<RomSegment<'a>, Error>> + 'a> {
         match self {
-            Chip::Esp8266 => Esp8266::get_flash_segments(image),
             Chip::Esp32 => Esp32::get_flash_segments(image),
+            Chip::Esp8266 => Esp8266::get_flash_segments(image),
         }
     }
 
     pub fn addr_is_flash(&self, addr: u32) -> bool {
         match self {
-            Chip::Esp8266 => Esp8266::addr_is_flash(addr),
             Chip::Esp32 => Esp32::addr_is_flash(addr),
+            Chip::Esp8266 => Esp8266::addr_is_flash(addr),
         }
     }
 
     pub fn spi_registers(&self) -> SpiRegisters {
         match self {
-            Chip::Esp8266 => Esp8266::SPI_REGISTERS,
             Chip::Esp32 => Esp32::SPI_REGISTERS,
+            Chip::Esp8266 => Esp8266::SPI_REGISTERS,
         }
     }
 }
@@ -132,3 +152,72 @@ struct SegmentHeader {
     addr: u32,
     length: u32,
 }
+
+fn encode_flash_size(size: FlashSize) -> Result<u8, Error> {
+    match size {
+        FlashSize::Flash256Kb => Err(Error::UnsupportedFlash(size as u8)),
+        FlashSize::Flash512Kb => Err(Error::UnsupportedFlash(size as u8)),
+        FlashSize::Flash1Mb => Ok(0x00),
+        FlashSize::Flash2Mb => Ok(0x10),
+        FlashSize::Flash4Mb => Ok(0x20),
+        FlashSize::Flash8Mb => Ok(0x30),
+        FlashSize::Flash16Mb => Ok(0x40),
+        FlashSize::FlashRetry => Err(Error::UnsupportedFlash(size as u8)),
+    }
+}
+
+const IROM_ALIGN: u32 = 65536;
+const SEG_HEADER_LEN: u32 = 8;
+
+/// Actual alignment (in data bytes) required for a segment header: positioned
+/// so that after we write the next 8 byte header, file_offs % IROM_ALIGN ==
+/// segment.addr % IROM_ALIGN
+///
+/// (this is because the segment's vaddr may not be IROM_ALIGNed, more likely is
+/// aligned IROM_ALIGN+0x18 to account for the binary file header
+fn get_segment_padding(offset: usize, segment: &CodeSegment) -> u32 {
+    let align_past = (segment.addr % IROM_ALIGN) - SEG_HEADER_LEN;
+    let pad_len = (IROM_ALIGN - ((offset as u32) % IROM_ALIGN)) + align_past;
+    if pad_len == 0 || pad_len == IROM_ALIGN {
+        0
+    } else if pad_len > SEG_HEADER_LEN {
+        pad_len - SEG_HEADER_LEN
+    } else {
+        pad_len + IROM_ALIGN - SEG_HEADER_LEN
+    }
+}
+
+fn save_flash_segment(
+    data: &mut Vec<u8>,
+    segment: &CodeSegment,
+    checksum: u8,
+) -> Result<u8, Error> {
+    let end_pos = (data.len() + segment.data.len()) as u32 + SEG_HEADER_LEN;
+    let segment_reminder = end_pos % IROM_ALIGN;
+
+    let checksum = save_segment(data, segment, checksum)?;
+
+    if segment_reminder < 0x24 {
+        // Work around a bug in ESP-IDF 2nd stage bootloader, that it didn't map the
+        // last MMU page, if an IROM/DROM segment was < 0x24 bytes over the page
+        // boundary.
+        data.write_all(&[0u8; 0x24][0..(0x24 - segment_reminder as usize)])?;
+    }
+    Ok(checksum)
+}
+
+fn save_segment(data: &mut Vec<u8>, segment: &CodeSegment, checksum: u8) -> Result<u8, Error> {
+    let padding = (4 - segment.data.len() % 4) % 4;
+
+    let header = SegmentHeader {
+        addr: segment.addr,
+        length: (segment.data.len() + padding) as u32,
+    };
+    data.write_all(bytes_of(&header))?;
+    data.write_all(segment.data)?;
+
+    let padding = &[0u8; 4][0..padding];
+    data.write_all(padding)?;
+
+    Ok(update_checksum(segment.data, checksum))
+}