Consistently support binary partition tables (#212)

* Move parsing either CSV or binary partition table into own associated function

Since this is done in multiple places it makes sense for PartitionTable
to just offer this as an associated function.

* Make save-image able to use binary partition tables

Fixes #211

Author: kyrias

espflash/src/cli/mod.rs

@@ -178,15 +178,15 @@ pub fn save_elf_as_image(
         };
 
         // If the '-T' option is provided, load the partition table from
-        // the CSV at the specified path.
+        // the CSV or binary file at the specified path.
         let partition_table = if let Some(partition_table_path) = partition_table_path {
             let path = fs::canonicalize(partition_table_path).into_diagnostic()?;
-            let data = fs::read_to_string(path)
+            let data = fs::read(path)
                 .into_diagnostic()
                 .wrap_err("Failed to open partition table")?;
 
             let table =
-                PartitionTable::try_from_str(data).wrap_err("Failed to parse partition table")?;
+                PartitionTable::try_from(data).wrap_err("Failed to parse partition table")?;
 
             Some(table)
         } else {
@@ -279,26 +279,14 @@ pub fn flash_elf_image(
     };
 
     // If the '--partition-table' option is provided, load the partition table from
-    // the CSV at the specified path.
+    // the CSV or binary file at the specified path.
     let partition_table = if let Some(path) = partition_table {
         let path = fs::canonicalize(path).into_diagnostic()?;
 
-        // If a partition table was detected from ESP-IDF (eg. using `esp-idf-sys`) then
-        // it will be passed in its _binary_ form. Otherwise, it will be provided as a
-        // CSV.
-        let table = if path.extension().map(|e| e.to_str().unwrap()) == Some("csv") {
-            let data = fs::read_to_string(path)
-                .into_diagnostic()
-                .wrap_err("Failed to open partition table")?;
-
-            PartitionTable::try_from_str(data).wrap_err("Failed to parse partition table")?
-        } else {
-            let data = fs::read(path)
-                .into_diagnostic()
-                .wrap_err("Failed to open partition table")?;
-
-            PartitionTable::try_from_bytes(data).wrap_err("Failed to parse partition table")?
-        };
+        let data = fs::read(path)
+            .into_diagnostic()
+            .wrap_err("Failed to open partition table")?;
+        let table = PartitionTable::try_from(data).wrap_err("Failed to parse partition table")?;
 
         Some(table)
     } else {
@@ -350,13 +338,7 @@ pub fn partition_table(opts: PartitionTableOpts) -> Result<()> {
 
         // Try getting the partition table from either the csv or the binary representation and
         // fail otherwise.
-        let part_table = if let Ok(part_table) =
-            PartitionTable::try_from_bytes(input.clone()).into_diagnostic()
-        {
-            part_table
-        } else if let Ok(part_table) =
-            PartitionTable::try_from_str(String::from_utf8(input).into_diagnostic()?)
-        {
+        let part_table = if let Ok(part_table) = PartitionTable::try_from(input).into_diagnostic() {
             part_table
         } else {
             return Err((InvalidPartitionTable {}).into());

espflash/src/partition_table.rs

@@ -14,9 +14,9 @@ use strum::IntoEnumIterator;
 use strum_macros::EnumIter;
 
 use crate::error::{
-    CSVError, DuplicatePartitionsError, InvalidChecksum, InvalidSubTypeError,
-    LengthNotMultipleOf32, NoAppError, NoEndMarker, OverlappingPartitionsError,
-    PartitionTableError, UnalignedPartitionError,
+    CSVError, DuplicatePartitionsError, InvalidChecksum, InvalidPartitionTable,
+    InvalidSubTypeError, LengthNotMultipleOf32, NoAppError, NoEndMarker,
+    OverlappingPartitionsError, PartitionTableError, UnalignedPartitionError,
 };
 
 const MAX_PARTITION_LENGTH: usize = 0xC00;
@@ -219,6 +219,30 @@ impl PartitionTable {
         }
     }
 
+    /// Attempt to parse either a binary or CSV partition table from the given input.
+    ///
+    /// For more information on the partition table format see:
+    /// https://docs.espressif.com/projects/esp-idf/en/latest/esp32/api-guides/partition-tables.html
+    pub fn try_from<S>(data: S) -> Result<Self, PartitionTableError>
+    where
+        S: Into<Vec<u8>>,
+    {
+        let input: Vec<u8> = data.into();
+
+        // If a partition table was detected from ESP-IDF (eg. using `esp-idf-sys`) then
+        // it will be passed in its _binary_ form. Otherwise, it will be provided as a
+        // CSV.
+        if let Ok(part_table) = Self::try_from_bytes(&*input) {
+            Ok(part_table)
+        } else if let Ok(part_table) =
+            Self::try_from_str(String::from_utf8(input).map_err(|_| InvalidPartitionTable)?)
+        {
+            Ok(part_table)
+        } else {
+            Err(InvalidPartitionTable.into())
+        }
+    }
+
     /// Attempt to parse a CSV partition table from the given string.
     ///
     /// For more information on the partition table format see:
@@ -812,6 +836,39 @@ phy_init, data, phy,     0xf000,  0x1000,
         assert_eq!(expected, result.as_slice());
     }
 
+    #[test]
+    fn test_from() {
+        let pt0 = PartitionTable::try_from(PTABLE_0);
+        assert!(pt0.is_ok());
+
+        let pt0 = pt0.unwrap();
+        let nvs = pt0.find("nvs").unwrap();
+        let fac = pt0.find("factory").unwrap();
+        assert_eq!(nvs.flags(), None);
+        assert_eq!(fac.flags(), Some(Flags::Encrypted));
+
+        let pt1 = PartitionTable::try_from(PTABLE_1);
+        assert!(pt1.is_ok());
+
+        let pt_spiffs = PartitionTable::try_from(PTABLE_SPIFFS);
+        assert!(pt_spiffs.is_ok());
+
+        PartitionTable::try_from(PTABLE_NO_FACTORY)
+            .expect("Failed to parse partition table without factory partition");
+
+        PartitionTable::try_from(PTABLE_NO_APP)
+            .expect_err("Failed to reject partition table without factory or ota partition");
+
+        use std::fs::{read, read_to_string};
+        let binary_table = read("./tests/data/partitions.bin").unwrap();
+        let binary_parsed = PartitionTable::try_from_bytes(binary_table).unwrap();
+
+        let csv_table = read_to_string("./tests/data/partitions.csv").unwrap();
+        let csv_parsed = PartitionTable::try_from(csv_table).unwrap();
+
+        assert_eq!(binary_parsed, csv_parsed);
+    }
+
     #[test]
     fn test_from_str() {
         let pt0 = PartitionTable::try_from_str(PTABLE_0);