diff --git a/tbf-parser/Cargo.toml b/tbf-parser/Cargo.toml index 015549b..2a80fc2 100644 --- a/tbf-parser/Cargo.toml +++ b/tbf-parser/Cargo.toml @@ -7,6 +7,7 @@ edition = "2021" [dependencies] log = "0.4.27" +bitflags = "2.11.1" [features] default = [] diff --git a/tbf-parser/src/types.rs b/tbf-parser/src/types.rs index e8b8697..af4ef77 100644 --- a/tbf-parser/src/types.rs +++ b/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 + } } diff --git a/tbf-parser/tests/serialization.rs b/tbf-parser/tests/serialization.rs new file mode 100644 index 0000000..fb4f274 --- /dev/null +++ b/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 = 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()); +}