btrfs_explorer/btrfs_explorer/src/http_tree.rs

use std::{
	str::FromStr,
};
use rouille::{Request, Response};
use crate::{
	btrfs_structs::{ItemType, Item, Key, ZERO_KEY, LAST_KEY},
	btrfs_lookup::Tree,
	render_tree::{render_table, TableResult},
	main_error::MainError,
};

enum TreeDisplayMode {
	// (x,y,z): Highlight key_id x, show y keys before (excluding x*), show z keys after (including x*)
	Highlight(u64, usize, usize),
	// (x, y): Show y keys starting at x, including x
	From(Key, usize),
	// (x, y): Show y keys before y, excluding y
	To(Key, usize),
}


fn http_tree_internal(tree: &Tree, tree_id: u64, mode: TreeDisplayMode) -> Response {
	let mut items: Vec<(Item, u64)>;
	let mut highlighted_key_id: Option<u64> = None;

	match mode {
		TreeDisplayMode::Highlight(key_id, before, after) => {
			let key = Key {key_id, key_type: ItemType::Invalid, key_offset: 0 };
			items = tree.range_with_node_addr(..key).rev().take(before).collect();
			items.reverse();
			items.extend(tree.range_with_node_addr(key..).take(after));
			highlighted_key_id = Some(key_id);
		},
		TreeDisplayMode::From(key, num_lines) => {
			items = tree.range_with_node_addr(key..).take(num_lines).collect();
			if items.len() < num_lines {
				items.reverse();
				items.extend(tree.range_with_node_addr(..key).rev().take(num_lines - items.len()));
				items.reverse();
			}
		},
		TreeDisplayMode::To(key, num_lines) => {
			items = tree.range_with_node_addr(..key).rev().take(num_lines).collect();
			items.reverse();
			if items.len() < num_lines {
				items.extend(tree.range_with_node_addr(key..).take(num_lines - items.len()));
			}
		}
	};

	let data_slice = |item: &Item, node_addr: u64| -> &[u8] {
		tree.reader.get_raw_data(node_addr, item.range.0, item.range.1).unwrap()
	};

	let table_result = TableResult {
		tree_id,
		tree_desc: root_key_desc(tree_id).map(|x|x.to_string()),
		key_id: highlighted_key_id,
		items: items.iter().map(|(it, addr)|(it, *addr, data_slice(it, *addr))).collect(),
		first_key: items.first().map(|x|x.0.key).unwrap_or(LAST_KEY),
		last_key: items.last().map(|x|x.0.key).unwrap_or(ZERO_KEY),
	};

	Response::html(render_table(table_result))
}

fn root_key_desc(id: u64) -> Option<&'static str> {
	match id {
		1 => Some("root"),
		2 => Some("extent"),
		3 => Some("chunk"),
		4 => Some("device"),
		5 => Some("filesystem"),
		6 => Some("root directory"),
		7 => Some("checksum"),
		8 => Some("quota"),
		9 => Some("UUID"),
		10 => Some("free space"),
		11 => Some("block group"),
		0xffff_ffff_ffff_fff7 => Some("data reloc"),
		_ => None,
	}
}

fn http_tree_parse_parameters(method: Option<&str>, key: Option<&str>) -> Result<TreeDisplayMode, MainError> {
	let result = match key {
		None => TreeDisplayMode::From(ZERO_KEY, 50),
		Some(key) => {
			let components: Vec<&str> = key.split('-').collect();

			match method {
				None => {
					if components.len() < 1 {
						return Err(MainError(format!("Invalid key: {key}")))
					}
					let key_id = u64::from_str_radix(components[0], 16)?;
					TreeDisplayMode::Highlight(key_id, 10, 40)
				},
				Some(method) => {
					if components.len() < 3 {
						return Err(MainError(format!("Invalid key: {key}")))
					}

					let key_id = u64::from_str_radix(components[0], 16)?;
					let key_type: ItemType = u8::from_str_radix(components[1], 16)?.into();
					let key_offset = u64::from_str_radix(components[2], 16)?;
					let key = Key {key_id, key_type, key_offset };

					if method == "from" {
						TreeDisplayMode::From(key, 50)
					} else if method == "to" {
						TreeDisplayMode::To(key, 50)
					} else {
						return Err(MainError(format!("not a valid method: {method}")))
					}
				}
			}
		}
	};

	Ok(result)
}

pub fn http_tree(image: &[u8], tree_id: &str, method: Option<&str>, key: Option<&str>, _req: &Request) -> Result<Response, MainError> {
	let tree_display_mode = http_tree_parse_parameters(method, key)?;

	let tree_id = u64::from_str(tree_id).unwrap();
	let tree = if tree_id == 1 {
		Tree::root(image).unwrap()
	} else if tree_id == 3 {
		Tree::chunk(image).unwrap()
	} else {
		Tree::new(image, tree_id).unwrap()
	};

	Ok(http_tree_internal(&tree, tree_id, tree_display_mode))
}

pub fn http_root(image: &[u8], _key: Option<&str>, _req: &Request) -> Response {
	let tree = Tree::root(image).unwrap();
	http_tree_internal(&tree, 1, TreeDisplayMode::From(ZERO_KEY, 100))
}