Add support for erasing otadata (#229)

* Add support for erasing otadata * Fix formatting * Look for proper OTA partition by type and subtype
esp-rs · Sep 13, 2022 · d64ff68 · d64ff68
1 parent 1b23940
commit d64ff68
Show file tree

Hide file tree

Showing 8 changed files with 103 additions and 5 deletions.
diff --git a/cargo-espflash/src/main.rs b/cargo-espflash/src/main.rs
@@ -143,10 +143,15 @@ fn main() -> Result<()> {
         .init();
 
     // Extract subcommand
-    let CargoSubCommand::Espflash(opts) = opts.subcommand;
+    let CargoSubCommand::Espflash(mut opts) = opts.subcommand;
 
     debug!("subcommand options: {:?}", opts);
 
+    // `erase_otadata` requires `use_stub`
+    if opts.flash_opts.erase_otadata {
+        opts.connect_opts.use_stub = true;
+    }
+
     // Load configuration and metadata
     let config = Config::load()?;
     let metadata = CargoEspFlashMeta::load("Cargo.toml")?;
@@ -226,6 +231,7 @@ fn flash(
             opts.build_opts.flash_config_opts.flash_mode,
             opts.build_opts.flash_config_opts.flash_size,
             opts.build_opts.flash_config_opts.flash_freq,
+            opts.flash_opts.erase_otadata,
         )?;
     }
 

diff --git a/espflash/src/cli/mod.rs b/espflash/src/cli/mod.rs
@@ -21,9 +21,10 @@ use crate::{
     cli::monitor::monitor,
     cli::serial::get_serial_port_info,
     elf::{ElfFirmwareImage, FlashFrequency, FlashMode},
-    error::Error,
+    error::{Error, NoOtadataError},
     flasher::FlashSize,
-    Chip, Flasher, ImageFormatId, InvalidPartitionTable, PartitionTable,
+    partition_table, Chip, Flasher, ImageFormatId, InvalidPartitionTable, MissingPartitionTable,
+    PartitionTable,
 };
 
 pub mod config;
@@ -59,6 +60,10 @@ pub struct FlashOpts {
     /// Open a serial monitor after flashing
     #[clap(long)]
     pub monitor: bool,
+    /// Erase the OTADATA partition
+    /// This is useful when using multiple OTA partitions and still wanting to be able to reflash via espflash
+    #[clap(long)]
+    pub erase_otadata: bool,
 }
 
 #[derive(Debug, Clone, Parser)]
@@ -283,6 +288,7 @@ pub fn flash_elf_image(
     flash_mode: Option<FlashMode>,
     flash_size: Option<FlashSize>,
     flash_freq: Option<FlashFrequency>,
+    erase_otadata: bool,
 ) -> Result<()> {
     // If the '--bootloader' option is provided, load the binary file at the
     // specified path.
@@ -310,6 +316,26 @@ pub fn flash_elf_image(
         None
     };
 
+    if erase_otadata {
+        let partition_table = match &partition_table {
+            Some(partition_table) => partition_table,
+            None => return Err((MissingPartitionTable {}).into()),
+        };
+
+        let otadata = match partition_table.find_by_subtype(
+            partition_table::Type::CoreType(partition_table::CoreType::Data),
+            partition_table::SubType::Data(partition_table::DataType::Ota),
+        ) {
+            Some(otadata) => otadata,
+            None => return Err((NoOtadataError {}).into()),
+        };
+
+        let offset = otadata.offset();
+        let size = otadata.size();
+
+        flasher.erase_region(offset, size)?;
+    }
+
     // Load the ELF data, optionally using the provider bootloader/partition
     // table/image format, to the device's flash memory.
     flasher.load_elf_to_flash_with_format(

diff --git a/espflash/src/command.rs b/espflash/src/command.rs
@@ -8,6 +8,7 @@ use strum_macros::Display;
 const DEFAULT_TIMEOUT: Duration = Duration::from_secs(3);
 const ERASE_REGION_TIMEOUT_PER_MB: Duration = Duration::from_secs(30);
 const ERASE_WRITE_TIMEOUT_PER_MB: Duration = Duration::from_secs(40);
+const ERASE_CHIP_TIMEOUT: Duration = Duration::from_secs(120);
 const MEM_END_TIMEOUT: Duration = Duration::from_millis(50);
 const SYNC_TIMEOUT: Duration = Duration::from_millis(100);
 
@@ -34,13 +35,17 @@ pub enum CommandType {
     FlashDeflateEnd = 0x12,
     FlashMd5 = 0x13,
     FlashDetect = 0x9f,
+    // Some commands supported by stub only
+    EraseFlash = 0xd0,
+    EraseRegion = 0xd1,
 }
 
 impl CommandType {
     pub fn timeout(&self) -> Duration {
         match self {
             CommandType::MemEnd => MEM_END_TIMEOUT,
             CommandType::Sync => SYNC_TIMEOUT,
+            CommandType::EraseFlash => ERASE_CHIP_TIMEOUT,
             _ => DEFAULT_TIMEOUT,
         }
     }
@@ -54,7 +59,7 @@ impl CommandType {
             )
         }
         match self {
-            CommandType::FlashBegin | CommandType::FlashDeflateBegin => {
+            CommandType::FlashBegin | CommandType::FlashDeflateBegin | CommandType::EraseRegion => {
                 calc_timeout(ERASE_REGION_TIMEOUT_PER_MB, size)
             }
             CommandType::FlashData | CommandType::FlashDeflateData => {
@@ -138,6 +143,11 @@ pub enum Command<'a> {
         reboot: bool,
     },
     FlashDetect,
+    EraseFlash,
+    EraseRegion {
+        offset: u32,
+        size: u32,
+    },
 }
 
 impl<'a> Command<'a> {
@@ -159,6 +169,8 @@ impl<'a> Command<'a> {
             Command::FlashDeflateData { .. } => CommandType::FlashDeflateData,
             Command::FlashDeflateEnd { .. } => CommandType::FlashDeflateEnd,
             Command::FlashDetect => CommandType::FlashDetect,
+            Command::EraseFlash { .. } => CommandType::EraseFlash,
+            Command::EraseRegion { .. } => CommandType::EraseRegion,
         }
     }
 
@@ -319,6 +331,18 @@ impl<'a> Command<'a> {
             Command::FlashDetect => {
                 write_basic(writer, &[], 0)?;
             }
+            Command::EraseFlash => {
+                write_basic(writer, &[], 0)?;
+            }
+            Command::EraseRegion { offset, size } => {
+                // length
+                writer.write_all(&(8u16.to_le_bytes()))?;
+                // checksum
+                writer.write_all(&(0u32.to_le_bytes()))?;
+                // data
+                writer.write_all(&offset.to_le_bytes())?;
+                writer.write_all(&size.to_le_bytes())?;
+            }
         };
         Ok(())
     }

diff --git a/espflash/src/error.rs b/espflash/src/error.rs
@@ -368,6 +368,9 @@ pub enum PartitionTableError {
     #[error(transparent)]
     #[diagnostic(transparent)]
     InvalidPartitionTable(#[from] InvalidPartitionTable),
+    #[error(transparent)]
+    #[diagnostic(transparent)]
+    MissingPartitionTable(#[from] MissingPartitionTable),
 }
 
 #[derive(Debug, Error, Diagnostic)]
@@ -538,7 +541,14 @@ impl NoAppError {
         }
     }
 }
+#[derive(Debug, Error, Diagnostic)]
+#[error("No otadata partition was found")]
+#[diagnostic(
+    code(espflash::partition_table::no_otadata),
+    help("Partition table must contain an otadata partition when trying to erase it")
+)]
 
+pub struct NoOtadataError;
 #[derive(Debug, Error, Diagnostic)]
 #[error("Unaligned partition")]
 #[diagnostic(code(espflash::partition_table::unaligned))]
@@ -578,6 +588,11 @@ pub struct NoEndMarker;
 #[diagnostic(code(espflash::partition_table::invalid_partition_table))]
 pub struct InvalidPartitionTable;
 
+#[derive(Debug, Error, Diagnostic)]
+#[error("Missing partition table")]
+#[diagnostic(code(espflash::partition_table::missing_partition_table))]
+pub struct MissingPartitionTable;
+
 #[derive(Debug, Error)]
 #[error("{0}")]
 pub struct ElfError(&'static str);

diff --git a/espflash/src/flasher.rs b/espflash/src/flasher.rs
@@ -641,6 +641,18 @@ impl Flasher {
     pub fn get_usb_pid(&self) -> Result<u16, Error> {
         self.connection.get_usb_pid()
     }
+
+    pub fn erase_region(&mut self, offset: u32, size: u32) -> Result<(), Error> {
+        debug!("Erasing region of 0x{:x}B at 0x{:08x}", size, offset);
+
+        self.connection
+            .with_timeout(CommandType::EraseRegion.timeout(), |connection| {
+                connection.command(Command::EraseRegion { offset, size })
+            })?;
+        std::thread::sleep(Duration::from_secs_f32(0.05));
+        self.connection.flush()?;
+        Ok(())
+    }
 }
 
 pub(crate) fn get_erase_size(offset: usize, size: usize) -> usize {

diff --git a/espflash/src/lib.rs b/espflash/src/lib.rs
@@ -1,7 +1,7 @@
 pub use chip::Chip;
 pub use cli::config::Config;
 pub use elf::{FlashFrequency, FlashMode};
-pub use error::{Error, InvalidPartitionTable};
+pub use error::{Error, InvalidPartitionTable, MissingPartitionTable};
 pub use flasher::{FlashSize, Flasher};
 pub use image_format::ImageFormatId;
 pub use partition_table::PartitionTable;

diff --git a/espflash/src/main.rs b/espflash/src/main.rs
@@ -95,6 +95,10 @@ fn main() -> Result<()> {
 
     debug!("options: {:?}", opts);
 
+    if opts.flash_opts.erase_otadata {
+        opts.connect_opts.use_stub = true;
+    }
+
     // Setup logging
     tracing_subscriber::fmt()
         .with_env_filter(EnvFilter::from_default_env().add_directive(opts.log_level.into()))
@@ -167,6 +171,7 @@ fn flash(opts: Opts, config: Config) -> Result<()> {
             opts.flash_config_opts.flash_mode,
             opts.flash_config_opts.flash_size,
             opts.flash_config_opts.flash_freq,
+            opts.flash_opts.erase_otadata,
         )?;
     }
 

diff --git a/espflash/src/partition_table.rs b/espflash/src/partition_table.rs
@@ -418,6 +418,12 @@ impl PartitionTable {
         self.partitions.iter().find(|&p| p.ty == ty)
     }
 
+    pub fn find_by_subtype(&self, ty: Type, sub_type: SubType) -> Option<&Partition> {
+        self.partitions
+            .iter()
+            .find(|&p| p.ty == ty && p.sub_type == sub_type)
+    }
+
     fn validate(&self, source: &str) -> Result<(), PartitionTableError> {
         for partition in &self.partitions {
             if let Some(line) = &partition.line {
@@ -662,6 +668,10 @@ impl Partition {
         self.offset
     }
 
+    pub fn size(&self) -> u32 {
+        self.size
+    }
+
     pub fn flags(&self) -> Option<Flags> {
         self.flags
     }