esp32 image format: allow override of header params

* Pull the default header from the esp-idf generated bootloader.
* Modify the header if required.
* Use the same header for the app, with the entry point changed.

Author: MabezDev

espflash/src/elf.rs

@@ -79,17 +79,17 @@ impl FromStr for FlashFrequency {
 pub struct FirmwareImage<'a> {
     pub entry: u32,
     pub elf: ElfFile<'a>,
-    pub flash_mode: FlashMode,
-    pub flash_size: FlashSize,
-    pub flash_frequency: FlashFrequency,
+    pub flash_mode: Option<FlashMode>,
+    pub flash_size: Option<FlashSize>,
+    pub flash_frequency: Option<FlashFrequency>,
 }
 
 impl<'a> FirmwareImage<'a> {
     pub fn new(
         elf: ElfFile<'a>,
-        flash_mode: FlashMode,
-        flash_size: FlashSize,
-        flash_frequency: FlashFrequency,
+        flash_mode: Option<FlashMode>,
+        flash_size: Option<FlashSize>,
+        flash_frequency: Option<FlashFrequency>,
     ) -> Self {
         Self {
             entry: elf.header.pt2.entry_point() as u32,
@@ -184,11 +184,7 @@ impl<'a> FirmwareImageBuilder<'a> {
     pub fn build(&self) -> Result<FirmwareImage<'a>, Error> {
         let elf = ElfFile::new(self.data).map_err(ElfError::from)?;
 
-        let flash_mode = self.flash_mode.unwrap_or(FlashMode::Dio);
-        let flash_size = self.flash_size.unwrap_or(FlashSize::Flash4Mb);
-        let flash_freq = self.flash_freq.unwrap_or(FlashFrequency::Flash40M);
-
-        let image = FirmwareImage::new(elf, flash_mode, flash_size, flash_freq);
+        let image = FirmwareImage::new(elf, self.flash_mode, self.flash_size, self.flash_freq);
 
         Ok(image)
     }

espflash/src/image_format/esp32bootloader.rs

@@ -1,6 +1,6 @@
 use std::{borrow::Cow, io::Write, iter::once};
 
-use bytemuck::{bytes_of, Pod, Zeroable, from_bytes};
+use bytemuck::{bytes_of, from_bytes, Pod, Zeroable};
 use sha2::{Digest, Sha256};
 
 use crate::{
@@ -41,22 +41,41 @@ impl<'a> Esp32BootloaderFormat<'a> {
 
         let mut data = Vec::new();
 
-        let header = EspCommonHeader {
-            magic: ESP_MAGIC,
-            segment_count: 0,
-            flash_mode: image.flash_mode as u8,
-            flash_config: encode_flash_size(image.flash_size)? + image.flash_frequency as u8,
-            entry: image.entry,
-        };
-
-        // Update the bootloader header
-        let current: EspCommonHeader = *from_bytes(&bootloader[0..8]);
-        if current.magic != ESP_MAGIC {
+        // fetch the generated header from the bootloader
+        let mut header: EspCommonHeader = *from_bytes(&bootloader[0..8]);
+        if header.magic != ESP_MAGIC {
             return Err(Error::InvalidBootloader);
         }
-        bootloader.to_mut()[2..4].copy_from_slice(&bytes_of(&header)[2..4]);
+        // update the header if a user has specified any custom arguments
+        if image.flash_frequency.is_some()
+            || image.flash_mode.is_some()
+            || image.flash_size.is_some()
+        {
+            if let Some(mode) = image.flash_mode {
+                header.flash_mode = mode as u8;
+                bootloader.to_mut()[2] = bytes_of(&header)[2];
+            }
+            match (image.flash_size, image.flash_frequency) {
+                (Some(s), Some(f)) => {
+                    header.flash_config = encode_flash_size(s)? + f as u8;
+                    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) + f as u8;
+                    bootloader.to_mut()[3] = bytes_of(&header)[3];
+                }
+                (None, None) => {} // nothing to update
+            }
+        }
 
         // 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 {

espflash/src/image_format/esp8266.rs

@@ -7,7 +7,7 @@ use crate::{
     error::{Error, FlashDetectError},
     flasher::FlashSize,
     image_format::{EspCommonHeader, ImageFormat, SegmentHeader, ESP_MAGIC},
-    Chip,
+    Chip, FlashFrequency, FlashMode,
 };
 
 const IROM_MAP_START: u32 = 0x40200000;
@@ -33,8 +33,9 @@ impl<'a> Esp8266Format<'a> {
         let header = EspCommonHeader {
             magic: ESP_MAGIC,
             segment_count: image.ram_segments(Chip::Esp8266).count() as u8,
-            flash_mode: image.flash_mode as u8,
-            flash_config: encode_flash_size(image.flash_size)? + image.flash_frequency as u8,
+            flash_mode: image.flash_mode.unwrap_or(FlashMode::Dio) as u8,
+            flash_config: encode_flash_size(image.flash_size.unwrap_or(FlashSize::Flash4Mb))?
+                + image.flash_frequency.unwrap_or(FlashFrequency::Flash40M) as u8,
             entry: image.entry,
         };
         common_data.write_all(bytes_of(&header))?;