feat: tbf serialization

tbf-parser/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 
 [dependencies]
 log = "0.4.27"
+bitflags = "2.11.1"
 
 [features]
 default = []

tbf-parser/src/types.rs

@@ -11,6 +11,9 @@ use core::convert::TryInto;
 use core::mem::size_of;
 use core::{fmt, str};
 
+// Flags crate
+use bitflags::bitflags;
+
 /// We only support up to a fixed number of storage permissions for each of read
 /// and modify. This simplification enables us to use fixed sized buffers.
 const NUM_STORAGE_PERMISSIONS: usize = 8;
@@ -126,6 +129,23 @@ pub struct TbfHeaderV2Base {
     pub(crate) checksum: u32,
 }
 
+bitflags! {
+    /// TBF flags
+    ///
+    /// The available flags can be found here. The remaining bits are reserved
+    /// https://book.tockos.org/doc/tock_binary_format?tbf-header-base#tbf-header-base
+    pub struct Flags: u32 {
+        const Enabled = 1;
+        const Sticky = 1 << 1;
+    }
+}
+
+impl Flags {
+    pub fn as_u32(&self) -> u32 {
+        self.bits()
+    }
+}
+
 /// Types in TLV structures for each optional block of the header.
 #[derive(Clone, Copy, Debug)]
 pub enum TbfHeaderTypes {
@@ -912,6 +932,67 @@ impl TbfHeader {
         }
     }
 
+    pub fn set_flags(&mut self, flags: Flags) {
+        match self {
+            TbfHeader::TbfHeaderV2(hd) => {
+                let old_flags = hd.base.flags;
+                hd.base.flags = flags.as_u32();
+                self.compute_checksum_flags(old_flags);
+            }
+            TbfHeader::Padding(hd) => {
+                let old_flags = hd.flags;
+                hd.flags = flags.as_u32();
+                self.compute_checksum_flags(old_flags);
+            }
+        }
+    }
+
+    pub fn set_enabled(&mut self, enabled: bool) {
+        match self {
+            TbfHeader::TbfHeaderV2(hd) => {
+                let old_flags = hd.base.flags;
+                if enabled {
+                    hd.base.flags |= Flags::Enabled.as_u32();
+                } else {
+                    hd.base.flags &= !Flags::Enabled.as_u32();
+                }
+                self.compute_checksum_flags(old_flags);
+            }
+            TbfHeader::Padding(hd) => {
+                let old_flags = hd.flags;
+                if enabled {
+                    hd.flags |= Flags::Enabled.as_u32();
+                } else {
+                    hd.flags &= !Flags::Enabled.as_u32();
+                }
+                self.compute_checksum_flags(old_flags);
+            }
+        }
+    }
+
+    pub fn set_sticky(&mut self, sticky: bool) {
+        match self {
+            TbfHeader::TbfHeaderV2(hd) => {
+                let old_flags = hd.base.flags;
+                if sticky {
+                    hd.base.flags |= Flags::Sticky.as_u32();
+                } else {
+                    hd.base.flags &= !Flags::Sticky.as_u32();
+                }
+                self.compute_checksum_flags(old_flags);
+            }
+            TbfHeader::Padding(hd) => {
+                let old_flags = hd.flags;
+                if sticky {
+                    hd.flags |= Flags::Sticky.as_u32();
+                } else {
+                    hd.flags &= !Flags::Sticky.as_u32();
+                }
+                self.compute_checksum_flags(old_flags);
+            }
+        }
+    }
+
     /// Return total size of the application.
     pub fn total_size(&self) -> u32 {
         match *self {
@@ -928,6 +1009,31 @@ impl TbfHeader {
         }
     }
 
+    pub fn compute_checksum_flags(&mut self, old_flags: u32) {
+        match self {
+            TbfHeader::TbfHeaderV2(hd) => {
+                let mut checksum = hd.base.checksum;
+
+                // remove contribution of old flags
+                checksum ^= old_flags;
+
+                // add new flags contribution
+                checksum ^= hd.base.flags;
+
+                // we cannot compute the full checksum using only the base, we reuse the old checksum
+                hd.base.checksum = checksum;
+            }
+            TbfHeader::Padding(hd) => {
+                let mut checksum = hd.checksum;
+
+                checksum ^= old_flags;
+                checksum ^= hd.flags;
+
+                hd.checksum = checksum;
+            }
+        }
+    }
+
     /// Return flags of the application
     ///
     /// DELTA: Originally did not exist
@@ -1210,4 +1316,18 @@ impl TbfHeader {
             _ => None,
         }
     }
+
+    pub fn serialize(&self) -> [u8; 16] {
+        let base = match self {
+            TbfHeader::TbfHeaderV2(hd) => &hd.base,
+            TbfHeader::Padding(base) => base,
+        };
+        let mut bytes = [0u8; 16];
+        bytes[0..2].copy_from_slice(&base.version.to_le_bytes());
+        bytes[2..4].copy_from_slice(&base.header_size.to_le_bytes());
+        bytes[4..8].copy_from_slice(&base.total_size.to_le_bytes());
+        bytes[8..12].copy_from_slice(&base.flags.to_le_bytes());
+        bytes[12..16].copy_from_slice(&base.checksum.to_le_bytes());
+        bytes
+    }
 }

tbf-parser/tests/serialization.rs

@@ -0,0 +1,158 @@
+use tbf_parser::parse::*;
+use tbf_parser::types::Flags;
+
+// Serialization
+
+#[test]
+fn serialize_identical_with_original() {
+    let buffer: Vec<u8> = include_bytes!("./flashes/simple.dat").to_vec();
+
+    let (_, header_len, _) = parse_tbf_header_lengths(&buffer[0..8].try_into().unwrap())
+        .ok()
+        .unwrap();
+
+    let header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    let serialized = header.serialize();
+
+    // Check if serialize matches original buffer
+    assert_eq!(&buffer[0..16], &serialized[..]);
+}
+
+// Flag modifications
+#[test]
+fn flags_modifications() {
+    let mut buffer = include_bytes!("./flashes/footerSHA256.dat").to_vec();
+    let (_, header_len, _) = parse_tbf_header_lengths(&buffer[0..8].try_into().unwrap())
+        .ok()
+        .unwrap();
+
+    let mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(header.enabled());
+    header.set_sticky(true);
+    // Set sticky without parsing
+    assert!(header.sticky());
+    // Unset
+    header.set_sticky(false);
+
+    // Disable
+    header.set_enabled(false);
+    let serialized = header.serialize();
+    buffer[0..16].copy_from_slice(&serialized);
+
+    let reparsed = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(!reparsed.enabled());
+
+    // Enable
+    let mut header = reparsed;
+    header.set_enabled(true);
+    let serialized = header.serialize();
+    buffer[0..16].copy_from_slice(&serialized);
+
+    let reparsed = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(reparsed.enabled());
+}
+
+#[test]
+fn padding_header_set_flags() {
+    let buffer = vec![
+        0x02, 0x00, 0x10, 0x00, 0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x12, 0x00, 0x10,
+        0x00,
+    ];
+
+    let (_, header_len, _) = parse_tbf_header_lengths(&buffer[0..8].try_into().unwrap())
+        .ok()
+        .unwrap();
+    let mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+
+    assert!(!header.is_app());
+
+    let flags = Flags::Enabled | Flags::Sticky;
+    header.set_flags(flags);
+
+    let serialized = header.serialize();
+    let flags = u32::from_le_bytes(serialized[8..12].try_into().unwrap());
+    assert_eq!(flags, (Flags::Enabled | Flags::Sticky).as_u32());
+}
+
+#[test]
+fn fields_preserved() {
+    let buffer = include_bytes!("./flashes/simple.dat").to_vec();
+    let (_, header_len, _) = parse_tbf_header_lengths(&buffer[0..8].try_into().unwrap())
+        .ok()
+        .unwrap();
+
+    let mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+
+    let header_size = header.header_size();
+    let total_size = header.total_size();
+
+    header.set_flags(Flags::Enabled | Flags::Sticky);
+    let serialized = header.serialize();
+
+    let _version = u16::from_le_bytes(serialized[0..2].try_into().unwrap());
+    let _header_size = u16::from_le_bytes(serialized[2..4].try_into().unwrap());
+    let _total_size = u32::from_le_bytes(serialized[4..8].try_into().unwrap());
+
+    // Check other fields are unchanged
+    assert_eq!(_version, 2);
+    assert_eq!(header_size, _header_size);
+    assert_eq!(total_size, _total_size);
+}
+
+// Complete use //
+#[test]
+fn serialization_multiple_checks() {
+    let mut buffer = include_bytes!("./flashes/footerRSA4096.dat").to_vec();
+    let (_, header_len, _) = parse_tbf_header_lengths(&buffer[0..8].try_into().unwrap())
+        .ok()
+        .unwrap();
+
+    let mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(header.enabled());
+
+    // Disable
+    header.set_enabled(false);
+    let serialized = header.serialize();
+    buffer[0..16].copy_from_slice(&serialized);
+
+    let header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(!header.enabled());
+
+    // Enable and set sticky
+    let mut header = header;
+    header.set_enabled(true);
+    header.set_sticky(true);
+    let serialized = header.serialize();
+    buffer[0..16].copy_from_slice(&serialized);
+
+    let header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(header.enabled());
+    assert!(header.sticky());
+
+    // Disable sticky with high bits
+    let flags = !Flags::Sticky;
+    let mut header = header;
+    header.set_flags(flags);
+    let serialized = header.serialize();
+    buffer[0..16].copy_from_slice(&serialized);
+
+    let header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(header.enabled());
+    assert!(!header.sticky());
+    let flags_buffer = u32::from_le_bytes(buffer[8..12].try_into().unwrap());
+    assert_eq!(flags_buffer, (!Flags::Sticky).as_u32());
+}
+
+#[test]
+fn corrupt() {
+    let mut buffer = include_bytes!("./flashes/simple.dat").to_vec();
+    let (_, header_len, _) = parse_tbf_header_lengths(&buffer[0..8].try_into().unwrap())
+        .ok()
+        .unwrap();
+
+    // Corrupt the checksum manually and check for parsing error
+    buffer[12] ^= 0x6D;
+
+    let result = parse_tbf_header(&buffer[0..header_len as usize], 2);
+    assert!(result.is_err());
+}