split node reading logic off from tree

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<AddressMap, ParseError> {
 		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<u64>;
+
+	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)
+		}
+	}
 }

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<AddressMap>,
+	pub reader: Rc<NodeReader<'a>>,
 	pub root_addr_log: u64,
 }
 
 impl<'a> Tree<'a> {
 	pub fn new<T: Into<u64>>(image: &'a [u8], tree_id: T) -> Result<Tree<'a>, 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<Tree<'a>, 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<Tree<'a>, 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<Arc<Node>, 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<Item, ParseError> {
-		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<Item, ParseError> {
-	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<Item, ParseError> {
 /// 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<Item, ParseError> {
-	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<Option<I
 	let mut next: Option<u64> = 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) {

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;

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!(
 				"<h3 style=\"text-align: center;\">{:x} - {:x} ({}, {})</h3><p>Physical: {}</p>\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(""),
 				}
 			)