From ad3f782c675e9158a7cb8578b92d69aba59f8a68 Mon Sep 17 00:00:00 2001 From: Florian Stecker Date: Tue, 13 Feb 2024 20:56:15 -0500 Subject: [PATCH] split node reading logic off from tree --- src/addrmap.rs | 17 ++++++++++++++--- src/btrfs_lookup.rs | 40 +++++++++++++++++----------------------- src/lib.rs | 1 + src/main.rs | 8 +++++--- 4 files changed, 37 insertions(+), 29 deletions(-) diff --git a/src/addrmap.rs b/src/addrmap.rs index 9643ea8..0db0a67 100644 --- a/src/addrmap.rs +++ b/src/addrmap.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use crate::btrfs_structs::{ParseBin, Key, ChunkItem, Value, Superblock, ParseError, NODE_SIZE}; use crate::btrfs_lookup::Tree; +use crate::nodereader::NodeReader; #[derive(Debug, Clone)] pub struct AddressMap(pub Vec<(u64,u64,Vec<(u64,u64)>)>); @@ -12,10 +13,11 @@ impl AddressMap { pub fn new(image: &[u8]) -> Result { let superblock = Superblock::parse(&image[0x10000..])?; let bootstrap_addr = AddressMap::from_superblock(&superblock)?; + let reader = Rc::new(NodeReader::with_addrmap(image, bootstrap_addr)?); let chunk_tree = Tree { - image: image, - addr_map: Arc::new(bootstrap_addr), + image, + reader, root_addr_log: superblock.chunk_root, }; @@ -90,7 +92,7 @@ impl LogToPhys for AddressMap { } } - +/* pub fn node_at_log<'a, T: LogToPhys>(image: &'a [u8], addr: &T, log: u64) -> Result<&'a [u8], ParseError> { if let Some(phys_addr) = addr.to_phys(log) { Ok(&image[phys_addr as usize .. phys_addr as usize + NODE_SIZE]) @@ -98,7 +100,16 @@ pub fn node_at_log<'a, T: LogToPhys>(image: &'a [u8], addr: &T, log: u64) -> Res err!("Logical address {:x} could not be translated to physical address", log) } } +*/ pub trait LogToPhys { fn to_phys(&self, log: u64) -> Option; + + fn node_at_log<'a>(&self, image: &'a [u8], log: u64) -> Result<&'a [u8], ParseError> { + if let Some(phys_addr) = self.to_phys(log) { + Ok(&image[phys_addr as usize .. phys_addr as usize + NODE_SIZE]) + } else { + err!("Logical address {:x} could not be translated to physical address", log) + } + } } diff --git a/src/btrfs_lookup.rs b/src/btrfs_lookup.rs index 843a17b..184887a 100644 --- a/src/btrfs_lookup.rs +++ b/src/btrfs_lookup.rs @@ -1,27 +1,28 @@ use std::convert::identity; use std::rc::Rc; -use std::sync::Arc; use std::ops::{Deref, RangeBounds, Bound}; use crate::btrfs_structs::{Leaf, Key, Item, InteriorNode, Node, ParseError, ParseBin, Value, Superblock, ItemType, ZERO_KEY}; -use crate::addrmap::{node_at_log, AddressMap}; +use crate::addrmap::{AddressMap, LogToPhys}; +use crate::nodereader::NodeReader; /// Represents a B-Tree inside a filesystem image. Can be used to look up keys, /// and handles the tree traversal and the virtual address translation. pub struct Tree<'a> { pub image: &'a [u8], - pub addr_map: Arc, + pub reader: Rc>, pub root_addr_log: u64, } impl<'a> Tree<'a> { pub fn new>(image: &'a [u8], tree_id: T) -> Result, ParseError> { - let addr_map = Arc::new(AddressMap::new(image)?); +// let addr_map = Arc::new(AddressMap::new(image)?); let superblock = Superblock::parse(&image[0x10000..])?; + let reader = Rc::new(NodeReader::new(image)?); let root_tree = Tree { image, - addr_map: Arc::clone(&addr_map), + reader: Rc::clone(&reader), root_addr_log: superblock.root }; let tree_root_item = root_tree.find_key(Key::new(tree_id.into(), ItemType::Root, 0))?; @@ -31,30 +32,23 @@ impl<'a> Tree<'a> { _ => return Err("root item invalid".into()) }; - Ok(Tree { image, addr_map, root_addr_log }) + Ok(Tree { image, reader: Rc::clone(&reader), root_addr_log }) } pub fn root(image: &'a [u8]) -> Result, ParseError> { - let addr_map = Arc::new(AddressMap::new(image)?); + let reader = Rc::new(NodeReader::new(image)?); +// let addr_map = Arc::new(AddressMap::new(image)?); let superblock = Superblock::parse(&image[0x10000..])?; - Ok(Tree { image, addr_map, root_addr_log: superblock.root }) + Ok(Tree { image, reader, root_addr_log: superblock.root }) } pub fn chunk(image: &'a [u8]) -> Result, ParseError> { - let addr_map = Arc::new(AddressMap::new(image)?); + let reader = Rc::new(NodeReader::new(image)?); +// let addr_map = Arc::new(AddressMap::new(image)?); let superblock = Superblock::parse(&image[0x10000..])?; - Ok(Tree { image, addr_map, root_addr_log: superblock.chunk_root }) - } - - /// Read a node given its logical address - pub fn get_node(&self, addr: u64) -> Result, ParseError> { -// eprintln!("Reading node at {:x}", addr); - - let node_data = node_at_log(self.image, self.addr_map.as_ref(), addr)?; - let node = Node::parse(node_data)?; - Ok(Arc::new(node)) + Ok(Tree { image, reader, root_addr_log: superblock.chunk_root }) } } @@ -98,7 +92,7 @@ impl Tree<'_> { /// Recursively traverse a tree to find a key, given they key and logical address /// of the tree root. Internal function, `Tree::find_key` is the public interface. fn find_key_in_node(&self, addr: u64, key: Key) -> Result { - let node = self.get_node(self.root_addr_log)?; + let node = self.reader.get_node(self.root_addr_log)?; match node.deref() { Node::Interior(interior_node) => { @@ -183,7 +177,7 @@ impl<'a> Tree<'a> { /// Get the first item under the node at logical address `addr`. /// This function panics if there are no items fn get_first_item(tree: &Tree, addr: u64) -> Result { - match tree.get_node(addr)?.deref() { + match tree.reader.get_node(addr)?.deref() { Node::Interior(intnode) => { get_first_item(tree, intnode.children[0].ptr) }, @@ -196,7 +190,7 @@ fn get_first_item(tree: &Tree, addr: u64) -> Result { /// Get the last item under the node at logical address `addr`. /// This function panics if there are no items fn get_last_item(tree: &Tree, addr: u64) -> Result { - match tree.get_node(addr)?.deref() { + match tree.reader.get_node(addr)?.deref() { Node::Interior(intnode) => { get_last_item(tree, intnode.children.last().unwrap().ptr) }, @@ -222,7 +216,7 @@ fn find_closest_key(tree: &Tree, key: Key, mode: FindKeyMode) -> Result = None; loop { - let node = tree.get_node(current)?; + let node = tree.reader.get_node(current)?; match node.deref() { Node::Interior(intnode) => { match intnode.find_key_or_previous(key) { diff --git a/src/lib.rs b/src/lib.rs index 129c3bc..a3fe97f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ pub mod util; pub mod btrfs_structs; pub mod btrfs_lookup; pub mod addrmap; +pub mod nodereader; #[cfg(test)] mod test; diff --git a/src/main.rs b/src/main.rs index e3d8c8f..b624f62 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,7 +10,8 @@ use memmap2::Mmap; use rouille::{Request, Response, router}; use parsebtrfs::{ btrfs_structs::{TreeID, Value::Extent, Value::BlockGroup, ParseError, NODE_SIZE, ItemType, Item, RootItem, Value, Key, ExtentDataBody}, - btrfs_lookup::Tree, key, + btrfs_lookup::Tree, + addrmap::AddressMap, }; const COLORS: &[&str] = &["#e6194b", "#3cb44b", "#ffe119", "#4363d8", "#f58231", "#911eb4", "#46f0f0", "#f032e6", "#bcf60c", "#fabebe", "#008080", "#e6beff", "#9a6324", "#fffac8", "#800000", "#aaffc3", "#808000", "#ffd8b1", "#000075", "#808080", "#000000"]; @@ -245,6 +246,7 @@ fn http_main_boxes(image: &[u8], _req: &Request) -> Response { }; // header + let addr_map: &AddressMap = extent_tree.reader.as_ref().addr_map(); result.push_str( &format!( "

{:x} - {:x} ({}, {})

Physical: {}

\n", @@ -264,8 +266,8 @@ fn http_main_boxes(image: &[u8], _req: &Request) -> Response { 0x04 => "Metadata", _ => "???", }, - match extent_tree.addr_map.as_ref().0.binary_search_by_key(&bg.key.key_id, |x|x.0) { - Ok(i) => format!("{:x?}", &extent_tree.addr_map.as_ref().0[i].2), + match addr_map.0.binary_search_by_key(&bg.key.key_id, |x|x.0) { + Ok(i) => format!("{:x?}", &addr_map.0[i].2), _ => String::from(""), } )