From 6109dd20fe01a1814ac3d736c116e5f2ae801e92 Mon Sep 17 00:00:00 2001 From: devttys0 Date: Sat, 30 Nov 2024 19:17:01 -0500 Subject: [PATCH] Replaced external DTB extractor with an internal one --- dependencies/ubuntu.sh | 1 - src/extractors/dtb.rs | 96 +++++++++++++++++++++++++++++++++++------- src/structures/dtb.rs | 94 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 174 insertions(+), 17 deletions(-) diff --git a/dependencies/ubuntu.sh b/dependencies/ubuntu.sh index 985f7fee8..3d3efb368 100755 --- a/dependencies/ubuntu.sh +++ b/dependencies/ubuntu.sh @@ -17,7 +17,6 @@ DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install \ git \ lz4 \ lzop \ - device-tree-compiler \ unrar \ unyaffs \ python3-pip \ diff --git a/src/extractors/dtb.rs b/src/extractors/dtb.rs index 1753ca0ab..e05e63b61 100644 --- a/src/extractors/dtb.rs +++ b/src/extractors/dtb.rs @@ -1,6 +1,9 @@ -use crate::extractors; +use crate::common::is_offset_safe; +use crate::extractors::common::{Chroot, ExtractionResult, Extractor, ExtractorType}; +use crate::structures::dtb::{parse_dtb_header, parse_dtb_node}; +use log::error; -/// Describes how to run the dtc utility to extract DTB files +/// Defines the internal extractor function for extracting Device Tree Blobs /// /// ``` /// use std::io::ErrorKind; @@ -22,20 +25,81 @@ use crate::extractors; /// } /// } /// ``` -pub fn dtb_extractor() -> extractors::common::Extractor { - extractors::common::Extractor { - utility: extractors::common::ExtractorType::External("dtc".to_string()), - extension: "dtb".to_string(), - arguments: vec![ - "-I".to_string(), // Input type: dtb - "dtb".to_string(), - "-O".to_string(), // Output type: dts - "dts".to_string(), - "-o".to_string(), // Output file name: system.dtb - "system.dtb".to_string(), - extractors::common::SOURCE_FILE_PLACEHOLDER.to_string(), - ], - exit_codes: vec![0], +pub fn dtb_extractor() -> Extractor { + Extractor { + utility: ExtractorType::Internal(extract_dtb), ..Default::default() } } + +/// Internal extractor for extracting Device Tree Blobs +pub fn extract_dtb( + file_data: &[u8], + offset: usize, + output_directory: Option<&String>, +) -> ExtractionResult { + let mut heirerarchy: Vec = Vec::new(); + + let mut result = ExtractionResult { + ..Default::default() + }; + + // Parse the DTB file header + if let Ok(dtb_header) = parse_dtb_header(&file_data[offset..]) { + // Get all the DTB data + if let Some(dtb_data) = file_data.get(offset..offset + dtb_header.total_size) { + // DTB node entries start at the structure offset specified in the DTB header + let mut entry_offset = dtb_header.struct_offset; + let mut previous_entry_offset = None; + let available_data = dtb_data.len(); + + // Loop over all DTB node entries + while is_offset_safe(available_data, entry_offset, previous_entry_offset) { + // Parse the next DTB node entry + let node = parse_dtb_node(&dtb_header, dtb_data, entry_offset); + + // Beginning of a node, add it to the heirerarchy list + if node.begin { + if !node.name.is_empty() { + heirerarchy.push(node.name.clone()); + } + // End of a node, remove it from the heirerarchy list + } else if node.end { + if !heirerarchy.is_empty() { + heirerarchy.pop(); + } + // End of the DTB structure, return success only if the whole DTB structure was parsed successfully up to the EOF marker + } else if node.eof { + result.success = true; + result.size = Some(available_data); + break; + // DTB property, extract it to disk + } else if node.property { + if output_directory.is_some() { + let chroot = Chroot::new(output_directory); + let dir_path = heirerarchy.join(std::path::MAIN_SEPARATOR_STR); + let file_path = chroot.safe_path_join(&dir_path, &node.name); + + if !chroot.create_directory(dir_path) { + break; + } + + if !chroot.create_file(file_path, &node.data) { + break; + } + } + // The only other supported node type is NOP + } else if !node.nop { + error!("Unknown or invalid DTB node"); + break; + } + + // Update offsets to parse the next DTB structure entry + previous_entry_offset = Some(entry_offset); + entry_offset += node.total_size; + } + } + } + + result +} diff --git a/src/structures/dtb.rs b/src/structures/dtb.rs index 141139cb0..eaf0378bf 100644 --- a/src/structures/dtb.rs +++ b/src/structures/dtb.rs @@ -1,3 +1,4 @@ +use crate::common::get_cstring; use crate::structures::common::{self, StructureError}; /// Struct to store DTB info @@ -67,3 +68,96 @@ pub fn parse_dtb_header(dtb_data: &[u8]) -> Result { Err(StructureError) } + +/// Describes a DTB node entry +#[derive(Debug, Default, Clone)] +pub struct DTBNode { + pub begin: bool, + pub end: bool, + pub eof: bool, + pub nop: bool, + pub property: bool, + pub name: String, + pub data: Vec, + pub total_size: usize, +} + +/// Parse a DTB node from the DTB data structure +pub fn parse_dtb_node(dtb_header: &DTBHeader, dtb_data: &[u8], node_offset: usize) -> DTBNode { + const FDT_BEGIN_NODE: usize = 1; + const FDT_END_NODE: usize = 2; + const FDT_PROP: usize = 3; + const FDT_NOP: usize = 4; + const FDT_END: usize = 9; + + let node_token = vec![("id", "u32")]; + let node_property = vec![("data_len", "u32"), ("name_offset", "u32")]; + + let mut node = DTBNode { + ..Default::default() + }; + + if let Some(node_data) = dtb_data.get(node_offset..) { + if let Ok(token) = common::parse(node_data, &node_token, "big") { + // Set total node size to the size of the token entry + node.total_size = common::size(&node_token); + + if token["id"] == FDT_END_NODE { + node.end = true; + } else if token["id"] == FDT_NOP { + node.nop = true; + } else if token["id"] == FDT_END { + node.eof = true; + // All other node types must include additional data, so the available data must be greater than just the token entry size + } else if node_data.len() > node.total_size { + if token["id"] == FDT_BEGIN_NODE { + // Begin nodes are immediately followed by a NULL-terminated name, padded to a 4-byte boundary if necessary + node.begin = true; + node.name = get_cstring(&node_data[node.total_size..]); + node.total_size += dtb_aligned(node.name.len() + 1); + } else if token["id"] == FDT_PROP { + // Property tokens are followed by a property structure + if let Ok(property) = + common::parse(&node_data[node.total_size..], &node_property, "big") + { + // Update the total node size to include the property structure + node.total_size += common::size(&node_property); + + // The property's data will immediately follow the property structure; property data may be NULL-padded for alignment + if let Some(property_data) = + node_data.get(node.total_size..node.total_size + property["data_len"]) + { + node.data = property_data.to_vec(); + node.total_size += dtb_aligned(node.data.len()); + + // Get the property name from the DTB strings table + if let Some(property_name_data) = + dtb_data.get(dtb_header.strings_offset + property["name_offset"]..) + { + node.name = get_cstring(property_name_data); + if !node.name.is_empty() { + node.property = true; + } + } + } + } + } + } + } + } + + node +} + +/// DTB entries must be aligned to 4-byte boundaries +fn dtb_aligned(len: usize) -> usize { + const ALIGNMENT: usize = 4; + + let rem = len % ALIGNMENT; + + if rem == 0 { + len + } else { + len + (ALIGNMENT - rem) + } +}