From 8dea4e67c4492d7b5db604dbebb1b6172dea8a4e Mon Sep 17 00:00:00 2001
From: Jesse Braham <jessebraham@users.noreply.github.com>
Date: Wed, 26 Mar 2025 12:54:18 +0100
Subject: [PATCH] Add chip detection using security info (#814)

* Add the ability to create a `Chip` from its corresponding chip ID

* Attempt to detect chip using security info first, and use magic value if this fails

* Update timeouts in HIL workflow

* Update `CHANGELOG.md`

* Increase more timeouts for HIL
---
 .github/workflows/hil.yml       | 10 +++---
 CHANGELOG.md                    |  1 +
 espflash/src/error.rs           |  4 +--
 espflash/src/flasher/mod.rs     | 63 +++++++++++++++++++++------------
 espflash/src/targets/esp32.rs   |  4 ++-
 espflash/src/targets/esp32c2.rs |  4 ++-
 espflash/src/targets/esp32c3.rs |  4 ++-
 espflash/src/targets/esp32c6.rs |  4 ++-
 espflash/src/targets/esp32h2.rs |  4 ++-
 espflash/src/targets/esp32p4.rs |  4 ++-
 espflash/src/targets/esp32s2.rs |  4 ++-
 espflash/src/targets/esp32s3.rs |  4 ++-
 espflash/src/targets/mod.rs     | 24 ++++++++++++-
 espflash/tests/scripts/flash.sh |  6 ++--
 14 files changed, 98 insertions(+), 42 deletions(-)

diff --git a/.github/workflows/hil.yml b/.github/workflows/hil.yml
index ffca619..391b3d6 100644
--- a/.github/workflows/hil.yml
+++ b/.github/workflows/hil.yml
@@ -103,7 +103,7 @@ jobs:
         run: timeout 10 bash espflash/tests/scripts/board-info.sh
 
       - name: flash test
-        run: timeout 25 bash espflash/tests/scripts/flash.sh ${{ matrix.board.mcu }}
+        run: timeout 60 bash espflash/tests/scripts/flash.sh ${{ matrix.board.mcu }}
 
       - name: monitor test
         run: timeout 10 bash espflash/tests/scripts/monitor.sh
@@ -116,13 +116,13 @@ jobs:
           timeout 90 bash espflash/tests/scripts/save-image_write-bin.sh ${{ matrix.board.mcu }}
 
       - name: erase-region test
-        run: timeout 10 bash espflash/tests/scripts/erase-region.sh
+        run: timeout 15 bash espflash/tests/scripts/erase-region.sh
 
       - name: hold-in-reset test
-        run: timeout 5 bash espflash/tests/scripts/hold-in-reset.sh
+        run: timeout 10 bash espflash/tests/scripts/hold-in-reset.sh
 
       - name: reset test
-        run: timeout 5 bash espflash/tests/scripts/reset.sh
+        run: timeout 10 bash espflash/tests/scripts/reset.sh
 
       - name: checksum-md5 test
         run: timeout 40 bash espflash/tests/scripts/checksum-md5.sh
@@ -134,4 +134,4 @@ jobs:
         run: timeout 20 bash espflash/tests/scripts/write-bin.sh
 
       - name: read-flash test
-        run: timeout 30 bash espflash/tests/scripts/read-flash.sh
+        run: timeout 60 bash espflash/tests/scripts/read-flash.sh
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0819ca8..3e735aa 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add `ROM` version of `read-flash` command (#812)
 - `espflash` can detect the log format automatically from ESP-HAL metadata. Reqires `esp-println` 0.14 (presumably, yet to be released) (#809)
 - Add `--output-format` option to monitor (#818)
+- Added chip detection based on security info, where supported (#814)
 
 ### Changed
 
diff --git a/espflash/src/error.rs b/espflash/src/error.rs
index 42a7aa1..43c1e84 100644
--- a/espflash/src/error.rs
+++ b/espflash/src/error.rs
@@ -31,14 +31,14 @@ pub enum Error {
     #[diagnostic(code(espflash::cancelled))]
     Cancelled,
 
-    #[error("Unrecognized magic value: {0:#x}")]
+    #[error("{0}")]
     #[diagnostic(
         code(espflash::chip_detect_error),
         help("Supported chips are: {}\n\
               If your chip is supported, try hard-resetting the device and try again",
              Chip::VARIANTS.join(", "))
     )]
-    ChipDetectError(u32),
+    ChipDetectError(String),
 
     #[error("Chip provided with `-c/--chip` ({0}) does not match the detected chip ({1})")]
     #[diagnostic(
diff --git a/espflash/src/flasher/mod.rs b/espflash/src/flasher/mod.rs
index 97eda27..4bf704d 100644
--- a/espflash/src/flasher/mod.rs
+++ b/espflash/src/flasher/mod.rs
@@ -676,8 +676,7 @@ impl Flasher {
 
         let detected_chip = if before_operation != ResetBeforeOperation::NoResetNoSync {
             // Detect which chip we are connected to.
-            let magic = connection.read_reg(CHIP_DETECT_MAGIC_REG_ADDR)?;
-            let detected_chip = Chip::from_magic(magic)?;
+            let detected_chip = detect_chip(&mut connection, use_stub)?;
             if let Some(chip) = chip {
                 if chip != detected_chip {
                     return Err(Error::ChipMismatch(
@@ -686,6 +685,7 @@ impl Flasher {
                     ));
                 }
             }
+
             detected_chip
         } else if before_operation == ResetBeforeOperation::NoResetNoSync && chip.is_some() {
             chip.unwrap()
@@ -792,8 +792,7 @@ impl Flasher {
         }?;
 
         // 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)?;
+        let chip = detect_chip(&mut self.connection, self.use_stub)?;
         debug!("Re-detected chip: {:?}", chip);
 
         Ok(())
@@ -1126,25 +1125,7 @@ impl Flasher {
 
     /// Get security info
     pub fn get_security_info(&mut self) -> Result<SecurityInfo, Error> {
-        self.connection
-            .with_timeout(CommandType::GetSecurityInfo.timeout(), |connection| {
-                let response = connection.command(Command::GetSecurityInfo)?;
-                // Extract raw bytes and convert them into `SecurityInfo`
-                if let crate::connection::CommandResponseValue::Vector(data) = response {
-                    // HACK: Not quite sure why there seem to be 4 extra bytes at the end of the
-                    //       response when the stub is not being used...
-                    let end = if self.use_stub {
-                        data.len()
-                    } else {
-                        data.len() - 4
-                    };
-                    SecurityInfo::try_from(&data[..end])
-                } else {
-                    Err(Error::InvalidResponse(
-                        "response was not a vector of bytes".into(),
-                    ))
-                }
-            })
+        get_security_info(&mut self.connection, self.use_stub)
     }
 
     pub fn change_baud(&mut self, speed: u32) -> Result<(), Error> {
@@ -1367,3 +1348,39 @@ impl Flasher {
         Ok(())
     }
 }
+
+#[cfg(feature = "serialport")]
+fn get_security_info(connection: &mut Connection, use_stub: bool) -> Result<SecurityInfo, Error> {
+    connection.with_timeout(CommandType::GetSecurityInfo.timeout(), |connection| {
+        let response = connection.command(Command::GetSecurityInfo)?;
+        // Extract raw bytes and convert them into `SecurityInfo`
+        if let crate::connection::CommandResponseValue::Vector(data) = response {
+            // HACK: Not quite sure why there seem to be 4 extra bytes at the end of the
+            //       response when the stub is not being used...
+            let end = if use_stub { data.len() } else { data.len() - 4 };
+            SecurityInfo::try_from(&data[..end])
+        } else {
+            Err(Error::InvalidResponse(
+                "response was not a vector of bytes".into(),
+            ))
+        }
+    })
+}
+
+#[cfg(feature = "serialport")]
+fn detect_chip(connection: &mut Connection, use_stub: bool) -> Result<Chip, Error> {
+    match get_security_info(connection, use_stub) {
+        Ok(info) if info.chip_id.is_some() => {
+            let chip_id = info.chip_id.unwrap() as u16;
+            let chip = Chip::try_from(chip_id)?;
+
+            Ok(chip)
+        }
+        _ => {
+            let magic = connection.read_reg(CHIP_DETECT_MAGIC_REG_ADDR)?;
+            let chip = Chip::from_magic(magic)?;
+
+            Ok(chip)
+        }
+    }
+}
diff --git a/espflash/src/targets/esp32.rs b/espflash/src/targets/esp32.rs
index 3516853..328b533 100644
--- a/espflash/src/targets/esp32.rs
+++ b/espflash/src/targets/esp32.rs
@@ -11,6 +11,8 @@ use crate::{
     Error,
 };
 
+pub(crate) const CHIP_ID: u16 = 0;
+
 const CHIP_DETECT_MAGIC_VALUES: &[u32] = &[0x00f0_1d83];
 
 const FLASH_RANGES: &[Range<u32>] = &[
@@ -179,7 +181,7 @@ impl Target for Esp32 {
             0x1000,
             0x1_0000,
             0x3f_0000,
-            0,
+            CHIP_ID,
             FlashFrequency::_40Mhz,
             booloader,
         );
diff --git a/espflash/src/targets/esp32c2.rs b/espflash/src/targets/esp32c2.rs
index 28dd3a9..7691ba8 100644
--- a/espflash/src/targets/esp32c2.rs
+++ b/espflash/src/targets/esp32c2.rs
@@ -12,6 +12,8 @@ use crate::{
     Error,
 };
 
+pub(crate) const CHIP_ID: u16 = 12;
+
 const CHIP_DETECT_MAGIC_VALUES: &[u32] = &[
     0x6f51_306f, // ECO0
     0x7c41_a06f, // ECO1
@@ -116,7 +118,7 @@ impl Target for Esp32c2 {
             0x0,
             0x1_0000,
             0x1f_0000,
-            12,
+            CHIP_ID,
             FlashFrequency::_30Mhz,
             booloader,
         );
diff --git a/espflash/src/targets/esp32c3.rs b/espflash/src/targets/esp32c3.rs
index e9a6913..c3c07f4 100644
--- a/espflash/src/targets/esp32c3.rs
+++ b/espflash/src/targets/esp32c3.rs
@@ -11,6 +11,8 @@ use crate::{
     Error,
 };
 
+pub(crate) const CHIP_ID: u16 = 5;
+
 const CHIP_DETECT_MAGIC_VALUES: &[u32] = &[
     0x6921_506f, // ECO1 + ECO2
     0x1b31_506f, // ECO3
@@ -27,7 +29,7 @@ const PARAMS: Esp32Params = Esp32Params::new(
     0x0,
     0x1_0000,
     0x3f_0000,
-    5,
+    CHIP_ID,
     FlashFrequency::_40Mhz,
     include_bytes!("../../resources/bootloaders/esp32c3-bootloader.bin"),
 );
diff --git a/espflash/src/targets/esp32c6.rs b/espflash/src/targets/esp32c6.rs
index fa23b08..a860f55 100644
--- a/espflash/src/targets/esp32c6.rs
+++ b/espflash/src/targets/esp32c6.rs
@@ -11,6 +11,8 @@ use crate::{
     Error,
 };
 
+pub(crate) const CHIP_ID: u16 = 13;
+
 const CHIP_DETECT_MAGIC_VALUES: &[u32] = &[0x2CE0_806F];
 
 const FLASH_RANGES: &[Range<u32>] = &[
@@ -22,7 +24,7 @@ const PARAMS: Esp32Params = Esp32Params::new(
     0x0,
     0x1_0000,
     0x3f_0000,
-    13,
+    CHIP_ID,
     FlashFrequency::_40Mhz,
     include_bytes!("../../resources/bootloaders/esp32c6-bootloader.bin"),
 );
diff --git a/espflash/src/targets/esp32h2.rs b/espflash/src/targets/esp32h2.rs
index 4cbfffb..ca8c20c 100644
--- a/espflash/src/targets/esp32h2.rs
+++ b/espflash/src/targets/esp32h2.rs
@@ -11,6 +11,8 @@ use crate::{
     Error,
 };
 
+pub(crate) const CHIP_ID: u16 = 16;
+
 const CHIP_DETECT_MAGIC_VALUES: &[u32] = &[0xD7B7_3E80];
 
 const FLASH_RANGES: &[Range<u32>] = &[
@@ -22,7 +24,7 @@ const PARAMS: Esp32Params = Esp32Params::new(
     0x0,
     0x1_0000,
     0x3f_0000,
-    16,
+    CHIP_ID,
     FlashFrequency::_24Mhz,
     include_bytes!("../../resources/bootloaders/esp32h2-bootloader.bin"),
 );
diff --git a/espflash/src/targets/esp32p4.rs b/espflash/src/targets/esp32p4.rs
index 7d0f478..f4bd6b2 100644
--- a/espflash/src/targets/esp32p4.rs
+++ b/espflash/src/targets/esp32p4.rs
@@ -11,6 +11,8 @@ use crate::{
     Error,
 };
 
+pub(crate) const CHIP_ID: u16 = 18;
+
 const CHIP_DETECT_MAGIC_VALUES: &[u32] = &[0x0, 0x0ADDBAD0];
 
 const FLASH_RANGES: &[Range<u32>] = &[
@@ -22,7 +24,7 @@ const PARAMS: Esp32Params = Esp32Params::new(
     0x2000,
     0x1_0000,
     0x3f_0000,
-    18,
+    CHIP_ID,
     FlashFrequency::_40Mhz,
     include_bytes!("../../resources/bootloaders/esp32p4-bootloader.bin"),
 );
diff --git a/espflash/src/targets/esp32s2.rs b/espflash/src/targets/esp32s2.rs
index 33c338a..86a1b8c 100644
--- a/espflash/src/targets/esp32s2.rs
+++ b/espflash/src/targets/esp32s2.rs
@@ -13,6 +13,8 @@ use crate::{
     Error,
 };
 
+pub(crate) const CHIP_ID: u16 = 2;
+
 const CHIP_DETECT_MAGIC_VALUES: &[u32] = &[0x0000_07c6];
 
 const FLASH_RANGES: &[Range<u32>] = &[
@@ -27,7 +29,7 @@ const PARAMS: Esp32Params = Esp32Params::new(
     0x1000,
     0x1_0000,
     0x10_0000,
-    2,
+    CHIP_ID,
     FlashFrequency::_40Mhz,
     include_bytes!("../../resources/bootloaders/esp32s2-bootloader.bin"),
 );
diff --git a/espflash/src/targets/esp32s3.rs b/espflash/src/targets/esp32s3.rs
index 21f6c6b..8df3be3 100644
--- a/espflash/src/targets/esp32s3.rs
+++ b/espflash/src/targets/esp32s3.rs
@@ -11,6 +11,8 @@ use crate::{
     Error,
 };
 
+pub(crate) const CHIP_ID: u16 = 9;
+
 const CHIP_DETECT_MAGIC_VALUES: &[u32] = &[0x9];
 
 const FLASH_RANGES: &[Range<u32>] = &[
@@ -22,7 +24,7 @@ const PARAMS: Esp32Params = Esp32Params::new(
     0x0,
     0x1_0000,
     0x10_0000,
-    9,
+    CHIP_ID,
     FlashFrequency::_40Mhz,
     include_bytes!("../../resources/bootloaders/esp32s3-bootloader.bin"),
 );
diff --git a/espflash/src/targets/mod.rs b/espflash/src/targets/mod.rs
index 33a2c07..f70c396 100644
--- a/espflash/src/targets/mod.rs
+++ b/espflash/src/targets/mod.rs
@@ -126,7 +126,9 @@ impl Chip {
         } else if Esp32s3::has_magic_value(magic) {
             Ok(Chip::Esp32s3)
         } else {
-            Err(Error::ChipDetectError(magic))
+            Err(Error::ChipDetectError(format!(
+                "unrecognized magic value: {magic:#x}"
+            )))
         }
     }
 
@@ -191,6 +193,26 @@ impl Chip {
     }
 }
 
+impl TryFrom<u16> for Chip {
+    type Error = Error;
+
+    fn try_from(value: u16) -> Result<Self, Self::Error> {
+        match value {
+            esp32::CHIP_ID => Ok(Chip::Esp32),
+            esp32c2::CHIP_ID => Ok(Chip::Esp32c2),
+            esp32c3::CHIP_ID => Ok(Chip::Esp32c3),
+            esp32c6::CHIP_ID => Ok(Chip::Esp32c6),
+            esp32h2::CHIP_ID => Ok(Chip::Esp32h2),
+            esp32p4::CHIP_ID => Ok(Chip::Esp32p4),
+            esp32s2::CHIP_ID => Ok(Chip::Esp32s2),
+            esp32s3::CHIP_ID => Ok(Chip::Esp32s3),
+            _ => Err(Error::ChipDetectError(format!(
+                "unrecognized chip ID: {value}"
+            ))),
+        }
+    }
+}
+
 /// Device-specific parameters
 #[derive(Debug, Clone, Copy)]
 pub struct Esp32Params {
diff --git a/espflash/tests/scripts/flash.sh b/espflash/tests/scripts/flash.sh
index 90a5551..4c40942 100644
--- a/espflash/tests/scripts/flash.sh
+++ b/espflash/tests/scripts/flash.sh
@@ -4,7 +4,7 @@ app="espflash/tests/data/$1"
 if [[ "$1" == "esp32c6" ]]; then
     # With manual log-format
     app_defmt="${app}_defmt"
-    result=$(timeout 8s espflash flash --no-skip --monitor --non-interactive $app_defmt --log-format defmt 2>&1)
+    result=$(timeout 15s espflash flash --no-skip --monitor --non-interactive $app_defmt --log-format defmt 2>&1)
     echo "$result"
     if [[ ! $result =~ "Flashing has completed!" ]]; then
         echo "Flashing failed!"
@@ -16,7 +16,7 @@ if [[ "$1" == "esp32c6" ]]; then
     fi
 
     # With auto-detected log-format
-    result=$(timeout 8s espflash flash --no-skip --monitor --non-interactive $app_defmt 2>&1)
+    result=$(timeout 15s espflash flash --no-skip --monitor --non-interactive $app_defmt 2>&1)
     echo "$result"
     if [[ ! $result =~ "Flashing has completed!" ]]; then
         echo "Flashing failed!"
@@ -28,7 +28,7 @@ if [[ "$1" == "esp32c6" ]]; then
     fi
 fi
 
-result=$(timeout 8s espflash flash --no-skip --monitor --non-interactive $app 2>&1)
+result=$(timeout 15s espflash flash --no-skip --monitor --non-interactive $app 2>&1)
 echo "$result"
 if [[ ! $result =~ "Flashing has completed!" ]]; then
     echo "Flashing failed!"