use crate::error::{Error, Result};
use crate::hash::{name_hash, HashVersion, NameHash};
const DX_ROOT_RESERVED_ZERO: u32 = 0;
#[derive(Debug, Clone, Copy)]
pub struct DxRootInfo {
pub reserved_zero: u32,
pub hash_version: HashVersion,
pub info_length: u8,
pub indirect_levels: u8,
pub unused_flags: u8,
}
#[derive(Debug, Clone, Copy)]
pub struct DxCountLimit {
pub limit: u16,
pub count: u16,
}
impl DxCountLimit {
pub fn parse(buf: &[u8]) -> Result<Self> {
if buf.len() < 4 {
return Err(Error::Corrupt("dx count_limit buffer too small"));
}
Ok(Self {
limit: u16::from_le_bytes(buf[0..2].try_into().unwrap()),
count: u16::from_le_bytes(buf[2..4].try_into().unwrap()),
})
}
}
#[derive(Debug, Clone, Copy)]
pub struct DxEntry {
pub hash: u32,
pub block: u32,
}
impl DxEntry {
pub const SIZE: usize = 8;
pub fn parse(buf: &[u8]) -> Result<Self> {
if buf.len() < Self::SIZE {
return Err(Error::Corrupt("dx entry buffer too small"));
}
Ok(Self {
hash: u32::from_le_bytes(buf[0..4].try_into().unwrap()),
block: u32::from_le_bytes(buf[4..8].try_into().unwrap()),
})
}
}
pub fn parse_root_info(block: &[u8]) -> Result<DxRootInfo> {
if block.len() < 32 {
return Err(Error::Corrupt("dx root block too small"));
}
let reserved = u32::from_le_bytes(block[24..28].try_into().unwrap());
if reserved != DX_ROOT_RESERVED_ZERO {
return Err(Error::Corrupt("dx root reserved must be 0"));
}
let hash_version_byte = block[28];
let hash_version = HashVersion::from_u8(hash_version_byte)
.ok_or(Error::Corrupt("unknown htree hash version"))?;
Ok(DxRootInfo {
reserved_zero: reserved,
hash_version,
info_length: block[29],
indirect_levels: block[30],
unused_flags: block[31],
})
}
pub fn parse_root_entries(block: &[u8]) -> Result<(DxCountLimit, Vec<DxEntry>)> {
let info = parse_root_info(block)?;
let cl_offset = 24 + info.info_length as usize;
if cl_offset + 4 > block.len() {
return Err(Error::Corrupt("dx root count_limit out of range"));
}
let cl = DxCountLimit::parse(&block[cl_offset..cl_offset + 4])?;
let (cl, mut entries) = parse_entries_from(block, cl_offset, cl)?;
if let Some(first) = entries.first_mut() {
first.hash = 0;
}
Ok((cl, entries))
}
pub fn parse_node_entries(block: &[u8]) -> Result<(DxCountLimit, Vec<DxEntry>)> {
if block.len() < 8 + 4 {
return Err(Error::Corrupt("dx node block too small"));
}
let cl = DxCountLimit::parse(&block[8..12])?;
let (cl, mut entries) = parse_entries_from(block, 8, cl)?;
if let Some(first) = entries.first_mut() {
first.hash = 0;
}
Ok((cl, entries))
}
fn parse_entries_from(
block: &[u8],
start: usize,
cl: DxCountLimit,
) -> Result<(DxCountLimit, Vec<DxEntry>)> {
if cl.count == 0 {
return Err(Error::Corrupt("dx node has zero entries"));
}
if cl.count > cl.limit {
return Err(Error::Corrupt("dx count > limit"));
}
let need_bytes = cl.count as usize * DxEntry::SIZE;
if start + need_bytes > block.len() {
return Err(Error::Corrupt("dx entries overflow block"));
}
let mut out = Vec::with_capacity(cl.count as usize);
for i in 0..cl.count as usize {
let off = start + i * DxEntry::SIZE;
out.push(DxEntry::parse(&block[off..off + DxEntry::SIZE])?);
}
Ok((cl, out))
}
pub fn find_entry_for_hash(entries: &[DxEntry], target: u32) -> &DxEntry {
debug_assert!(!entries.is_empty(), "dx node must have ≥1 entry");
let mut chosen = &entries[0];
for e in &entries[1..] {
if e.hash <= target {
chosen = e;
} else {
break;
}
}
chosen
}
pub enum NextStep {
Leaf(u32),
Inner(u32),
}
pub fn target_hash(name: &[u8], root_block: &[u8], hash_seed: &[u32; 4]) -> Result<NameHash> {
let info = parse_root_info(root_block)?;
Ok(name_hash(name, info.hash_version, hash_seed))
}
pub fn lookup_leaf<R>(
name: &[u8],
root_block: &[u8],
hash_seed: &[u32; 4],
mut read_dx_block: R,
) -> Result<Option<u32>>
where
R: FnMut(u32) -> Result<Vec<u8>>,
{
let info = parse_root_info(root_block)?;
let hash = name_hash(name, info.hash_version, hash_seed);
let (_cl, entries) = parse_root_entries(root_block)?;
if entries.is_empty() {
return Ok(None);
}
let mut target = find_entry_for_hash(&entries, hash.major).block;
for _level in 0..info.indirect_levels {
let block = read_dx_block(target)?;
let (_cl, entries) = parse_node_entries(&block)?;
if entries.is_empty() {
return Ok(None);
}
target = find_entry_for_hash(&entries, hash.major).block;
}
Ok(Some(target))
}
#[cfg(test)]
mod tests {
use super::*;
fn synth_root_block() -> Vec<u8> {
let block_size = 4096;
let mut buf = vec![0u8; block_size];
buf[0..4].copy_from_slice(&2u32.to_le_bytes());
buf[4..6].copy_from_slice(&12u16.to_le_bytes());
buf[6] = 1;
buf[7] = 2;
buf[8] = b'.';
buf[12..16].copy_from_slice(&2u32.to_le_bytes());
buf[16..18].copy_from_slice(&((block_size - 12) as u16).to_le_bytes());
buf[18] = 2;
buf[19] = 2;
buf[20] = b'.';
buf[21] = b'.';
buf[24..28].copy_from_slice(&0u32.to_le_bytes());
buf[28] = 2; buf[29] = 8; buf[30] = 0; buf[31] = 0;
buf[32..34].copy_from_slice(&200u16.to_le_bytes()); buf[34..36].copy_from_slice(&2u16.to_le_bytes()); buf[36..40].copy_from_slice(&1u32.to_le_bytes());
buf[40..44].copy_from_slice(&0x8000_0000u32.to_le_bytes());
buf[44..48].copy_from_slice(&2u32.to_le_bytes());
buf
}
#[test]
fn parses_root_info_correctly() {
let buf = synth_root_block();
let info = parse_root_info(&buf).expect("parse");
assert_eq!(info.hash_version, HashVersion::Tea);
assert_eq!(info.info_length, 8);
assert_eq!(info.indirect_levels, 0);
}
#[test]
fn parses_root_entries_count_2() {
let buf = synth_root_block();
let (cl, entries) = parse_root_entries(&buf).expect("parse entries");
assert_eq!(cl.count, 2);
assert_eq!(cl.limit, 200);
assert_eq!(entries.len(), 2);
assert_eq!(entries[0].block, 1);
assert_eq!(entries[1].block, 2);
assert_eq!(entries[1].hash, 0x8000_0000);
}
#[test]
fn find_entry_for_hash_picks_correct_slot() {
let buf = synth_root_block();
let (_cl, entries) = parse_root_entries(&buf).unwrap();
let e = find_entry_for_hash(&entries, 0x1000);
assert_eq!(e.block, 1);
let e = find_entry_for_hash(&entries, 0x9000_0000);
assert_eq!(e.block, 2);
let e = find_entry_for_hash(&entries, 0x8000_0000);
assert_eq!(e.block, 2);
}
#[test]
fn lookup_leaf_zero_indirect_picks_correct_block() {
let buf = synth_root_block();
let seed = [0u32; 4];
let result = lookup_leaf(b"any-name", &buf, &seed, |_| {
unreachable!("should not descend with 0 indirect levels")
})
.expect("lookup");
assert!(result.is_some());
let block = result.unwrap();
assert!(block == 1 || block == 2, "expected 1 or 2, got {block}");
}
#[test]
fn rejects_bad_hash_version() {
let mut buf = synth_root_block();
buf[28] = 99; let result = parse_root_info(&buf);
assert!(matches!(result, Err(Error::Corrupt(_))));
}
#[test]
fn rejects_nonzero_reserved() {
let mut buf = synth_root_block();
buf[24] = 1; let result = parse_root_info(&buf);
assert!(matches!(result, Err(Error::Corrupt(_))));
}
}