From bbc1970bbebb76a15811964dd023193b3ab238a7 Mon Sep 17 00:00:00 2001 From: Florian Stecker Date: Sun, 4 Feb 2024 17:52:46 -0500 Subject: [PATCH] last somewhat working version before reimplementing tree traversal --- src/btrfs_lookup.rs | 48 ++++++++++++++++++++++++++++++++------------ src/btrfs_structs.rs | 26 ++++++++++++++++++++---- src/main.rs | 26 +++++++++++------------- src/test.rs | 20 +++++++++++++++++- 4 files changed, 88 insertions(+), 32 deletions(-) diff --git a/src/btrfs_lookup.rs b/src/btrfs_lookup.rs index 0c409bb..80bef60 100644 --- a/src/btrfs_lookup.rs +++ b/src/btrfs_lookup.rs @@ -1,9 +1,11 @@ use std::rc::Rc; -use std::ops::Deref; +use std::ops::{Deref, RangeBounds}; use crate::btrfs_structs::{Leaf, Key, Item, InteriorNode, Node, ParseError, ParseBin, Value, Superblock, ItemType}; use crate::addrmap::{node_at_log, LogToPhys, AddressMap}; +/// 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: Rc, @@ -16,7 +18,7 @@ impl<'a> Tree<'a> { let superblock = Superblock::parse(&image[0x10000..])?; let root_tree = Tree { - image: image, + image, addr_map: Rc::clone(&addr_map), root_addr_log: superblock.root }; @@ -94,7 +96,7 @@ impl Tree<'_> { /***** iterator *****/ -pub struct Iter<'a> { +pub struct RangeIter<'a, R: RangeBounds, F: Fn(Key) -> Key = fn(Key) -> Key> { tree: &'a Tree<'a>, // path to the last returned item @@ -102,27 +104,28 @@ pub struct Iter<'a> { leaf: Option>, indices: Vec, - lower_limit: Option, - upper_limit: Option, + bounds: R, + skip_fn: F, } impl Tree<'_> { - pub fn iter<'a>(&'a self) -> Iter<'a> { + pub fn iter<'a>(&'a self) -> RangeIter<'a> { self.range(None, None) } - pub fn range<'a>(&'a self, lower: Option, upper: Option) -> Iter<'a> { - Iter { + pub fn range<'a>(&'a self, lower: Option, upper: Option) -> RangeIter<'a> { + RangeIter { tree: self, nodes: Vec::new(), leaf: None, indices: Vec::new(), // in nodes and leaf lower_limit: lower, upper_limit: upper, + skip_fn: |x|x } } - pub fn range_id<'a>(&'a self, id: u64) -> Iter<'a> { + pub fn range_id<'a>(&'a self, id: u64) -> RangeIter<'a> { if id == u64::MAX { self.range( Some(Key::new(id, ItemType::Invalid, 0)), @@ -135,9 +138,29 @@ impl Tree<'_> { ) } } + + /// given a tree, a range of indices, and two "skip functions", produces a double + /// ended iterator which iterates through the keys contained in the range, in ascending + /// or descending order. + + /// the skip functions are ignored for now, but are intended as an optimization: + /// after a key `k` was returned by the iterator (or the reverse iterator), all keys + /// strictly lower than `forward_skip_fn(k)` are skipped (resp. all keys strictly above + /// `backward_skip_fn` are skipped. + pub fn range_with_skip<'a, R, F>(&'a self, range: R, forward_skip_fn: F, backward_skip_fn: F) -> RangeIter<'a, F> + where + R: RangeBounds, + F: Fn(Key) -> Key { + RangeIter { + tree: self, + nodes: Vec::new(), + leaf: None, + indices: Vec::new(), + } + } } -impl Iter<'_> { +impl Key> RangeIter<'_, F> { fn move_down_and_get_first_item(&mut self, mut node_addr: u64) -> Option { loop { let node = Node::parse(node_at_log(self.tree.image, self.tree.addr_map.deref(), node_addr).ok()?).ok()?; @@ -190,10 +213,9 @@ impl Iter<'_> { } } } - } -impl Iterator for Iter<'_> { +impl Key> Iterator for RangeIter<'_, F> { type Item = Item; // for now we just silently stop when we encounter an error, maybe that isn't the best solution @@ -267,7 +289,7 @@ impl Iterator for Iter<'_> { }; // first first item under this node - self.move_down_and_get_first_item(node_addr) + return self.move_down_and_get_first_item(node_addr) .filter(|item|self.upper_limit.is_none() || item.key < self.upper_limit.unwrap()) } } diff --git a/src/btrfs_structs.rs b/src/btrfs_structs.rs index ecd17d5..9bcc3d6 100644 --- a/src/btrfs_structs.rs +++ b/src/btrfs_structs.rs @@ -63,6 +63,16 @@ pub struct Key { pub key_offset: u64, } +impl Key { + pub fn new(key_id: u64, key_type: ItemType, key_offset: u64) -> Key { + Key { key_id, key_type, key_offset } + } + + pub fn id(key_id: u64) -> Key { + Key { key_id, key_type: ItemType::Invalid, key_offset: 0 } + } +} + #[allow(unused)] #[derive(Debug,Clone)] pub enum Value { @@ -709,8 +719,16 @@ impl fmt::Debug for Checksum { } } -impl Key { - pub fn new(key_id: u64, key_type: ItemType, key_offset: u64) -> Key { - Key { key_id, key_type, key_offset } - } +macro_rules! key { + ($arg1:expr) => { + btrfs_structs::Key { key_id: $arg1, key_type: btrfs_structs::ItemType::Invalid, key_offset: 0 } + }; + ($arg1:expr, $arg2:expr) => { + btrfs_structs::Key { key_id: $arg1, key_type: $arg2, key_offset: 0 } + }; + ($arg1:expr, $arg2:expr, $arg3:expr) => { + btrfs_structs::Key { key_id: $arg1, key_type: $arg2, key_offset: $arg3 } + }; } + +pub(crate) use key; diff --git a/src/main.rs b/src/main.rs index c55be77..c93f4ca 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,17 +1,15 @@ -use memmap2::{Mmap, MmapOptions}; -use std::fs::File; -use rouille::Request; -use rouille::Response; -use rouille::router; -use std::iter; -use std::env; -use std::{fs::OpenOptions, os::unix::fs::OpenOptionsExt}; -use std::collections::HashMap; - -use parsebtrfs::btrfs_structs::{TreeID, Value::Extent, Value::BlockGroup, ParseError, NODE_SIZE, ItemType, Node, ParseBin}; -use parsebtrfs::btrfs_lookup::Tree; - -use parsebtrfs::addrmap::{AddressMap, LogToPhys}; +use std::{ + iter, + env, + fs::OpenOptions, + collections::HashMap, +}; +use memmap2::Mmap; +use rouille::{Request, Response, router}; +use parsebtrfs::{ + btrfs_structs::{TreeID, Value::Extent, Value::BlockGroup, ParseError, NODE_SIZE, ItemType}, + btrfs_lookup::Tree, +}; const COLORS: &[&str] = &["#e6194b", "#3cb44b", "#ffe119", "#4363d8", "#f58231", "#911eb4", "#46f0f0", "#f032e6", "#bcf60c", "#fabebe", "#008080", "#e6beff", "#9a6324", "#fffac8", "#800000", "#aaffc3", "#808000", "#ffd8b1", "#000075", "#808080", "#000000"]; diff --git a/src/test.rs b/src/test.rs index 65451cb..5c4397d 100644 --- a/src/test.rs +++ b/src/test.rs @@ -1,5 +1,9 @@ use super::*; -use btrfs_structs::{Key, ItemType}; +use btrfs_structs::key; +use btrfs_structs::{Key, ItemType, TreeID}; +use btrfs_lookup::Tree; +use memmap2::Mmap; +use std::fs::OpenOptions; #[test] fn test_key_new() { @@ -8,3 +12,17 @@ fn test_key_new() { Key { key_id: 1, key_type: ItemType::Root, key_offset: 2 } ); } + +#[test] +fn image_test() { + let file = OpenOptions::new().read(true).open("../image").unwrap(); + let image = unsafe { Mmap::map(&file).unwrap() }; + + let extent_tree = Tree::new(&image, TreeID::Extent).unwrap(); + + let skip = |Key { key_id, key_type, key_offset }| Key { key_id: key_id + 0x10000, key_type, key_offset }; + + for item in extent_tree.range_with_skip(Some(key!(0x1d70000)), Some(key!(0x2000000)), skip) { + println!("{item:x?}"); + } +}