From 847d9c5666dc2456790a5d76856758f2e7b84ba1 Mon Sep 17 00:00:00 2001
From: Thomas Guignard <thomas.guignard@epita.fr>
Date: Wed, 1 Oct 2025 11:34:31 +0200
Subject: [PATCH 1/4] feat: add set_flags function with serialization support

---
 tbf-parser/src/types.rs | 73 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 73 insertions(+)

diff --git a/tbf-parser/src/types.rs b/tbf-parser/src/types.rs
index 5f4dca45..24771a0c 100644
--- a/tbf-parser/src/types.rs
+++ b/tbf-parser/src/types.rs
@@ -1183,4 +1183,77 @@ impl TbfHeader {
             _ => None,
         }
     }
+
+    pub fn compute_checksum(header: &[u8], new_flags: u32) -> Result<u32, TbfParseError> {
+        let mut checksum: u32 = 0;
+
+        let header_iter = header.chunks_exact(4);
+
+        // Iterate all chunks and XOR the chunks to compute the checksum.
+        for (i, chunk) in header_iter.enumerate() {
+            let word = if i == 2 {
+                new_flags
+            } else if i == 3 {
+                continue;
+            } else {
+                u32::from_le_bytes(chunk.try_into()?)
+            };
+            checksum ^= word;
+        }
+        Ok(checksum)
+    }
+
+    pub fn set_flags(&mut self, flags: u32, header: &[u8]) -> Result<[u8; 16], TbfParseError> {
+        let new_checksum = Self::compute_checksum(header, flags)?;
+        match self {
+            TbfHeader::TbfHeaderV2(hd) => {
+                hd.base.flags = flags;
+                hd.base.checksum = new_checksum;
+            }
+            TbfHeader::Padding(base) => {
+                base.flags = flags;
+                base.checksum = new_checksum;
+            }
+        }
+
+        self.serialize()
+    }
+
+    pub fn serialize(&self) -> Result<[u8; 16], TbfParseError> {
+        match self {
+            TbfHeader::TbfHeaderV2(hd) => {
+                let base = &hd.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());
+                Ok(bytes)
+            }
+            TbfHeader::Padding(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());
+                Ok(bytes)
+            }
+        }
+    }
+
+    pub fn update_from_serialize(&mut self, bytes: &[u8]) -> Result<(), TbfParseError> {
+        let new_base = TbfHeaderV2Base::try_from(&bytes[..])?;
+        match self {
+            TbfHeader::TbfHeaderV2(hd) => {
+                hd.base = new_base;
+                Ok(())
+            }
+            TbfHeader::Padding(base) => {
+                *base = new_base;
+                Ok(())
+            }
+        }
+    }
 }

From 8430e2387191332687941d3e39a7ca19e493f02a Mon Sep 17 00:00:00 2001
From: Thomas Guignard <thomas.guignard@epita.fr>
Date: Thu, 2 Oct 2025 17:27:42 +0200
Subject: [PATCH 2/4] fix(types): removed unsued pub function, followed the
 suggestions feat(tests/serialization): add tests for TBF header serialization

---
 tbf-parser/src/types.rs           |  54 ++--
 tbf-parser/tests/serialization.rs | 437 ++++++++++++++++++++++++++++++
 2 files changed, 454 insertions(+), 37 deletions(-)
 create mode 100644 tbf-parser/tests/serialization.rs

diff --git a/tbf-parser/src/types.rs b/tbf-parser/src/types.rs
index 24771a0c..d4fade07 100644
--- a/tbf-parser/src/types.rs
+++ b/tbf-parser/src/types.rs
@@ -1184,6 +1184,8 @@ impl TbfHeader {
         }
     }
 
+    /// Returns the checksum for TBF header according to the `parse_tbf_header` function
+    /// `new_flags` is the new value we want to set
     pub fn compute_checksum(header: &[u8], new_flags: u32) -> Result<u32, TbfParseError> {
         let mut checksum: u32 = 0;
 
@@ -1203,7 +1205,8 @@ impl TbfHeader {
         Ok(checksum)
     }
 
-    pub fn set_flags(&mut self, flags: u32, header: &[u8]) -> Result<[u8; 16], TbfParseError> {
+    /// Sets the flag field and updates the checksum, it modifies the state accordingly
+    pub fn set_flags(&mut self, flags: u32, header: &[u8]) -> Result<(), TbfParseError> {
         let new_checksum = Self::compute_checksum(header, flags)?;
         match self {
             TbfHeader::TbfHeaderV2(hd) => {
@@ -1216,44 +1219,21 @@ impl TbfHeader {
             }
         }
 
-        self.serialize()
+        Ok(())
     }
 
+    /// Returns a 16 byte array with the serialized base header (little-endian format)
     pub fn serialize(&self) -> Result<[u8; 16], TbfParseError> {
-        match self {
-            TbfHeader::TbfHeaderV2(hd) => {
-                let base = &hd.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());
-                Ok(bytes)
-            }
-            TbfHeader::Padding(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());
-                Ok(bytes)
-            }
-        }
-    }
-
-    pub fn update_from_serialize(&mut self, bytes: &[u8]) -> Result<(), TbfParseError> {
-        let new_base = TbfHeaderV2Base::try_from(&bytes[..])?;
-        match self {
-            TbfHeader::TbfHeaderV2(hd) => {
-                hd.base = new_base;
-                Ok(())
-            }
-            TbfHeader::Padding(base) => {
-                *base = new_base;
-                Ok(())
-            }
-        }
+        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());
+        Ok(bytes)
     }
 }
diff --git a/tbf-parser/tests/serialization.rs b/tbf-parser/tests/serialization.rs
new file mode 100644
index 00000000..5f5e674a
--- /dev/null
+++ b/tbf-parser/tests/serialization.rs
@@ -0,0 +1,437 @@
+use tbf_parser::parse::*;
+use tbf_parser::types::TbfHeader;
+
+// Serialization
+
+#[test]
+fn simple_tbf() {
+    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().unwrap();
+
+    // Check if serialize matches original buffer
+    assert_eq!(&buffer[0..16], &serialized[..]);
+}
+
+#[test]
+fn footer_sha256() {
+    let buffer: Vec<u8> = include_bytes!("./flashes/footerSHA256.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().unwrap();
+
+    // Check if serialize matches original buffer
+    assert_eq!(&buffer[0..16], &serialized[..]);
+}
+
+#[test]
+fn footer_rsa4096() {
+    let buffer: Vec<u8> = include_bytes!("./flashes/footerRSA4096.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().unwrap();
+
+    // Check if serialize matches original buffer
+    assert_eq!(&buffer[0..16], &serialized[..]);
+}
+
+// Flag modifications
+#[test]
+fn disable_simple_tbf() {
+    let mut 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 mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+
+    assert!(header.enabled());
+    assert_eq!(header.get_package_name().unwrap(), "_heart");
+
+    // Set flag to 0 to disable the app
+    header
+        .set_flags(0x00000000, &buffer[0..header_len as usize])
+        .unwrap();
+    let serialized = header.serialize().unwrap();
+    buffer[0..16].copy_from_slice(&serialized);
+
+    // Parse again and check if disable
+    let reparsed = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(!reparsed.enabled());
+    assert_eq!(reparsed.get_package_name().unwrap(), "_heart");
+    assert_eq!(reparsed.get_minimum_app_ram_size(), 4848);
+    assert_eq!(reparsed.get_init_function_offset(), 41 + header_len as u32);
+    assert_eq!(reparsed.get_protected_size(), header_len as u32);
+    assert_eq!(reparsed.get_kernel_version().unwrap(), (2, 0));
+}
+
+#[test]
+fn enable_disable_footer_sha256() {
+    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());
+
+    // Disable
+    header
+        .set_flags(0x00000000, &buffer[0..header_len as usize])
+        .unwrap();
+    let serialized = header.serialize().unwrap();
+    buffer[0..16].copy_from_slice(&serialized);
+
+    let reparsed = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(!reparsed.enabled());
+
+    // Check
+    let mut header = reparsed;
+    header
+        .set_flags(0x00000001, &buffer[0..header_len as usize])
+        .unwrap();
+    let serialized = header.serialize().unwrap();
+    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() {
+    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 header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+
+    assert!(!header.enabled());
+    assert!(!header.is_app());
+    let serialized = header.serialize().unwrap();
+    assert_eq!(&serialized[..], &buffer[0..16]);
+}
+
+#[test]
+fn padding_header2() {
+    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 serialized = header.serialize();
+    assert!(serialized.is_ok());
+
+    assert!(header.set_flags(0x06000001, &buffer).is_ok());
+
+    let serialized = header.serialize().unwrap();
+    let flags = u32::from_le_bytes(serialized[8..12].try_into().unwrap());
+    assert_eq!(flags, 0x06000001);
+}
+
+#[test]
+fn fields() {
+    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(0x00000003, &buffer[0..header_len as usize])
+        .unwrap();
+    let serialized = header.serialize().unwrap();
+
+    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);
+}
+
+#[test]
+fn multiple_set() {
+    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();
+
+    // Trying to set multiple times to check consistency
+    for i in 1..21 {
+        header
+            .set_flags(i, &buffer[0..header_len as usize])
+            .unwrap();
+        assert_eq!(header.enabled(), i % 2 == 1);
+    }
+}
+
+#[test]
+fn no_parsing() {
+    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();
+    header
+        .set_flags(0x00000003, &buffer[0..header_len as usize])
+        .unwrap();
+
+    assert!(header.enabled());
+    assert!(header.sticky());
+}
+
+#[test]
+fn unset_set() {
+    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 s1 = header.serialize().unwrap();
+    header
+        .set_flags(0x80000001, &buffer[0..header_len as usize])
+        .unwrap();
+    let s2 = header.serialize().unwrap();
+    header
+        .set_flags(0x00000001, &buffer[0..header_len as usize])
+        .unwrap();
+    let s3 = header.serialize().unwrap();
+
+    assert_ne!(s1, s2);
+    assert_eq!(s1, s3);
+    assert_ne!(s2, s3);
+}
+
+// Checksum //
+#[test]
+fn checksum() {
+    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();
+
+    let mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+
+    // Test with array of multiple flags for checksum validation
+    for flags in [0x000FABCD, 0x00000001, 0x00000002, 0x00000003] {
+        header
+            .set_flags(flags, &buffer[0..header_len as usize])
+            .unwrap();
+        let serialized = header.serialize().unwrap();
+        buffer[0..16].copy_from_slice(&serialized);
+
+        let result = parse_tbf_header(&buffer[0..header_len as usize], 2);
+        assert!(result.is_ok(), "Checksum validation failed for {flags}");
+    }
+}
+
+#[test]
+fn checksum_footer_sha256() {
+    let buffer = include_bytes!("./flashes/footerSHA256.dat").to_vec();
+    let (_, header_len, _) = parse_tbf_header_lengths(&buffer[0..8].try_into().unwrap())
+        .ok()
+        .unwrap();
+
+    let slice = &buffer[0..header_len as usize];
+    let flags = u32::from_le_bytes(slice[8..12].try_into().unwrap());
+
+    // Compute checksum
+    let computed = TbfHeader::compute_checksum(slice, flags).unwrap();
+    let stored = u32::from_le_bytes(slice[12..16].try_into().unwrap());
+
+    // Compared with the one stored
+    assert_eq!(computed, stored);
+}
+
+#[test]
+fn empty_buffer() {
+    let empty_buffer: Vec<u8> = vec![];
+    let result = TbfHeader::compute_checksum(&empty_buffer, 0x00000006D);
+    assert_eq!(result.unwrap(), 0x00000000);
+}
+
+// Complete use //
+#[test]
+fn all_simple_tbf() {
+    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();
+
+    let mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(header.enabled());
+
+    // Disable
+    header
+        .set_flags(0x00000000, &buffer[0..header_len as usize])
+        .unwrap();
+    let serialized = header.serialize().unwrap();
+    buffer[0..16].copy_from_slice(&serialized);
+
+    let header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(!header.enabled());
+
+    // Enable
+    let mut header = header;
+    header
+        .set_flags(0x00000001, &buffer[0..header_len as usize])
+        .unwrap();
+    let serialized = header.serialize().unwrap();
+    buffer[0..16].copy_from_slice(&serialized);
+
+    let header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(header.enabled());
+}
+
+#[test]
+fn all_rsa4096() {
+    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());
+    assert_eq!(header.get_package_name().unwrap(), "c_hello");
+
+    // Disable
+    header
+        .set_flags(0x00000000, &buffer[0..header_len as usize])
+        .unwrap();
+    let serialized = header.serialize().unwrap();
+    buffer[0..16].copy_from_slice(&serialized);
+
+    let header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(!header.enabled());
+    assert_eq!(header.get_package_name().unwrap(), "c_hello");
+
+    // Enable and set sticky
+    let mut header = header;
+    header
+        .set_flags(0x00000003, &buffer[0..header_len as usize])
+        .unwrap();
+    let serialized = header.serialize().unwrap();
+    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());
+}
+
+#[test]
+fn all_together_high_bits() {
+    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();
+
+    let mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(header.enabled());
+    assert_eq!(header.get_package_name().unwrap(), "_heart");
+
+    header
+        .set_flags(0xFFFF0001, &buffer[0..header_len as usize])
+        .unwrap();
+    let serialized = header.serialize().unwrap();
+    buffer[0..16].copy_from_slice(&serialized);
+
+    let header2 = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
+    assert!(header2.enabled());
+
+    // Verify flags are well set
+    let flags_buffer = u32::from_le_bytes(buffer[8..12].try_into().unwrap());
+    assert_eq!(flags_buffer, 0xFFFF0001);
+
+    // Set to sticky and disabled
+    let mut header3 = header;
+    header3
+        .set_flags(0xFFFF0002, &buffer[0..header_len as usize])
+        .unwrap();
+    let serialized = header3.serialize().unwrap();
+    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());
+
+    // Enable sticky and change the high bits
+    let flags = 0xD6D60003;
+    let mut header = header;
+    header
+        .set_flags(flags, &buffer[0..header_len as usize])
+        .unwrap();
+    let serialized = header.serialize().unwrap();
+    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);
+
+    // Clear high bits
+    let mut header = header;
+    header
+        .set_flags(0x00000001, &buffer[0..header_len as usize])
+        .unwrap();
+    let serialized = header.serialize().unwrap();
+    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 final_flags = u32::from_le_bytes(buffer[8..12].try_into().unwrap());
+    assert_eq!(final_flags, 0x00000001);
+
+    assert_eq!(header.get_package_name().unwrap(), "_heart");
+    assert_eq!(header.get_kernel_version().unwrap(), (2, 0));
+}
+
+#[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());
+}

From 208fc1463b68f49e9ecf8ffd3fb91612433b4cdf Mon Sep 17 00:00:00 2001
From: Thomas Guignard <thomas.guignard@epita.fr>
Date: Mon, 6 Oct 2025 16:00:01 +0200
Subject: [PATCH 3/4] feat(types): add setters for enable and sticky
 application's flags fix(tests/serialization): removed redundant tests and
 make enable and sticky tests explicit

---
 tbf-parser/src/types.rs           |  32 +++-
 tbf-parser/tests/serialization.rs | 269 +++---------------------------
 2 files changed, 56 insertions(+), 245 deletions(-)

diff --git a/tbf-parser/src/types.rs b/tbf-parser/src/types.rs
index d4fade07..4a271177 100644
--- a/tbf-parser/src/types.rs
+++ b/tbf-parser/src/types.rs
@@ -1222,7 +1222,37 @@ impl TbfHeader {
         Ok(())
     }
 
-    /// Returns a 16 byte array with the serialized base header (little-endian format)
+    /// Returns the TBF Header's flags (avoid duplication code for the setting functions below).
+    fn get_flags(&self) -> u32 {
+        match self {
+            TbfHeader::TbfHeaderV2(hd) => hd.base.flags,
+            TbfHeader::Padding(base) => base.flags,
+        }
+    }
+
+    /// Enables or disables the application by setting the enabled flag.
+    pub fn set_enabled(&mut self, enabled: bool, header: &[u8]) -> Result<(), TbfParseError> {
+        let flags:u32 = if enabled {
+            self.get_flags() | 0x00000001
+        }
+        else {
+            self.get_flags() & !0x00000001
+        };
+        self.set_flags(flags, header)
+    }
+
+    /// Enable or disables erase confirmation by setting the sticky flag.
+    pub fn set_sticky(&mut self, enabled: bool, header: &[u8]) -> Result<(), TbfParseError> {
+        let flags:u32 = if enabled {
+            self.get_flags() | 0x00000002
+        }
+        else {
+            self.get_flags() & !0x00000002
+        };
+        self.set_flags(flags, header)
+    }
+
+    /// Returns a 16 byte array with the serialized base header (little-endian format).
     pub fn serialize(&self) -> Result<[u8; 16], TbfParseError> {
         let base = match self {
             TbfHeader::TbfHeaderV2(hd) => &hd.base,
diff --git a/tbf-parser/tests/serialization.rs b/tbf-parser/tests/serialization.rs
index 5f5e674a..26c6dbd9 100644
--- a/tbf-parser/tests/serialization.rs
+++ b/tbf-parser/tests/serialization.rs
@@ -4,7 +4,7 @@ use tbf_parser::types::TbfHeader;
 // Serialization
 
 #[test]
-fn simple_tbf() {
+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())
@@ -18,67 +18,9 @@ fn simple_tbf() {
     assert_eq!(&buffer[0..16], &serialized[..]);
 }
 
-#[test]
-fn footer_sha256() {
-    let buffer: Vec<u8> = include_bytes!("./flashes/footerSHA256.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().unwrap();
-
-    // Check if serialize matches original buffer
-    assert_eq!(&buffer[0..16], &serialized[..]);
-}
-
-#[test]
-fn footer_rsa4096() {
-    let buffer: Vec<u8> = include_bytes!("./flashes/footerRSA4096.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().unwrap();
-
-    // Check if serialize matches original buffer
-    assert_eq!(&buffer[0..16], &serialized[..]);
-}
-
 // Flag modifications
 #[test]
-fn disable_simple_tbf() {
-    let mut 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 mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
-
-    assert!(header.enabled());
-    assert_eq!(header.get_package_name().unwrap(), "_heart");
-
-    // Set flag to 0 to disable the app
-    header
-        .set_flags(0x00000000, &buffer[0..header_len as usize])
-        .unwrap();
-    let serialized = header.serialize().unwrap();
-    buffer[0..16].copy_from_slice(&serialized);
-
-    // Parse again and check if disable
-    let reparsed = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
-    assert!(!reparsed.enabled());
-    assert_eq!(reparsed.get_package_name().unwrap(), "_heart");
-    assert_eq!(reparsed.get_minimum_app_ram_size(), 4848);
-    assert_eq!(reparsed.get_init_function_offset(), 41 + header_len as u32);
-    assert_eq!(reparsed.get_protected_size(), header_len as u32);
-    assert_eq!(reparsed.get_kernel_version().unwrap(), (2, 0));
-}
-
-#[test]
-fn enable_disable_footer_sha256() {
+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()
@@ -86,10 +28,15 @@ fn enable_disable_footer_sha256() {
 
     let mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
     assert!(header.enabled());
+    header.set_sticky(true, &buffer[0..header_len as usize]).unwrap();
+    // Set sticky without parsing
+    assert!(header.sticky());
+    // Unset
+    header.set_sticky(false, &buffer[0..header_len as usize]).unwrap();
 
     // Disable
     header
-        .set_flags(0x00000000, &buffer[0..header_len as usize])
+        .set_enabled(false, &buffer[0..header_len as usize])
         .unwrap();
     let serialized = header.serialize().unwrap();
     buffer[0..16].copy_from_slice(&serialized);
@@ -97,10 +44,10 @@ fn enable_disable_footer_sha256() {
     let reparsed = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
     assert!(!reparsed.enabled());
 
-    // Check
+    // Enable
     let mut header = reparsed;
     header
-        .set_flags(0x00000001, &buffer[0..header_len as usize])
+        .set_enabled(true, &buffer[0..header_len as usize])
         .unwrap();
     let serialized = header.serialize().unwrap();
     buffer[0..16].copy_from_slice(&serialized);
@@ -110,24 +57,7 @@ fn enable_disable_footer_sha256() {
 }
 
 #[test]
-fn padding_header() {
-    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 header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
-
-    assert!(!header.enabled());
-    assert!(!header.is_app());
-    let serialized = header.serialize().unwrap();
-    assert_eq!(&serialized[..], &buffer[0..16]);
-}
-
-#[test]
-fn padding_header2() {
+fn padding_header_set_flags() {
     let buffer = vec![
         0x02, 0x00, 0x10, 0x00, 0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x12, 0x00, 0x10,
         0x00,
@@ -151,7 +81,7 @@ fn padding_header2() {
 }
 
 #[test]
-fn fields() {
+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()
@@ -178,7 +108,7 @@ fn fields() {
 }
 
 #[test]
-fn multiple_set() {
+fn multiple_flags_set() {
     let buffer = include_bytes!("./flashes/simple.dat").to_vec();
     let (_, header_len, _) = parse_tbf_header_lengths(&buffer[0..8].try_into().unwrap())
         .ok()
@@ -195,49 +125,14 @@ fn multiple_set() {
     }
 }
 
-#[test]
-fn no_parsing() {
-    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();
-    header
-        .set_flags(0x00000003, &buffer[0..header_len as usize])
-        .unwrap();
-
-    assert!(header.enabled());
-    assert!(header.sticky());
-}
-
-#[test]
-fn unset_set() {
-    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 s1 = header.serialize().unwrap();
-    header
-        .set_flags(0x80000001, &buffer[0..header_len as usize])
-        .unwrap();
-    let s2 = header.serialize().unwrap();
-    header
-        .set_flags(0x00000001, &buffer[0..header_len as usize])
-        .unwrap();
-    let s3 = header.serialize().unwrap();
-
-    assert_ne!(s1, s2);
-    assert_eq!(s1, s3);
-    assert_ne!(s2, s3);
-}
-
 // Checksum //
 #[test]
 fn checksum() {
+    // Try with empty_buffer
+    let empty_buffer: Vec<u8> = vec![];
+    let result = TbfHeader::compute_checksum(&empty_buffer, 0x00000006D);
+    assert_eq!(result.unwrap(), 0x00000000);
+
     let mut buffer = include_bytes!("./flashes/simple.dat").to_vec();
     let (_, header_len, _) = parse_tbf_header_lengths(&buffer[0..8].try_into().unwrap())
         .ok()
@@ -258,66 +153,9 @@ fn checksum() {
     }
 }
 
-#[test]
-fn checksum_footer_sha256() {
-    let buffer = include_bytes!("./flashes/footerSHA256.dat").to_vec();
-    let (_, header_len, _) = parse_tbf_header_lengths(&buffer[0..8].try_into().unwrap())
-        .ok()
-        .unwrap();
-
-    let slice = &buffer[0..header_len as usize];
-    let flags = u32::from_le_bytes(slice[8..12].try_into().unwrap());
-
-    // Compute checksum
-    let computed = TbfHeader::compute_checksum(slice, flags).unwrap();
-    let stored = u32::from_le_bytes(slice[12..16].try_into().unwrap());
-
-    // Compared with the one stored
-    assert_eq!(computed, stored);
-}
-
-#[test]
-fn empty_buffer() {
-    let empty_buffer: Vec<u8> = vec![];
-    let result = TbfHeader::compute_checksum(&empty_buffer, 0x00000006D);
-    assert_eq!(result.unwrap(), 0x00000000);
-}
-
 // Complete use //
 #[test]
-fn all_simple_tbf() {
-    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();
-
-    let mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
-    assert!(header.enabled());
-
-    // Disable
-    header
-        .set_flags(0x00000000, &buffer[0..header_len as usize])
-        .unwrap();
-    let serialized = header.serialize().unwrap();
-    buffer[0..16].copy_from_slice(&serialized);
-
-    let header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
-    assert!(!header.enabled());
-
-    // Enable
-    let mut header = header;
-    header
-        .set_flags(0x00000001, &buffer[0..header_len as usize])
-        .unwrap();
-    let serialized = header.serialize().unwrap();
-    buffer[0..16].copy_from_slice(&serialized);
-
-    let header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
-    assert!(header.enabled());
-}
-
-#[test]
-fn all_rsa4096() {
+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()
@@ -325,70 +163,32 @@ fn all_rsa4096() {
 
     let mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
     assert!(header.enabled());
-    assert_eq!(header.get_package_name().unwrap(), "c_hello");
 
     // Disable
     header
-        .set_flags(0x00000000, &buffer[0..header_len as usize])
+        .set_enabled(false, &buffer[0..header_len as usize])
         .unwrap();
     let serialized = header.serialize().unwrap();
     buffer[0..16].copy_from_slice(&serialized);
 
     let header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
     assert!(!header.enabled());
-    assert_eq!(header.get_package_name().unwrap(), "c_hello");
 
     // Enable and set sticky
     let mut header = header;
     header
-        .set_flags(0x00000003, &buffer[0..header_len as usize])
+        .set_enabled(true, &buffer[0..header_len as usize])
         .unwrap();
+    header.set_sticky(true, &buffer[0..header_len as usize]).unwrap();
     let serialized = header.serialize().unwrap();
     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());
-}
-
-#[test]
-fn all_together_high_bits() {
-    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();
-
-    let mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
-    assert!(header.enabled());
-    assert_eq!(header.get_package_name().unwrap(), "_heart");
-
-    header
-        .set_flags(0xFFFF0001, &buffer[0..header_len as usize])
-        .unwrap();
-    let serialized = header.serialize().unwrap();
-    buffer[0..16].copy_from_slice(&serialized);
-
-    let header2 = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
-    assert!(header2.enabled());
-
-    // Verify flags are well set
-    let flags_buffer = u32::from_le_bytes(buffer[8..12].try_into().unwrap());
-    assert_eq!(flags_buffer, 0xFFFF0001);
-
-    // Set to sticky and disabled
-    let mut header3 = header;
-    header3
-        .set_flags(0xFFFF0002, &buffer[0..header_len as usize])
-        .unwrap();
-    let serialized = header3.serialize().unwrap();
-    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());
 
-    // Enable sticky and change the high bits
-    let flags = 0xD6D60003;
+    // Disable sticky with high bits
+    let flags = 0xD6D6FFF1;
     let mut header = header;
     header
         .set_flags(flags, &buffer[0..header_len as usize])
@@ -398,28 +198,9 @@ fn all_together_high_bits() {
 
     let header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
     assert!(header.enabled());
-    assert!(header.sticky());
-
+    assert!(!header.sticky());
     let flags_buffer = u32::from_le_bytes(buffer[8..12].try_into().unwrap());
     assert_eq!(flags_buffer, flags);
-
-    // Clear high bits
-    let mut header = header;
-    header
-        .set_flags(0x00000001, &buffer[0..header_len as usize])
-        .unwrap();
-    let serialized = header.serialize().unwrap();
-    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 final_flags = u32::from_le_bytes(buffer[8..12].try_into().unwrap());
-    assert_eq!(final_flags, 0x00000001);
-
-    assert_eq!(header.get_package_name().unwrap(), "_heart");
-    assert_eq!(header.get_kernel_version().unwrap(), (2, 0));
 }
 
 #[test]

From 956175084bd6aa7111a041c853db504df4ba38ea Mon Sep 17 00:00:00 2001
From: Thomas Guignard <thomas.guignard@epita.fr>
Date: Mon, 6 Oct 2025 16:04:13 +0200
Subject: [PATCH 4/4] fix(): ran

---
 tbf-parser/src/types.rs           | 10 ++++------
 tbf-parser/tests/serialization.rs | 14 ++++++++++----
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/tbf-parser/src/types.rs b/tbf-parser/src/types.rs
index 4a271177..8491924c 100644
--- a/tbf-parser/src/types.rs
+++ b/tbf-parser/src/types.rs
@@ -1232,10 +1232,9 @@ impl TbfHeader {
 
     /// Enables or disables the application by setting the enabled flag.
     pub fn set_enabled(&mut self, enabled: bool, header: &[u8]) -> Result<(), TbfParseError> {
-        let flags:u32 = if enabled {
+        let flags: u32 = if enabled {
             self.get_flags() | 0x00000001
-        }
-        else {
+        } else {
             self.get_flags() & !0x00000001
         };
         self.set_flags(flags, header)
@@ -1243,10 +1242,9 @@ impl TbfHeader {
 
     /// Enable or disables erase confirmation by setting the sticky flag.
     pub fn set_sticky(&mut self, enabled: bool, header: &[u8]) -> Result<(), TbfParseError> {
-        let flags:u32 = if enabled {
+        let flags: u32 = if enabled {
             self.get_flags() | 0x00000002
-        }
-        else {
+        } else {
             self.get_flags() & !0x00000002
         };
         self.set_flags(flags, header)
diff --git a/tbf-parser/tests/serialization.rs b/tbf-parser/tests/serialization.rs
index 26c6dbd9..42173b02 100644
--- a/tbf-parser/tests/serialization.rs
+++ b/tbf-parser/tests/serialization.rs
@@ -28,11 +28,15 @@ fn flags_modifications() {
 
     let mut header = parse_tbf_header(&buffer[0..header_len as usize], 2).unwrap();
     assert!(header.enabled());
-    header.set_sticky(true, &buffer[0..header_len as usize]).unwrap();
+    header
+        .set_sticky(true, &buffer[0..header_len as usize])
+        .unwrap();
     // Set sticky without parsing
     assert!(header.sticky());
     // Unset
-    header.set_sticky(false, &buffer[0..header_len as usize]).unwrap();
+    header
+        .set_sticky(false, &buffer[0..header_len as usize])
+        .unwrap();
 
     // Disable
     header
@@ -155,7 +159,7 @@ fn checksum() {
 
 // Complete use //
 #[test]
-fn serialization_multiple_checks(){
+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()
@@ -179,7 +183,9 @@ fn serialization_multiple_checks(){
     header
         .set_enabled(true, &buffer[0..header_len as usize])
         .unwrap();
-    header.set_sticky(true, &buffer[0..header_len as usize]).unwrap();
+    header
+        .set_sticky(true, &buffer[0..header_len as usize])
+        .unwrap();
     let serialized = header.serialize().unwrap();
     buffer[0..16].copy_from_slice(&serialized);