add support for flashing direct-boot images to esp32c3

Author: icewind1991

cargo-espflash/src/main.rs

@@ -8,7 +8,7 @@ use std::{
 use cargo_metadata::Message;
 use clap::{App, Arg, ArgMatches, SubCommand};
 use error::Error;
-use espflash::{Chip, Config, FirmwareImage, Flasher, PartitionTable};
+use espflash::{Chip, Config, FirmwareImage, Flasher, ImageFormatId, PartitionTable};
 use miette::{IntoDiagnostic, Result, WrapErr};
 use monitor::monitor;
 use package_metadata::CargoEspFlashMeta;
@@ -17,6 +17,7 @@ use serial::{BaudRate, FlowControl, SerialPort};
 use crate::cargo_config::CargoConfig;
 use crate::error::NoTargetError;
 use crate::{cargo_config::parse_cargo_config, error::UnsupportedTargetError};
+use std::str::FromStr;
 
 mod cargo_config;
 mod error;
@@ -42,6 +43,11 @@ fn main() -> Result<()> {
             .takes_value(true)
             .value_name("FEATURES")
             .help("Comma delimited list of build features"),
+        Arg::with_name("format")
+            .long("format")
+            .takes_value(true)
+            .value_name("image format")
+            .help("Image format to flash"),
     ];
     let connect_args = [Arg::with_name("serial")
         .takes_value(true)
@@ -228,12 +234,18 @@ fn flash(
         None
     };
 
+    let image_format = matches
+        .value_of("format")
+        .map(ImageFormatId::from_str)
+        .transpose()
+        .into_diagnostic()?;
+
     // Read the ELF data from the build path and load it to the target.
     let elf_data = fs::read(path).into_diagnostic()?;
     if matches.is_present("ram") {
         flasher.load_elf_to_ram(&elf_data)?;
     } else {
-        flasher.load_elf_to_flash(&elf_data, bootloader, partition_table)?;
+        flasher.load_elf_to_flash(&elf_data, bootloader, partition_table, image_format)?;
     }
     println!("\nFlashing has completed!");
 
@@ -266,13 +278,6 @@ fn build(
     cargo_config: &CargoConfig,
     chip: Option<Chip>,
 ) -> Result<PathBuf> {
-    // The 'build-std' unstable cargo feature is required to enable
-    // cross-compilation. If it has not been set then we cannot build the
-    // application.
-    if !cargo_config.has_build_std() {
-        return Err(Error::NoBuildStd.into());
-    };
-
     let target = cargo_config
         .target()
         .ok_or_else(|| NoTargetError::new(chip))?;
@@ -281,6 +286,13 @@ fn build(
             return Err(Error::UnsupportedTarget(UnsupportedTargetError::new(target, chip)).into());
         }
     }
+    // The 'build-std' unstable cargo feature is required to enable
+    // cross-compilation for xtensa targets.
+    // If it has not been set then we cannot build the
+    // application.
+    if !cargo_config.has_build_std() && target.starts_with("xtensa-") {
+        return Err(Error::NoBuildStd.into());
+    };
 
     // Build the list of arguments to pass to 'cargo build'.
     let mut args = vec![];
@@ -370,7 +382,13 @@ fn save_image(
 
     let image = FirmwareImage::from_data(&elf_data)?;
 
-    let flash_image = chip.get_flash_image(&image, None, None, None)?;
+    let image_format = matches
+        .value_of("format")
+        .map(ImageFormatId::from_str)
+        .transpose()
+        .into_diagnostic()?;
+
+    let flash_image = chip.get_flash_image(&image, None, None, image_format)?;
     let parts: Vec<_> = flash_image.ota_segments().collect();
 
     let out_path = matches.value_of("file").unwrap();

espflash/src/chip/esp32/esp32c3.rs

@@ -1,6 +1,7 @@
 use std::ops::Range;
 
 use super::Esp32Params;
+use crate::image_format::Esp32DirectBootFormat;
 use crate::{
     chip::{bytes_to_mac_addr, ChipType, ReadEFuse, SpiRegisters},
     connection::Connection,
@@ -79,9 +80,7 @@ impl ChipType for Esp32c3 {
                 partition_table,
                 bootloader,
             )?)),
-            ImageFormatId::DirectBoot => {
-                todo!()
-            }
+            ImageFormatId::DirectBoot => Ok(Box::new(Esp32DirectBootFormat::new(image)?)),
         }
     }
 

espflash/src/elf.rs

@@ -6,6 +6,7 @@ use crate::error::{ElfError, Error};
 use crate::flasher::FlashSize;
 use std::fmt::{Debug, Formatter};
 use std::mem::take;
+use std::ops::AddAssign;
 use xmas_elf::sections::{SectionData, ShType};
 use xmas_elf::ElfFile;
 
@@ -88,7 +89,7 @@ impl<'a> FirmwareImage<'a> {
     }
 }
 
-#[derive(Eq, Clone)]
+#[derive(Eq, Clone, Default)]
 /// A segment of code from the source elf
 pub struct CodeSegment<'a> {
     pub addr: u32,
@@ -97,21 +98,12 @@ pub struct CodeSegment<'a> {
 
 impl<'a> CodeSegment<'a> {
     pub fn new(addr: u32, data: &'a [u8]) -> Self {
-        // pad to 4 byte
-        let padding = (4 - data.len() % 4) % 4;
-        if padding == 0 {
-            CodeSegment {
-                addr,
-                data: Cow::Borrowed(data),
-            }
-        } else {
-            let mut data = data.to_vec();
-            data.extend_from_slice(&[0; 4][0..padding]);
-            CodeSegment {
-                addr,
-                data: Cow::Owned(data),
-            }
-        }
+        let mut segment = CodeSegment {
+            addr,
+            data: Cow::Borrowed(data),
+        };
+        segment.pad_align(4);
+        segment
     }
 
     /// Split of the first `count` bytes into a new segment, adjusting the remaining segment as needed
@@ -142,19 +134,41 @@ impl<'a> CodeSegment<'a> {
         }
     }
 
-    pub fn add(&mut self, extend: &[u8]) {
-        let mut data = take(&mut self.data).into_owned();
-        data.extend_from_slice(extend);
-        self.data = Cow::Owned(data);
-    }
-
     pub fn size(&self) -> u32 {
         self.data.len() as u32
     }
 
     pub fn data(&self) -> &[u8] {
         self.data.as_ref()
     }
+
+    pub fn pad_align(&mut self, align: usize) {
+        let padding = (align - self.data.len() % align) % align;
+        if padding > 0 {
+            let mut data = take(&mut self.data).into_owned();
+            data.extend_from_slice(&[0; 4][0..padding]);
+            self.data = Cow::Owned(data);
+        }
+    }
+}
+
+impl<'a> AddAssign<&'_ [u8]> for CodeSegment<'a> {
+    fn add_assign(&mut self, rhs: &'_ [u8]) {
+        let mut data = take(&mut self.data).into_owned();
+        data.extend_from_slice(rhs);
+        self.data = Cow::Owned(data);
+    }
+}
+
+impl<'a> AddAssign<&'_ CodeSegment<'_>> for CodeSegment<'a> {
+    fn add_assign(&mut self, rhs: &'_ CodeSegment<'_>) {
+        let mut data = take(&mut self.data).into_owned();
+        // pad or truncate
+        #[allow(clippy::suspicious_op_assign_impl)]
+        data.resize((rhs.addr - self.addr) as usize, 0);
+        data.extend_from_slice(rhs.data());
+        self.data = Cow::Owned(data);
+    }
 }
 
 impl Debug for CodeSegment<'_> {
@@ -184,6 +198,7 @@ impl Ord for CodeSegment<'_> {
     }
 }
 
+#[derive(Clone)]
 /// A segment of data to write to the flash
 pub struct RomSegment<'a> {
     pub addr: u32,
@@ -219,14 +234,14 @@ pub fn update_checksum(data: &[u8], mut checksum: u8) -> u8 {
     checksum
 }
 
-pub fn merge_segments(mut segments: Vec<CodeSegment>) -> Vec<CodeSegment> {
+pub fn merge_adjacent_segments(mut segments: Vec<CodeSegment>) -> Vec<CodeSegment> {
     segments.sort();
 
     let mut merged: Vec<CodeSegment> = Vec::with_capacity(segments.len());
     for segment in segments {
         match merged.last_mut() {
             Some(last) if last.addr + last.size() == segment.addr => {
-                last.add(segment.data());
+                *last += segment.data();
             }
             _ => {
                 merged.push(segment);

espflash/src/flasher.rs

@@ -4,6 +4,7 @@ use bytemuck::{__core::time::Duration, bytes_of, Pod, Zeroable};
 use serial::{BaudRate, SystemPort};
 use strum_macros::Display;
 
+use crate::image_format::ImageFormatId;
 use crate::{
     chip::Chip,
     connection::Connection,
@@ -533,16 +534,17 @@ impl Flasher {
         elf_data: &[u8],
         bootloader: Option<Vec<u8>>,
         partition_table: Option<PartitionTable>,
+        image_format: Option<ImageFormatId>,
     ) -> Result<(), Error> {
         let mut image = FirmwareImage::from_data(elf_data)?;
         image.flash_size = self.flash_size();
 
         let mut target = self.chip.flash_target(self.spi_params);
         target.begin(&mut self.connection, &image).flashing()?;
 
-        let flash_image = self
-            .chip
-            .get_flash_image(&image, bootloader, partition_table, None)?;
+        let flash_image =
+            self.chip
+                .get_flash_image(&image, bootloader, partition_table, image_format)?;
 
         for segment in flash_image.flash_segments() {
             target

espflash/src/image_format/esp32bootloader.rs

@@ -6,7 +6,8 @@ use sha2::{Digest, Sha256};
 use crate::{
     chip::Esp32Params,
     elf::{
-        merge_segments, update_checksum, CodeSegment, FirmwareImage, RomSegment, ESP_CHECKSUM_MAGIC,
+        merge_adjacent_segments, update_checksum, CodeSegment, FirmwareImage, RomSegment,
+        ESP_CHECKSUM_MAGIC,
     },
     error::{Error, FlashDetectError},
     flasher::FlashSize,
@@ -62,8 +63,8 @@ impl<'a> Esp32BootloaderFormat<'a> {
 
         let mut checksum = ESP_CHECKSUM_MAGIC;
 
-        let flash_segments: Vec<_> = merge_segments(image.rom_segments(chip).collect());
-        let mut ram_segments: Vec<_> = merge_segments(image.ram_segments(chip).collect());
+        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());
 
         let mut segment_count = 0;
 

espflash/src/image_format/esp32directboot.rs

@@ -0,0 +1,49 @@
+use crate::elf::CodeSegment;
+use crate::{
+    elf::{FirmwareImage, RomSegment},
+    error::Error,
+    image_format::ImageFormat,
+};
+use std::iter::once;
+
+/// Image format for esp32 family chips using a 2nd stage bootloader
+pub struct Esp32DirectBootFormat<'a> {
+    segment: RomSegment<'a>,
+}
+
+impl<'a> Esp32DirectBootFormat<'a> {
+    pub fn new(image: &'a FirmwareImage) -> Result<Self, Error> {
+        let mut segment = image
+            .segments()
+            .map(|mut segment| {
+                // map address to the first 4MB
+                segment.addr %= 0x400000;
+                segment
+            })
+            .fold(CodeSegment::default(), |mut a, b| {
+                a += &b;
+                a
+            });
+        segment.pad_align(4);
+
+        Ok(Self {
+            segment: segment.into(),
+        })
+    }
+}
+
+impl<'a> ImageFormat<'a> for Esp32DirectBootFormat<'a> {
+    fn flash_segments<'b>(&'b self) -> Box<dyn Iterator<Item = RomSegment<'b>> + 'b>
+    where
+        'a: 'b,
+    {
+        Box::new(once(self.segment.borrow()))
+    }
+
+    fn ota_segments<'b>(&'b self) -> Box<dyn Iterator<Item = RomSegment<'b>> + 'b>
+    where
+        'a: 'b,
+    {
+        Box::new(once(self.segment.borrow()))
+    }
+}

espflash/src/image_format/mod.rs

@@ -1,9 +1,11 @@
 mod esp32bootloader;
+mod esp32directboot;
 mod esp8266;
 
 use crate::elf::RomSegment;
 use bytemuck::{Pod, Zeroable};
 pub use esp32bootloader::*;
+pub use esp32directboot::*;
 pub use esp8266::*;
 
 use strum_macros::{AsStaticStr, Display, EnumString};

espflash/src/lib.rs

@@ -14,4 +14,5 @@ pub use config::Config;
 pub use elf::FirmwareImage;
 pub use error::Error;
 pub use flasher::Flasher;
+pub use image_format::ImageFormatId;
 pub use partition_table::PartitionTable;

espflash/src/main.rs

@@ -1,13 +1,14 @@
 use std::fs::{read, read_to_string};
 
-use espflash::{Config, Error, Flasher, PartitionTable};
+use espflash::{Config, Error, Flasher, ImageFormatId, PartitionTable};
 use miette::{IntoDiagnostic, Result, WrapErr};
 use pico_args::Arguments;
 use serial::{BaudRate, FlowControl, SerialPort};
+use std::str::FromStr;
 
 #[allow(clippy::unnecessary_wraps)]
 fn help() -> Result<()> {
-    println!("Usage: espflash [--board-info] [--ram] [--partition-table partition.csv] [--bootloader boot.bin] <serial> <elf image>");
+    println!("Usage: espflash [--board-info] [--ram] [--partition-table partition.csv] [--bootloader boot.bin] [--format <bootloader|direct-boot> <serial> <elf image>");
     Ok(())
 }
 
@@ -27,6 +28,9 @@ fn main() -> Result<()> {
     let partition_table_path = args
         .opt_value_from_str::<_, String>("--partition-table")
         .into_diagnostic()?;
+    let image_format_string = args
+        .opt_value_from_str::<_, String>("--format")
+        .into_diagnostic()?;
 
     let mut serial: Option<String> = args.opt_free_from_str().into_diagnostic()?;
     let mut elf: Option<String> = args.opt_free_from_str().into_diagnostic()?;
@@ -83,6 +87,14 @@ fn main() -> Result<()> {
                     bootloader_path.unwrap()
                 )
             })?;
+        let image_format = image_format_string
+            .as_deref()
+            .map(ImageFormatId::from_str)
+            .transpose()
+            .into_diagnostic()
+            .wrap_err_with(|| {
+                format!("Unrecognized image format {}", image_format_string.unwrap())
+            })?;
         let partition_table = partition_table_path
             .as_deref()
             .map(|path| {
@@ -97,7 +109,7 @@ fn main() -> Result<()> {
                     partition_table_path.unwrap()
                 )
             })?;
-        flasher.load_elf_to_flash(&input_bytes, bootloader, partition_table)?;
+        flasher.load_elf_to_flash(&input_bytes, bootloader, partition_table, image_format)?;
     }
 
     Ok(())