Merge pull request ReFirmLabs#787 from ReFirmLabs/dts

Replaced external DTB extractor with an internal one

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 \

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<String> = 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
+}

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<DTBHeader, StructureError> {
 
     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<u8>,
+    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)
+    }
+}