Correct the handling or partition table flags, simplification and cleanup

Author: jessebraham

espflash/src/partition_table.rs

@@ -8,6 +8,8 @@ use std::{
 use md5::{Context, Digest};
 use regex::Regex;
 use serde::{Deserialize, Deserializer, Serialize};
+use strum::IntoEnumIterator;
+use strum_macros::EnumIter;
 
 use crate::error::{
     CSVError, DuplicatePartitionsError, InvalidSubTypeError, NoAppError,
@@ -16,10 +18,11 @@ use crate::error::{
 
 const MAX_PARTITION_LENGTH: usize = 0xC00;
 const PARTITION_TABLE_SIZE: usize = 0x1000;
+const PARTITION_SIZE: usize = 32;
+const PARTITION_ALIGNMENT: u32 = 0x10000;
 
 #[derive(Copy, Clone, Debug, Deserialize, Serialize, PartialEq)]
 #[repr(u8)]
-#[allow(dead_code)]
 #[serde(rename_all = "lowercase")]
 pub enum Type {
     App = 0x00,
@@ -29,21 +32,14 @@ pub enum Type {
 impl Type {
     pub fn subtype_hint(&self) -> String {
         match self {
-            Type::App => "'factory', 'ota_0' through 'ota_15' and 'test'".into(),
+            Type::App => "'factory', 'ota_0' through 'ota_15', and 'test'".into(),
             Type::Data => {
-                use DataType::*;
-                let types = [
-                    Ota, Phy, Nvs, CoreDump, NvsKeys, EFuse, EspHttpd, Fat, Spiffs,
-                ];
-
-                let mut out = format!("'{}'", serde_plain::to_string(&types[0]).unwrap());
-                for ty in &types[1..types.len() - 2] {
-                    let ser = serde_plain::to_string(&ty).unwrap();
-                    write!(&mut out, ", '{}'", ser).unwrap();
-                }
+                let types = DataType::iter()
+                    .map(|dt| format!("'{}'", serde_plain::to_string(&dt).unwrap()))
+                    .collect::<Vec<_>>();
 
-                let ser = serde_plain::to_string(&types[types.len() - 1]).unwrap();
-                write!(&mut out, " and '{}'", ser).unwrap();
+                let mut out = types[0..types.len() - 2].join(", ");
+                write!(&mut out, ", and {}", types[types.len() - 1]).unwrap();
 
                 out
             }
@@ -52,16 +48,15 @@ impl Type {
 }
 
 impl Display for Type {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
         write!(f, "{}", serde_plain::to_string(self).unwrap())
     }
 }
 
 #[derive(Copy, Clone, Debug, Deserialize, Serialize, PartialEq)]
 #[repr(u8)]
-#[allow(dead_code)]
-#[serde(rename_all = "lowercase")]
 pub enum AppType {
+    #[serde(rename = "factory")]
     Factory = 0x00,
     #[serde(rename = "ota_0")]
     Ota0 = 0x10,
@@ -95,12 +90,12 @@ pub enum AppType {
     Ota14 = 0x1e,
     #[serde(rename = "ota_15")]
     Ota15 = 0x1f,
+    #[serde(rename = "test")]
     Test = 0x20,
 }
 
-#[derive(Copy, Clone, Debug, Deserialize, Serialize, PartialEq)]
+#[derive(Copy, Clone, Debug, Deserialize, EnumIter, Serialize, PartialEq)]
 #[repr(u8)]
-#[allow(dead_code)]
 #[serde(rename_all = "lowercase")]
 pub enum DataType {
     Ota = 0x00,
@@ -122,20 +117,20 @@ impl DataType {
 }
 
 #[derive(Debug, Deserialize, PartialEq, Copy, Clone)]
-#[allow(dead_code)]
 #[serde(untagged)]
 pub enum SubType {
     App(AppType),
     Data(DataType),
 }
 
 impl Display for SubType {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
         let ser = match self {
             SubType::App(sub) => serde_plain::to_string(sub),
             SubType::Data(sub) => serde_plain::to_string(sub),
         }
         .unwrap();
+
         write!(f, "{}", ser)
     }
 }
@@ -156,6 +151,18 @@ impl SubType {
     }
 }
 
+#[derive(Copy, Clone, Debug, Deserialize, Serialize, PartialEq)]
+#[serde(rename_all = "lowercase")]
+pub enum Flags {
+    Encrypted = 0x1,
+}
+
+impl Flags {
+    pub fn as_u32(&self) -> u32 {
+        *self as u32
+    }
+}
+
 #[derive(Debug)]
 pub struct PartitionTable {
     partitions: Vec<Partition>,
@@ -202,7 +209,10 @@ impl PartitionTable {
     /// Attempt to parse a partition table from the given string. For more
     /// information on the partition table CSV format see:
     /// https://docs.espressif.com/projects/esp-idf/en/latest/esp32/api-guides/partition-tables.html
-    pub fn try_from_str<S: Into<String>>(data: S) -> Result<Self, PartitionTableError> {
+    pub fn try_from_str<S>(data: S) -> Result<Self, PartitionTableError>
+    where
+        S: Into<String>,
+    {
         let data = data.into();
         let mut reader = csv::ReaderBuilder::new()
             .comment(Some(b'#'))
@@ -213,32 +223,39 @@ impl PartitionTable {
         // Default offset is 0x8000 in esp-idf, partition table size is 0x1000
         let mut offset = 0x9000;
         let mut partitions = Vec::with_capacity(data.lines().count());
+
         for record in reader.records() {
             let record = record.map_err(|e| CSVError::new(e, data.clone()))?;
             let position = record.position();
+
             let mut partition: DeserializedPartition = record
                 .deserialize(None)
                 .map_err(|e| CSVError::new(e, data.clone()))?;
-
             partition.fixup_offset(&mut offset);
 
             let mut partition = Partition::from(partition);
             partition.line = position.map(|pos| pos.line() as usize);
+
             partitions.push(partition);
         }
 
         let table = Self { partitions };
         table.validate(&data)?;
+
         Ok(table)
     }
 
     pub fn to_bytes(&self) -> Vec<u8> {
         let mut result = Vec::with_capacity(PARTITION_TABLE_SIZE);
         self.save(&mut result).unwrap();
+
         result
     }
 
-    pub fn save<W: Write>(&self, writer: &mut W) -> std::io::Result<()> {
+    pub fn save<W>(&self, writer: &mut W) -> std::io::Result<()>
+    where
+        W: Write,
+    {
         let mut hasher = HashWriter::new(writer);
         for partition in &self.partitions {
             partition.save(&mut hasher)?;
@@ -286,7 +303,7 @@ impl PartitionTable {
                     .into());
                 }
 
-                if partition.ty == Type::App && partition.offset.rem(0x10000) != 0 {
+                if partition.ty == Type::App && partition.offset.rem(PARTITION_ALIGNMENT) != 0 {
                     return Err(UnalignedPartitionError::new(source, *line).into());
                 }
             }
@@ -330,8 +347,6 @@ impl PartitionTable {
     }
 }
 
-const PARTITION_SIZE: usize = 32;
-
 #[derive(Debug, Deserialize)]
 pub struct DeserializedPartition {
     #[serde(deserialize_with = "deserialize_partition_name")]
@@ -342,13 +357,13 @@ pub struct DeserializedPartition {
     offset: Option<u32>,
     #[serde(deserialize_with = "deserialize_partition_size")]
     size: u32,
-    flags: Option<u32>,
+    flags: Option<Flags>,
 }
 
 impl DeserializedPartition {
     fn align(offset: u32, ty: Type) -> u32 {
         let pad = match ty {
-            Type::App => 0x10000,
+            Type::App => PARTITION_ALIGNMENT,
             Type::Data => 4,
         };
 
@@ -368,28 +383,14 @@ impl DeserializedPartition {
     }
 }
 
-impl From<DeserializedPartition> for Partition {
-    fn from(part: DeserializedPartition) -> Self {
-        Partition {
-            name: part.name,
-            ty: part.ty,
-            sub_type: part.sub_type,
-            offset: part.offset.unwrap(),
-            size: part.size,
-            flags: part.flags,
-            line: None,
-        }
-    }
-}
-
 #[derive(Debug)]
 pub struct Partition {
     name: String,
     ty: Type,
     sub_type: SubType,
     offset: u32,
     size: u32,
-    flags: Option<u32>,
+    flags: Option<Flags>,
     line: Option<usize>,
 }
 
@@ -399,7 +400,7 @@ impl Partition {
         sub_type: SubType,
         offset: u32,
         size: u32,
-        flags: Option<u32>,
+        flags: Option<Flags>,
     ) -> Self {
         Partition {
             name,
@@ -415,7 +416,10 @@ impl Partition {
         }
     }
 
-    pub fn save<W: Write>(&self, writer: &mut W) -> std::io::Result<()> {
+    pub fn save<W>(&self, writer: &mut W) -> std::io::Result<()>
+    where
+        W: Write,
+    {
         writer.write_all(&[0xAA, 0x50])?;
         writer.write_all(&[self.ty as u8, self.sub_type.as_u8()])?;
         writer.write_all(&self.offset.to_le_bytes())?;
@@ -428,7 +432,7 @@ impl Partition {
         writer.write_all(&name_bytes)?;
 
         let flags = match &self.flags {
-            Some(f) => f.to_le_bytes(),
+            Some(f) => f.as_u32().to_le_bytes(),
             None => 0u32.to_le_bytes(),
         };
         writer.write_all(&flags)?;
@@ -440,11 +444,29 @@ impl Partition {
         self.offset
     }
 
+    pub fn flags(&self) -> Option<Flags> {
+        self.flags
+    }
+
     fn overlaps(&self, other: &Partition) -> bool {
         max(self.offset, other.offset) < min(self.offset + self.size, other.offset + other.size)
     }
 }
 
+impl From<DeserializedPartition> for Partition {
+    fn from(p: DeserializedPartition) -> Self {
+        Partition {
+            name: p.name,
+            ty: p.ty,
+            sub_type: p.sub_type,
+            offset: p.offset.unwrap(),
+            size: p.size,
+            flags: p.flags,
+            line: None,
+        }
+    }
+}
+
 fn deserialize_partition_name<'de, D>(deserializer: D) -> Result<String, D::Error>
 where
     D: Deserializer<'de>,
@@ -511,16 +533,20 @@ where
     D: Deserializer<'de>,
 {
     use serde::de::Error;
+
     let deserialized = deserialize_partition_offset_or_size(deserializer)?;
     deserialized.ok_or_else(|| Error::custom("invalid partition size/offset format"))
 }
 
-struct HashWriter<W: Write> {
+struct HashWriter<W> {
     inner: W,
     hasher: Context,
 }
 
-impl<W: Write> Write for HashWriter<W> {
+impl<W> Write for HashWriter<W>
+where
+    W: Write,
+{
     fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
         self.hasher.write_all(buf)?;
         self.inner.write(buf)
@@ -531,7 +557,10 @@ impl<W: Write> Write for HashWriter<W> {
     }
 }
 
-impl<W: Write> HashWriter<W> {
+impl<W> HashWriter<W>
+where
+    W: Write,
+{
     pub fn new(inner: W) -> Self {
         HashWriter {
             inner,
@@ -553,7 +582,7 @@ mod tests {
 # Name,   Type, SubType, Offset,  Size, Flags
 nvs,      data, nvs,     0x9000,  0x6000,
 phy_init, data, phy,     0xf000,  0x1000,
-factory,  app,  factory, 0x10000, 1M,
+factory,  app,  factory, 0x10000, 1M, encrypted
 ";
 
     const PTABLE_1: &str = "
@@ -644,6 +673,12 @@ phy_init, data, phy,     0xf000,  0x1000,
         let pt0 = PartitionTable::try_from_str(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_str(PTABLE_1);
         assert!(pt1.is_ok());
 
@@ -652,6 +687,7 @@ phy_init, data, phy,     0xf000,  0x1000,
 
         PartitionTable::try_from_str(PTABLE_NO_FACTORY)
             .expect("Failed to parse partition table without factory partition");
+
         PartitionTable::try_from_str(PTABLE_NO_APP)
             .expect_err("Failed to reject partition table without factory or ota partition");
     }