last somewhat working version before reimplementing tree traversal

This commit is contained in:
Florian Stecker 2024-02-04 17:52:46 -05:00
parent d5f7ff4c14
commit bbc1970bbe
4 changed files with 88 additions and 32 deletions

View File

@ -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<AddressMap>,
@ -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<Key>, 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<Box<Leaf>>,
indices: Vec<usize>,
lower_limit: Option<Key>,
upper_limit: Option<Key>,
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<Key>, upper: Option<Key>) -> Iter<'a> {
Iter {
pub fn range<'a>(&'a self, lower: Option<Key>, upper: Option<Key>) -> 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<Key>,
F: Fn(Key) -> Key {
RangeIter {
tree: self,
nodes: Vec::new(),
leaf: None,
indices: Vec::new(),
}
}
}
impl Iter<'_> {
impl<F: Fn(Key) -> Key> RangeIter<'_, F> {
fn move_down_and_get_first_item(&mut self, mut node_addr: u64) -> Option<Item> {
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<F: Fn(Key) -> 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())
}
}

View File

@ -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;

View File

@ -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"];

View File

@ -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?}");
}
}