enable support for loader stubs (via --use-stub arg)

Author: ryankurte

Cargo.lock

@@ -29,6 +29,7 @@ bin-dir = "{ bin }{ binary-ext }"
 pkg-fmt = "zip"
 
 [dependencies]
+base64 = "0.13.0"
 binread = "2.2"
 bytemuck = { version = "1.11", features = ["derive"] }
 clap = { version = "3.2", features = ["derive"] }
@@ -47,6 +48,7 @@ parse_int = "0.6"
 regex = "1.6"
 serde = { version = "1.0", features = ["derive"] }
 serde-hex = "0.1"
+serde_json = "1.0.83"
 serde_plain = "1.0"
 serialport = "4.2"
 sha2 = "0.10"

espflash/Cargo.toml

@@ -31,16 +31,21 @@ pub mod monitor;
 
 mod serial;
 
-#[derive(Parser)]
+#[derive(Parser, Debug)]
 pub struct ConnectOpts {
     /// Serial port connected to target device
     pub serial: Option<String>,
+
     /// Baud rate at which to flash target device
     #[clap(long)]
     pub speed: Option<u32>,
+
+    /// Use RAM stub for loading
+    #[clap(long)]
+    pub use_stub: bool,
 }
 
-#[derive(Parser)]
+#[derive(Parser, Debug)]
 pub struct FlashOpts {
     /// Load the application to RAM instead of Flash
     #[clap(long)]
@@ -56,7 +61,7 @@ pub struct FlashOpts {
     pub monitor: bool,
 }
 
-#[derive(Parser)]
+#[derive(Parser, Debug)]
 pub struct FlashConfigOpts {
     /// Flash mode to use
     #[clap(short = 'm', long, possible_values = FlashMode::VARIANTS, value_name = "MODE")]
@@ -69,7 +74,7 @@ pub struct FlashConfigOpts {
     pub flash_freq: Option<FlashFrequency>,
 }
 
-#[derive(Parser)]
+#[derive(Parser, Debug)]
 pub struct PartitionTableOpts {
     /// Convert CSV parition table to binary representation
     #[clap(long, required_unless_present_any = ["info", "to-csv"])]
@@ -87,13 +92,15 @@ pub struct PartitionTableOpts {
     output: Option<PathBuf>,
 }
 
-#[derive(Parser)]
+#[derive(Parser, Debug)]
 pub struct WriteBinToFlashOpts {
     /// Address at which to write the binary file
     #[clap(value_parser = parse_u32)]
     addr: u32,
+
     /// File containing the binary data to write
     bin_file: String,
+
     #[clap(flatten)]
     connect_opts: ConnectOpts,
 }
@@ -131,7 +138,7 @@ pub fn connect(opts: &ConnectOpts, config: &Config) -> Result<Flasher> {
         _ => unreachable!(),
     };
 
-    Ok(Flasher::connect(serial, port_info, opts.speed)?)
+    Ok(Flasher::connect(serial, port_info, opts.speed, opts.use_stub)?)
 }
 
 pub fn board_info(opts: ConnectOpts, config: Config) -> Result<()> {

espflash/src/cli/mod.rs

@@ -86,7 +86,7 @@ impl Connection {
         Err(Error::Connection(ConnectionError::ConnectionFailed))
     }
 
-    fn sync(&mut self) -> Result<(), Error> {
+    pub(crate) fn sync(&mut self) -> Result<(), Error> {
         self.with_timeout(CommandType::Sync.timeout(), |connection| {
             connection.write_command(Command::Sync)?;
             connection.flush()?;
@@ -254,7 +254,7 @@ impl Connection {
         Ok(())
     }
 
-    fn read(&mut self, len: usize) -> Result<Option<Vec<u8>>, Error> {
+    pub(crate) fn read(&mut self, len: usize) -> Result<Option<Vec<u8>>, Error> {
         let mut tmp = Vec::with_capacity(1024);
         loop {
             self.decoder.decode(&mut self.serial, &mut tmp)?;

espflash/src/connection.rs

@@ -21,7 +21,7 @@ use crate::{
 
 pub const ESP_CHECKSUM_MAGIC: u8 = 0xef;
 
-#[derive(Copy, Clone, EnumVariantNames)]
+#[derive(Copy, Clone, Debug, EnumVariantNames)]
 #[strum(serialize_all = "UPPERCASE")]
 pub enum FlashMode {
     Qio,

espflash/src/elf.rs

@@ -11,7 +11,7 @@ use crate::{
     elf::{ElfFirmwareImage, FirmwareImage, FlashFrequency, FlashMode, RomSegment},
     error::{ConnectionError, FlashDetectError, ResultExt},
     image_format::ImageFormatId,
-    Error, PartitionTable,
+    Error, PartitionTable, stubs::FlashStub,
 };
 
 const DEFAULT_TIMEOUT: Duration = Duration::from_secs(3);
@@ -192,6 +192,7 @@ impl Flasher {
         serial: Box<dyn SerialPort>,
         port_info: UsbPortInfo,
         speed: Option<u32>,
+        use_stub: bool,
     ) -> Result<Self, Error> {
         // Establish a connection to the device using the default baud rate of 115,200
         // and timeout of 3 seconds.
@@ -209,6 +210,13 @@ impl Flasher {
             flash_size: FlashSize::Flash4Mb,
             spi_params: SpiAttachParams::default(),
         };
+
+        // Load flash stub if enabled
+        if use_stub {
+            println!("Using flash stub");
+            flasher.load_stub()?;
+        }
+
         flasher.spi_autodetect()?;
 
         // Now that we have established a connection and detected the chip and flash
@@ -228,11 +236,65 @@ impl Flasher {
         Ok(flasher)
     }
 
+    /// Load flash stub
+    fn load_stub(&mut self) -> Result<(), Error> {
+
+        println!("Loading flash stub for chip: {:?}", self.chip);
+
+        // Load flash stub
+        let stub = FlashStub::get(self.chip).unwrap();
+
+        let mut ram_target = self.chip.ram_target(Some(stub.entry()));
+        ram_target.begin(&mut self.connection).flashing()?;
+
+        let (text_addr, text) = stub.text();
+        println!("Write {} byte stub text", text.len());
+
+        ram_target.write_segment(
+                &mut self.connection,
+                RomSegment {
+                    addr: text_addr,
+                    data: Cow::Borrowed(&text),
+                },
+            )
+            .flashing()?;
+
+
+        let (data_addr, data) = stub.data();
+        println!("Write {} byte stub data", data.len());
+
+        ram_target.write_segment(
+                &mut self.connection,
+                RomSegment {
+                    addr: data_addr,
+                    data: Cow::Borrowed(&data),
+                },
+            )
+            .flashing()?;
+
+        println!("Finish stub write");
+        ram_target.finish(&mut self.connection, true).flashing()?;
+        
+        println!("Stub written...");
+
+        // Re-sync connection
+        self.connection.sync()?;
+
+        // Re-detect chip to check stub is up
+        let magic = self.connection.read_reg(CHIP_DETECT_MAGIC_REG_ADDR)?;
+        let chip = Chip::from_magic(magic)?;
+        println!("Re-detected chip: {:?}", chip);
+
+        Ok(())
+    }
+
     fn spi_autodetect(&mut self) -> Result<(), Error> {
         // Loop over all available SPI parameters until we find one that successfully
         // reads the flash size.
         for spi_params in TRY_SPI_PARAMS.iter().copied() {
-            self.enable_flash(spi_params)?;
+            if let Err(_e) = self.enable_flash(spi_params) {
+                continue;
+            }
             if let Some(flash_size) = self.flash_detect()? {
                 // Flash detection was successful, so save the flash size and SPI parameters and
                 // return.

espflash/src/flasher.rs

@@ -19,3 +19,5 @@ pub mod partition_table;
 
 #[doc(hidden)]
 pub mod cli;
+
+pub mod stubs;

espflash/src/lib.rs

@@ -12,25 +12,30 @@ use espflash::{
 use miette::{IntoDiagnostic, Result, WrapErr};
 use strum::VariantNames;
 
-#[derive(Parser)]
+#[derive(Debug, Parser)]
 #[clap(version, propagate_version = true)]
 struct Opts {
     /// Image format to flash
     #[clap(long, possible_values = &["bootloader", "direct-boot"])]
     pub format: Option<String>,
+
     #[clap(flatten)]
     pub flash_config_opts: FlashConfigOpts,
+
     #[clap(flatten)]
     flash_opts: FlashOpts,
+
     #[clap(flatten)]
     connect_opts: ConnectOpts,
+
     /// ELF image to flash
     image: Option<String>,
+
     #[clap(subcommand)]
     subcommand: Option<SubCommand>,
 }
 
-#[derive(Parser)]
+#[derive(Debug, Parser)]
 pub enum SubCommand {
     /// Display information about the connected board and exit without flashing
     BoardInfo(ConnectOpts),
@@ -44,7 +49,7 @@ pub enum SubCommand {
     WriteBinToFlash(WriteBinToFlashOpts),
 }
 
-#[derive(Parser)]
+#[derive(Debug, Parser)]
 pub struct SaveImageOpts {
     #[clap(flatten)]
     pub flash_config_opts: FlashConfigOpts,

espflash/src/main.rs

@@ -0,0 +1,56 @@
+use serde::{Serialize, Deserialize};
+
+use crate::Chip;
+
+
+#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
+pub struct FlashStub {
+    entry: u32,
+    text: String,
+    text_start: u32,
+    data: String,
+    data_start: u32,
+}
+
+const STUB_32:   &str = include_str!("../../../stubs/stub_flasher_32.json");
+const STUB_32C2: &str = include_str!("../../../stubs/stub_flasher_32c2.json");
+const STUB_32C3: &str = include_str!("../../../stubs/stub_flasher_32c3.json");
+const STUB_32S2: &str = include_str!("../../../stubs/stub_flasher_32s2.json");
+const STUB_32S3: &str = include_str!("../../../stubs/stub_flasher_32s3.json");
+const STUB_8266: &str = include_str!("../../../stubs/stub_flasher_8266.json");
+
+impl FlashStub {
+    /// Fetch flash stub for the provided chip
+    pub fn get(chip: Chip) -> Result<FlashStub, ()> {
+        let s = match chip {
+            Chip::Esp32 => STUB_32,
+            Chip::Esp32c2 => STUB_32C2,
+            Chip::Esp32c3 => STUB_32C3,
+            Chip::Esp32s2 => STUB_32S2,
+            Chip::Esp32s3 => STUB_32S3,
+            Chip::Esp8266 => STUB_8266,
+        };
+
+        let stub: FlashStub = serde_json::from_str(s)
+            .map_err(|_| () )?;
+
+        Ok(stub)
+    }
+
+    /// Fetch stub entry point
+    pub fn entry(&self) -> u32 {
+        self.entry
+    }
+
+    /// Fetch text start address and bytes
+    pub fn text(&self) -> (u32, Vec<u8>) {
+        let v = base64::decode(&self.text).unwrap();
+        (self.text_start, v)
+    }
+
+    /// Fetch data start address and bytes
+    pub fn data(&self) -> (u32, Vec<u8>) {
+        let v = base64::decode(&self.data).unwrap();
+        (self.data_start, v)
+    }
+}

espflash/src/stubs/mod.rs

@@ -0,0 +1,2 @@
+
+Flasher stubs from https://github.com/espressif/esptool/tree/master/esptool/targets/stub_flasher