Image header improvements and bugfixes (#375)

* Image header improvements and bugfixes

* Review fixes

Author: maxwase

espflash/src/flasher/mod.rs

@@ -68,6 +68,21 @@ pub enum FlashFrequency {
     _80Mhz,
 }
 
+impl FlashFrequency {
+    /// Encodes flash frequency into the format used by the bootloader.
+    pub fn encode_flash_frequency(self: FlashFrequency, chip: Chip) -> Result<u8, Error> {
+        let encodings = chip.into_target().flash_frequency_encodings();
+        if let Some(&f) = encodings.get(&self) {
+            Ok(f)
+        } else {
+            Err(Error::UnsupportedFlashFrequency {
+                chip,
+                frequency: self,
+            })
+        }
+    }
+}
+
 /// Supported flash modes
 #[cfg_attr(feature = "cli", derive(clap::ValueEnum))]
 #[derive(Copy, Clone, Debug, Default, EnumVariantNames)]
@@ -93,50 +108,95 @@ pub enum FlashMode {
 #[non_exhaustive]
 #[repr(u8)]
 #[strum(serialize_all = "SCREAMING_SNAKE_CASE")]
+#[doc(alias("esp_image_flash_size_t"))]
 pub enum FlashSize {
     /// 256 KB
-    _256Kb = 0x12,
+    _256Kb,
     /// 512 KB
-    _512Kb = 0x13,
+    _512Kb,
     /// 1 MB
-    _1Mb = 0x14,
+    _1Mb,
     /// 2 MB
-    _2Mb = 0x15,
+    _2Mb,
     /// 4 MB
     #[default]
-    _4Mb = 0x16,
+    _4Mb,
     /// 8 MB
-    _8Mb = 0x17,
+    _8Mb,
     /// 16 MB
-    _16Mb = 0x18,
+    _16Mb,
     /// 32 MB
-    _32Mb = 0x19,
+    _32Mb,
     /// 64 MB
-    _64Mb = 0x1a,
+    _64Mb,
     /// 128 MB
-    _128Mb = 0x21,
+    _128Mb,
+    /// 256 MB
+    _256Mb,
 }
 
 impl FlashSize {
-    /// Create a `FlashSize` from an `u8`
-    fn from(value: u8) -> Result<FlashSize, Error> {
+    /// Encodes flash size into the format used by the bootloader.
+    ///
+    /// ## Values:
+    /// * [Esp8266](https://docs.espressif.com/projects/esptool/en/latest/esp8266/advanced-topics/firmware-image-format.html#file-header)
+    /// * [Others](https://docs.espressif.com/projects/esptool/en/latest/esp32s3/advanced-topics/firmware-image-format.html#file-header)
+    pub const fn encode_flash_size(self: FlashSize, chip: Chip) -> Result<u8, Error> {
+        use FlashSize::*;
+
+        let encoded = match chip {
+            Chip::Esp8266 => match self {
+                _256Kb => 1,
+                _512Kb => 0,
+                _1Mb => 2,
+                _2Mb => 3,
+                _4Mb => 4,
+                // Currently not supported
+                // _2Mb_c1 => 5,
+                // _4Mb_c1 => 6,
+                _8Mb => 8,
+                _16Mb => 9,
+                _ => return Err(Error::UnsupportedFlash(self as u8)),
+            },
+            _ => match self {
+                _1Mb => 0,
+                _2Mb => 1,
+                _4Mb => 2,
+                _8Mb => 3,
+                _16Mb => 4,
+                _32Mb => 5,
+                _64Mb => 6,
+                _128Mb => 7,
+                _256Mb => 8,
+                _ => return Err(Error::UnsupportedFlash(self as u8)),
+            },
+        };
+
+        Ok(encoded)
+    }
+
+    /// Create a [FlashSize] from an [u8]
+    ///
+    /// [source](https://github.com/espressif/esptool/blob/f4d2510e2c897621884f433ef3f191e8fc5ff184/esptool/cmds.py#L42)
+    const fn from_detected(value: u8) -> Result<FlashSize, Error> {
         match value {
-            0x12 => Ok(FlashSize::_256Kb),
-            0x13 => Ok(FlashSize::_512Kb),
-            0x14 => Ok(FlashSize::_1Mb),
-            0x15 => Ok(FlashSize::_2Mb),
-            0x16 => Ok(FlashSize::_4Mb),
-            0x17 => Ok(FlashSize::_8Mb),
-            0x18 => Ok(FlashSize::_16Mb),
-            0x19 => Ok(FlashSize::_32Mb),
-            0x1a => Ok(FlashSize::_64Mb),
-            0x21 => Ok(FlashSize::_128Mb),
+            0x12 | 0x32 => Ok(FlashSize::_256Kb),
+            0x13 | 0x33 => Ok(FlashSize::_512Kb),
+            0x14 | 0x34 => Ok(FlashSize::_1Mb),
+            0x15 | 0x35 => Ok(FlashSize::_2Mb),
+            0x16 | 0x36 => Ok(FlashSize::_4Mb),
+            0x17 | 0x37 => Ok(FlashSize::_8Mb),
+            0x18 | 0x38 => Ok(FlashSize::_16Mb),
+            0x19 | 0x39 => Ok(FlashSize::_32Mb),
+            0x20 | 0x1A | 0x3A => Ok(FlashSize::_64Mb),
+            0x21 | 0x1B => Ok(FlashSize::_128Mb),
+            0x22 | 0x1C => Ok(FlashSize::_256Mb),
             _ => Err(Error::UnsupportedFlash(value)),
         }
     }
 
-    /// Return the flash size in bytes
-    pub fn size(self) -> u32 {
+    /// Returns the flash size in bytes
+    pub const fn size(self) -> u32 {
         match self {
             FlashSize::_256Kb => 0x0040000,
             FlashSize::_512Kb => 0x0080000,
@@ -148,6 +208,7 @@ impl FlashSize {
             FlashSize::_32Mb => 0x2000000,
             FlashSize::_64Mb => 0x4000000,
             FlashSize::_128Mb => 0x8000000,
+            FlashSize::_256Mb => 0x10000000,
         }
     }
 }
@@ -450,15 +511,15 @@ impl Flasher {
             return Ok(None);
         }
 
-        let flash_size = match FlashSize::from(size_id) {
+        let flash_size = match FlashSize::from_detected(size_id) {
             Ok(size) => size,
             Err(_) => {
                 warn!(
                     "Could not detect flash size (FlashID=0x{:02X}, SizeID=0x{:02X}), defaulting to 4MB",
                     flash_id,
                     size_id
                 );
-                FlashSize::_4Mb
+                FlashSize::default()
             }
         };
 

espflash/src/image_format/esp8266.rs

@@ -7,8 +7,7 @@ use crate::{
     error::Error,
     flasher::{FlashFrequency, FlashMode, FlashSize},
     image_format::{
-        encode_flash_frequency, update_checksum, EspCommonHeader, ImageFormat, SegmentHeader,
-        ESP_CHECKSUM_MAGIC, ESP_MAGIC,
+        update_checksum, ImageFormat, ImageHeader, SegmentHeader, ESP_CHECKSUM_MAGIC, ESP_MAGIC,
     },
     targets::Chip,
 };
@@ -39,22 +38,25 @@ impl<'a> Esp8266Format<'a> {
 
         // Common header
         let flash_mode = flash_mode.unwrap_or_default() as u8;
-        let flash_freq = flash_freq.unwrap_or_default();
-        let flash_size = flash_size.unwrap_or_default();
-        let flash_config =
-            encode_flash_size(flash_size)? + encode_flash_frequency(Chip::Esp8266, flash_freq)?;
         let segment_count = image.ram_segments(Chip::Esp8266).count() as u8;
 
-        let header = EspCommonHeader {
+        let mut header = ImageHeader {
             magic: ESP_MAGIC,
             segment_count,
             flash_mode,
-            flash_config,
             entry: image.entry(),
+            ..Default::default()
         };
-        common_data.write_all(bytes_of(&header))?;
+        header.write_flash_config(
+            flash_size.unwrap_or_default(),
+            flash_freq.unwrap_or_default(),
+            Chip::Esp8266,
+        )?;
+
+        // Esp8266 does not have extended header
+        let mut total_len = ImageHeader::COMMON_HEADER_LEN;
+        common_data.write_all(&bytes_of(&header)[0..total_len as usize])?;
 
-        let mut total_len = 8;
         let mut checksum = ESP_CHECKSUM_MAGIC;
 
         for segment in image.ram_segments(Chip::Esp8266) {
@@ -163,21 +165,6 @@ fn merge_rom_segments<'a>(
     })
 }
 
-fn encode_flash_size(size: FlashSize) -> Result<u8, Error> {
-    use FlashSize::*;
-
-    match size {
-        _256Kb => Ok(0x10),
-        _512Kb => Ok(0x00),
-        _1Mb => Ok(0x20),
-        _2Mb => Ok(0x30),
-        _4Mb => Ok(0x40),
-        _8Mb => Ok(0x80),
-        _16Mb => Ok(0x90),
-        _ => Err(Error::UnsupportedFlash(size as u8)),
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use std::fs;

espflash/src/image_format/idf_bootloader.rs

@@ -1,6 +1,6 @@
-use std::{borrow::Cow, io::Write, iter::once};
+use std::{borrow::Cow, io::Write, iter::once, mem::size_of};
 
-use bytemuck::{bytes_of, from_bytes, Pod, Zeroable};
+use bytemuck::{bytes_of, from_bytes};
 use esp_idf_part::{PartitionTable, Type};
 use sha2::{Digest, Sha256};
 
@@ -9,29 +9,15 @@ use crate::{
     error::Error,
     flasher::{FlashFrequency, FlashMode, FlashSize},
     image_format::{
-        encode_flash_frequency, update_checksum, EspCommonHeader, ImageFormat, SegmentHeader,
-        ESP_CHECKSUM_MAGIC, ESP_MAGIC, WP_PIN_DISABLED,
+        update_checksum, ImageFormat, ImageHeader, SegmentHeader, ESP_CHECKSUM_MAGIC, ESP_MAGIC,
+        WP_PIN_DISABLED,
     },
     targets::{Chip, Esp32Params},
 };
 
 const IROM_ALIGN: u32 = 0x10000;
 const SEG_HEADER_LEN: u32 = 8;
 
-/// Wrapper for the extended header used by the ESP-IDF bootloader
-#[derive(Debug, Default, Clone, Copy, Pod, Zeroable)]
-#[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,
-}
-
 /// Image format for ESP32 family chips using the second-stage bootloader from
 /// ESP-IDF
 pub struct IdfBootloaderFormat<'a> {
@@ -62,52 +48,38 @@ impl<'a> IdfBootloaderFormat<'a> {
             Cow::Borrowed(params.default_bootloader)
         };
 
-        let mut data = Vec::new();
-
         // fetch the generated header from the bootloader
-        let mut header: EspCommonHeader = *from_bytes(&bootloader[0..8]);
+        let mut header: ImageHeader = *from_bytes(&bootloader[0..size_of::<ImageHeader>()]);
         if header.magic != ESP_MAGIC {
             return Err(Error::InvalidBootloader);
         }
 
         // update the header if a user has specified any custom arguments
         if let Some(mode) = flash_mode {
             header.flash_mode = mode as u8;
-            bootloader.to_mut()[2] = bytes_of(&header)[2];
         }
 
-        match (flash_size, flash_freq) {
-            (Some(s), Some(f)) => {
-                header.flash_config = encode_flash_size(s)? + encode_flash_frequency(chip, f)?;
-                bootloader.to_mut()[3] = bytes_of(&header)[3];
-            }
-            (Some(s), None) => {
-                header.flash_config = encode_flash_size(s)? + (header.flash_config & 0x0F);
-                bootloader.to_mut()[3] = bytes_of(&header)[3];
-            }
-            (None, Some(f)) => {
-                header.flash_config =
-                    (header.flash_config & 0xF0) + encode_flash_frequency(chip, f)?;
-                bootloader.to_mut()[3] = bytes_of(&header)[3];
-            }
-            (None, None) => {} // nothing to update
-        }
+        header.write_flash_config(
+            flash_size.unwrap_or_default(),
+            flash_freq.unwrap_or_default(),
+            chip,
+        )?;
+
+        bootloader.to_mut().splice(
+            0..size_of::<ImageHeader>(),
+            bytes_of(&header).iter().copied(),
+        );
 
         // write the header of the app
         // use the same settings as the bootloader
         // just update the entry point
         header.entry = image.entry();
-        data.write_all(bytes_of(&header))?;
-
-        let extended_header = ExtendedHeader {
-            wp_pin: WP_PIN_DISABLED,
-            chip_id: params.chip_id,
-            append_digest: 1,
 
-            ..ExtendedHeader::default()
-        };
+        header.wp_pin = WP_PIN_DISABLED;
+        header.chip_id = params.chip_id;
+        header.append_digest = 1;
 
-        data.write_all(bytes_of(&extended_header))?;
+        let mut data = bytes_of(&header).to_vec();
 
         let flash_segments: Vec<_> = merge_adjacent_segments(image.rom_segments(chip).collect());
         let mut ram_segments: Vec<_> = merge_adjacent_segments(image.ram_segments(chip).collect());
@@ -240,23 +212,6 @@ impl<'a> ImageFormat<'a> for IdfBootloaderFormat<'a> {
     }
 }
 
-/// Enconde the flash size into the format used by the bootloader.
-fn encode_flash_size(size: FlashSize) -> Result<u8, Error> {
-    use FlashSize::*;
-
-    match size {
-        _1Mb => Ok(0x00),
-        _2Mb => Ok(0x10),
-        _4Mb => Ok(0x20),
-        _8Mb => Ok(0x30),
-        _16Mb => Ok(0x40),
-        _32Mb => Ok(0x19),
-        _64Mb => Ok(0x1a),
-        _128Mb => Ok(0x21),
-        _ => Err(Error::UnsupportedFlash(size as u8)),
-    }
-}
-
 /// 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