From 239ea0e378f5a1c050c1fc457242cb0b956616b1 Mon Sep 17 00:00:00 2001 From: Hunter Beast Date: Sat, 27 Jun 2026 01:42:12 -0600 Subject: [PATCH 01/12] Add keyed bao methods. --- src/io/fsm.rs | 305 +++++++++++- src/io/outboard.rs | 31 ++ src/io/sync.rs | 269 ++++++++++- src/lib.rs | 77 ++- src/rec.rs | 449 ++++++++++++++++- src/tests.rs | 1140 +++++++++++++++++++++++++++++++++++++++++++- src/tests2.rs | 541 ++++++++++++++++++++- 7 files changed, 2748 insertions(+), 64 deletions(-) diff --git a/src/io/fsm.rs b/src/io/fsm.rs index 5c64e01..00eba8d 100644 --- a/src/io/fsm.rs +++ b/src/io/fsm.rs @@ -22,14 +22,14 @@ use smallvec::SmallVec; pub use super::BaoContentItem; use super::{combine_hash_pair, DecodeError}; use crate::{ - blake3, hash_subtree, + blake3, hash_subtree_with_key, io::{ error::EncodeError, outboard::{PostOrderOutboard, PreOrderOutboard}, Leaf, Parent, }, iter::{BaoChunk, ResponseIter}, - parent_cv, + parent_cv_with_key, rec::{encode_selected_rec, truncate_ranges, truncate_ranges_owned}, BaoTree, BlockSize, ChunkRanges, ChunkRangesRef, TreeNode, }; @@ -129,6 +129,37 @@ pub trait CreateOutboard { /// /// It will only include data up the the current tree size. fn init_from(&mut self, data: impl AsyncStreamReader) -> impl Future>; + + /// Create a keyed outboard from a seekable data source. + #[allow(async_fn_in_trait)] + async fn create_keyed( + mut data: impl AsyncSliceReader, + block_size: BlockSize, + key: &[u8; 32], + ) -> io::Result + where + Self: Default + Sized, + { + let size = data.size().await?; + Self::create_sized_keyed(Cursor::new(data), size, block_size, key).await + } + + /// Create a keyed outboard from a data source with a known size. + fn create_sized_keyed( + data: impl AsyncStreamReader, + size: u64, + block_size: BlockSize, + key: &[u8; 32], + ) -> impl Future> + where + Self: Default + Sized; + + /// Init a keyed outboard from a data source. + fn init_from_keyed( + &mut self, + data: impl AsyncStreamReader, + key: &[u8; 32], + ) -> impl Future>; } impl Outboard for &mut O { @@ -250,6 +281,35 @@ impl CreateOutboard for PreOrderOutboard { this.sync().await?; Ok(()) } + + async fn create_sized_keyed( + data: impl AsyncStreamReader, + size: u64, + block_size: BlockSize, + key: &[u8; 32], + ) -> io::Result + where + Self: Default + Sized, + { + let mut res = Self { + tree: BaoTree::new(size, block_size), + ..Self::default() + }; + res.init_from_keyed(data, key).await?; + Ok(res) + } + + async fn init_from_keyed( + &mut self, + data: impl AsyncStreamReader, + key: &[u8; 32], + ) -> io::Result<()> { + let mut this = self; + let root = keyed_outboard(data, this.tree, &mut this, key).await?; + this.root = root; + this.sync().await?; + Ok(()) + } } impl CreateOutboard for PostOrderOutboard { @@ -276,6 +336,35 @@ impl CreateOutboard for PostOrderOutboard { this.sync().await?; Ok(()) } + + async fn create_sized_keyed( + data: impl AsyncStreamReader, + size: u64, + block_size: BlockSize, + key: &[u8; 32], + ) -> io::Result + where + Self: Default + Sized, + { + let mut res = Self { + tree: BaoTree::new(size, block_size), + ..Self::default() + }; + res.init_from_keyed(data, key).await?; + Ok(res) + } + + async fn init_from_keyed( + &mut self, + data: impl AsyncStreamReader, + key: &[u8; 32], + ) -> io::Result<()> { + let mut this = self; + let root = keyed_outboard(data, this.tree, &mut this, key).await?; + this.root = root; + this.sync().await?; + Ok(()) + } } impl Outboard for PostOrderOutboard { @@ -318,16 +407,28 @@ struct ResponseDecoderInner { iter: ResponseIter, stack: SmallVec<[blake3::Hash; 10]>, encoded: R, + key: Option<[u8; 32]>, } impl ResponseDecoderInner { fn new(tree: BaoTree, hash: blake3::Hash, ranges: ChunkRanges, encoded: R) -> Self { + Self::new_with_key(tree, hash, ranges, encoded, None) + } + + fn new_with_key( + tree: BaoTree, + hash: blake3::Hash, + ranges: ChunkRanges, + encoded: R, + key: Option<[u8; 32]>, + ) -> Self { // now that we know the size, we can canonicalize the ranges let ranges = truncate_ranges_owned(ranges, tree.size()); let mut res = Self { iter: ResponseIter::new(tree, ranges), stack: SmallVec::new(), encoded, + key, }; res.stack.push(hash); res @@ -362,6 +463,23 @@ impl ResponseDecoder { ))) } + /// Create a new keyed response decoder. + pub fn new_keyed( + hash: blake3::Hash, + ranges: ChunkRanges, + tree: BaoTree, + encoded: R, + key: &[u8; 32], + ) -> Self { + Self(Box::new(ResponseDecoderInner::new_with_key( + tree, + hash, + ranges, + encoded, + Some(*key), + ))) + } + /// Proceed to the next state by reading the next chunk from the stream. pub async fn next(mut self) -> ResponseDecoderNext { if let Some(chunk) = self.0.iter.next() { @@ -404,7 +522,7 @@ impl ResponseDecoder { .map_err(|e| DecodeError::maybe_parent_not_found(e, node))?; let pair @ (l_hash, r_hash) = read_parent(&buf); let parent_hash = this.stack.pop().unwrap(); - let actual = parent_cv(&l_hash, &r_hash, is_root); + let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, this.key.as_ref()); // Push the children in reverse order so they are popped in the correct order // only push right if the range intersects with the right child if right { @@ -434,7 +552,8 @@ impl ResponseDecoder { .await .map_err(|e| DecodeError::maybe_leaf_not_found(e, start_chunk))?; let leaf_hash = this.stack.pop().unwrap(); - let actual = hash_subtree(start_chunk.0, &data, is_root); + let actual = + hash_subtree_with_key(start_chunk.0, &data, is_root, this.key.as_ref()); if leaf_hash != actual { return Err(DecodeError::LeafHashMismatch(start_chunk)); } @@ -501,10 +620,41 @@ where /// This will either succeed if the requested ranges are all present, or fail /// as soon as a range is missing. pub async fn encode_ranges_validated( + data: D, + outboard: O, + ranges: &ChunkRangesRef, + encoded: W, +) -> result::Result<(), EncodeError> +where + D: AsyncSliceReader, + O: Outboard, + W: AsyncStreamWriter, +{ + encode_ranges_validated_with_key(data, outboard, ranges, encoded, None).await +} + +/// Encode ranges with BLAKE3 keyed hash validation. +pub async fn keyed_encode_ranges_validated( + data: D, + outboard: O, + ranges: &ChunkRangesRef, + encoded: W, + key: &[u8; 32], +) -> result::Result<(), EncodeError> +where + D: AsyncSliceReader, + O: Outboard, + W: AsyncStreamWriter, +{ + encode_ranges_validated_with_key(data, outboard, ranges, encoded, Some(key)).await +} + +async fn encode_ranges_validated_with_key( mut data: D, mut outboard: O, ranges: &ChunkRangesRef, encoded: W, + key: Option<&[u8; 32]>, ) -> result::Result<(), EncodeError> where D: AsyncSliceReader, @@ -529,7 +679,7 @@ where .. } => { let (l_hash, r_hash) = outboard.load(node).await?.unwrap(); - let actual = parent_cv(&l_hash, &r_hash, is_root); + let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, key); let expected = stack.pop().unwrap(); if actual != expected { return Err(EncodeError::ParentHashMismatch(node)); @@ -570,10 +720,11 @@ where tree.block_size.to_u32(), true, &mut out_buf, + key, ); (actual, out_buf.clone().into()) } else { - let actual = hash_subtree(start_chunk.0, &bytes, is_root); + let actual = hash_subtree_with_key(start_chunk.0, &bytes, is_root, key); (actual, bytes) }; if actual != expected { @@ -594,17 +745,53 @@ where /// If you do not want to update an outboard, use [super::outboard::EmptyOutboard] as /// the outboard. pub async fn decode_ranges( + encoded: R, + ranges: ChunkRanges, + target: W, + outboard: O, +) -> std::result::Result<(), DecodeError> +where + O: OutboardMut + Outboard, + R: AsyncStreamReader, + W: AsyncSliceWriter, +{ + decode_ranges_with_key(encoded, ranges, target, outboard, None).await +} + +/// Decode a keyed response into a file while updating an outboard. +pub async fn keyed_decode_ranges( + encoded: R, + ranges: ChunkRanges, + target: W, + outboard: O, + key: &[u8; 32], +) -> std::result::Result<(), DecodeError> +where + O: OutboardMut + Outboard, + R: AsyncStreamReader, + W: AsyncSliceWriter, +{ + decode_ranges_with_key(encoded, ranges, target, outboard, Some(key)).await +} + +async fn decode_ranges_with_key( encoded: R, ranges: ChunkRanges, mut target: W, mut outboard: O, + key: Option<&[u8; 32]>, ) -> std::result::Result<(), DecodeError> where O: OutboardMut + Outboard, R: AsyncStreamReader, W: AsyncSliceWriter, { - let mut reading = ResponseDecoder::new(outboard.root(), ranges, outboard.tree(), encoded); + let mut reading = match key { + None => ResponseDecoder::new(outboard.root(), ranges, outboard.tree(), encoded), + Some(key) => { + ResponseDecoder::new_keyed(outboard.root(), ranges, outboard.tree(), encoded, key) + } + }; loop { let item = match reading.next().await { ResponseDecoderNext::Done(_reader) => break, @@ -635,12 +822,31 @@ fn read_parent(buf: &[u8]) -> (blake3::Hash, blake3::Hash) { /// Unlike [outboard_post_order], this will work with any outboard /// implementation, but it is not guaranteed that writes are sequential. pub async fn outboard( + data: impl AsyncStreamReader, + tree: BaoTree, + outboard: impl OutboardMut, +) -> io::Result { + outboard_with_key(data, tree, outboard, None).await +} + +/// Compute the keyed outboard for the given data. +pub async fn keyed_outboard( + data: impl AsyncStreamReader, + tree: BaoTree, + outboard: impl OutboardMut, + key: &[u8; 32], +) -> io::Result { + outboard_with_key(data, tree, outboard, Some(key)).await +} + +async fn outboard_with_key( data: impl AsyncStreamReader, tree: BaoTree, mut outboard: impl OutboardMut, + key: Option<&[u8; 32]>, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - let hash = outboard_impl(tree, data, &mut outboard, &mut buffer).await?; + let hash = outboard_impl(tree, data, &mut outboard, &mut buffer, key).await?; Ok(hash) } @@ -650,6 +856,7 @@ async fn outboard_impl( mut data: impl AsyncStreamReader, mut outboard: impl OutboardMut, buffer: &mut [u8], + key: Option<&[u8; 32]>, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -660,7 +867,7 @@ async fn outboard_impl( let right_hash = stack.pop().unwrap(); let left_hash = stack.pop().unwrap(); outboard.save(node, &(left_hash, right_hash)).await?; - let parent = parent_cv(&left_hash, &right_hash, is_root); + let parent = parent_cv_with_key(&left_hash, &right_hash, is_root, key); stack.push(parent); } BaoChunk::Leaf { @@ -670,7 +877,7 @@ async fn outboard_impl( .. } => { let buf = data.read_bytes_exact(size).await?; - let hash = hash_subtree(start_chunk.0, &buf, is_root); + let hash = hash_subtree_with_key(start_chunk.0, &buf, is_root, key); stack.push(hash); } } @@ -687,12 +894,31 @@ async fn outboard_impl( /// This will not add the size to the output. You need to store it somewhere else /// or append it yourself. pub async fn outboard_post_order( + data: impl AsyncStreamReader, + tree: BaoTree, + outboard: impl AsyncStreamWriter, +) -> io::Result { + outboard_post_order_with_key(data, tree, outboard, None).await +} + +/// Compute the keyed post order outboard for the given data. +pub async fn keyed_outboard_post_order( + data: impl AsyncStreamReader, + tree: BaoTree, + outboard: impl AsyncStreamWriter, + key: &[u8; 32], +) -> io::Result { + outboard_post_order_with_key(data, tree, outboard, Some(key)).await +} + +async fn outboard_post_order_with_key( data: impl AsyncStreamReader, tree: BaoTree, mut outboard: impl AsyncStreamWriter, + key: Option<&[u8; 32]>, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - let hash = outboard_post_order_impl(tree, data, &mut outboard, &mut buffer).await?; + let hash = outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, key).await?; Ok(hash) } @@ -702,6 +928,7 @@ async fn outboard_post_order_impl( mut data: impl AsyncStreamReader, mut outboard: impl AsyncStreamWriter, buffer: &mut [u8], + key: Option<&[u8; 32]>, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -713,7 +940,7 @@ async fn outboard_post_order_impl( let left_hash = stack.pop().unwrap(); outboard.write(left_hash.as_bytes()).await?; outboard.write(right_hash.as_bytes()).await?; - let parent = parent_cv(&left_hash, &right_hash, is_root); + let parent = parent_cv_with_key(&left_hash, &right_hash, is_root, key); stack.push(parent); } BaoChunk::Leaf { @@ -723,7 +950,7 @@ async fn outboard_post_order_impl( .. } => { let buf = data.read_bytes_exact(size).await?; - let hash = hash_subtree(start_chunk.0, &buf, is_root); + let hash = hash_subtree_with_key(start_chunk.0, &buf, is_root, key); stack.push(hash); } } @@ -757,8 +984,8 @@ mod validate { use super::Outboard; use crate::{ - blake3, hash_subtree, io::LocalBoxFuture, parent_cv, rec::truncate_ranges, split, BaoTree, - ChunkNum, ChunkRangesRef, TreeNode, + blake3, hash_subtree_with_key, io::LocalBoxFuture, parent_cv_with_key, + rec::truncate_ranges, split, BaoTree, ChunkNum, ChunkRangesRef, TreeNode, }; /// Given a data file and an outboard, compute all valid ranges. @@ -771,12 +998,40 @@ mod validate { data: D, ranges: &'a ChunkRangesRef, ) -> impl Stream>> + 'a + where + O: Outboard + 'a, + D: AsyncSliceReader + 'a, + { + valid_ranges_with_key(outboard, data, ranges, None) + } + + /// Given a data file and a keyed outboard, compute all valid ranges. + pub fn keyed_valid_ranges<'a, O, D>( + outboard: O, + data: D, + ranges: &'a ChunkRangesRef, + key: &'a [u8; 32], + ) -> impl Stream>> + 'a + where + O: Outboard + 'a, + D: AsyncSliceReader + 'a, + { + valid_ranges_with_key(outboard, data, ranges, Some(key)) + } + + fn valid_ranges_with_key<'a, O, D>( + outboard: O, + data: D, + ranges: &'a ChunkRangesRef, + key: Option<&'a [u8; 32]>, + ) -> impl Stream>> + 'a where O: Outboard + 'a, D: AsyncSliceReader + 'a, { Gen::new(move |co| async move { - if let Err(cause) = RecursiveDataValidator::validate(outboard, data, ranges, &co).await + if let Err(cause) = + RecursiveDataValidator::validate(outboard, data, ranges, &co, key).await { co.yield_(Err(cause)).await; } @@ -789,6 +1044,7 @@ mod validate { outboard: O, data: D, co: &'a Co>>, + key: Option<&'a [u8; 32]>, } impl RecursiveDataValidator<'_, O, D> { @@ -797,6 +1053,7 @@ mod validate { data: D, ranges: &ChunkRangesRef, co: &Co>>, + key: Option<&[u8; 32]>, ) -> io::Result<()> { let tree = outboard.tree(); if tree.blocks() == 1 { @@ -805,7 +1062,7 @@ mod validate { let data = data .read_exact_at(0, tree.size().try_into().unwrap()) .await?; - let actual = hash_subtree(0, &data, true); + let actual = hash_subtree_with_key(0, &data, true, key); if actual == outboard.root() { co.yield_(Ok(ChunkNum(0)..tree.chunks())).await; } @@ -820,6 +1077,7 @@ mod validate { outboard, data, co, + key, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -835,7 +1093,12 @@ mod validate { let len = (range.end - range.start).try_into().unwrap(); let data = self.data.read_exact_at(range.start, len).await?; // is_root is always false because the case of a single chunk group is handled before calling this function - let actual = hash_subtree(ChunkNum::full_chunks(range.start).0, &data, is_root); + let actual = hash_subtree_with_key( + ChunkNum::full_chunks(range.start).0, + &data, + is_root, + self.key, + ); if &actual == hash { // yield the left range self.co @@ -869,7 +1132,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = parent_cv(&l_hash, &r_hash, is_root); + let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, self.key); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -972,7 +1235,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = parent_cv(&l_hash, &r_hash, is_root); + let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, None); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -998,4 +1261,4 @@ mod validate { } } #[cfg(feature = "validate")] -pub use validate::{valid_outboard_ranges, valid_ranges}; +pub use validate::{keyed_valid_ranges, valid_outboard_ranges, valid_ranges}; diff --git a/src/io/outboard.rs b/src/io/outboard.rs index 327456f..b9d12af 100644 --- a/src/io/outboard.rs +++ b/src/io/outboard.rs @@ -193,6 +193,21 @@ impl PostOrderMemOutboard { } } + /// Create a keyed outboard from `data` and a `block_size`. + pub fn create_keyed(data: impl AsRef<[u8]>, block_size: BlockSize, key: &[u8; 32]) -> Self { + let data = data.as_ref(); + let size = data.len() as u64; + let tree = BaoTree::new(size, block_size); + let mut outboard = Vec::with_capacity(tree.outboard_size().try_into().unwrap()); + let root = + crate::io::sync::keyed_outboard_post_order(data, tree, &mut outboard, key).unwrap(); + Self { + root, + tree, + data: outboard, + } + } + /// returns the outboard data, with the length suffix. pub fn into_inner_with_suffix(self) -> Vec { let mut res = self.data; @@ -366,6 +381,22 @@ impl PreOrderMemOutboard { res.root = root; res } + + /// Create a keyed outboard from `data` and a `block_size`. + pub fn create_keyed(data: impl AsRef<[u8]>, block_size: BlockSize, key: &[u8; 32]) -> Self { + let data = data.as_ref(); + let size = data.len() as u64; + let tree = BaoTree::new(size, block_size); + let outboard = vec![0u8; tree.outboard_size().try_into().unwrap()]; + let mut res = Self { + root: blake3::Hash::from([0; 32]), + tree, + data: outboard, + }; + let root = crate::io::sync::keyed_outboard(data, tree, &mut res, key).unwrap(); + res.root = root; + res + } } impl PreOrderMemOutboard { diff --git a/src/io/sync.rs b/src/io/sync.rs index ba715a5..853bcfe 100644 --- a/src/io/sync.rs +++ b/src/io/sync.rs @@ -14,14 +14,14 @@ use smallvec::SmallVec; use super::{combine_hash_pair, BaoContentItem, DecodeError}; pub use crate::rec::truncate_ranges; use crate::{ - blake3, hash_subtree, + blake3, hash_subtree_with_key, io::{ error::EncodeError, outboard::{parse_hash_pair, PostOrderOutboard, PreOrderOutboard}, Leaf, Parent, }, iter::{BaoChunk, ResponseIterRef}, - parent_cv, + parent_cv_with_key, rec::encode_selected_rec, BaoTree, BlockSize, ChunkRangesRef, TreeNode, }; @@ -97,6 +97,33 @@ pub trait CreateOutboard { /// /// It will only include data up the the current tree size. fn init_from(&mut self, data: impl Read) -> io::Result<()>; + + /// Create a keyed outboard from a data source. + fn create_keyed( + mut data: impl Read + Seek, + block_size: BlockSize, + key: &[u8; 32], + ) -> io::Result + where + Self: Default + Sized, + { + let size = data.seek(io::SeekFrom::End(0))?; + data.rewind()?; + Self::create_sized_keyed(data, size, block_size, key) + } + + /// Create a keyed outboard from a data source with a known size. + fn create_sized_keyed( + data: impl Read, + size: u64, + block_size: BlockSize, + key: &[u8; 32], + ) -> io::Result + where + Self: Default + Sized; + + /// Init a keyed outboard from a data source. + fn init_from_keyed(&mut self, data: impl Read, key: &[u8; 32]) -> io::Result<()>; } impl OutboardMut for &mut O { @@ -193,6 +220,33 @@ impl CreateOutboard for PreOrderOutboard { this.sync()?; Ok(()) } + + fn create_sized_keyed( + data: impl Read, + size: u64, + block_size: BlockSize, + key: &[u8; 32], + ) -> io::Result + where + Self: Default + Sized, + { + let tree = BaoTree::new(size, block_size); + let mut res = Self { + tree, + ..Default::default() + }; + res.init_from_keyed(data, key)?; + res.sync()?; + Ok(res) + } + + fn init_from_keyed(&mut self, data: impl Read, key: &[u8; 32]) -> io::Result<()> { + let mut this = self; + let root = keyed_outboard(data, this.tree, &mut this, key)?; + this.root = root; + this.sync()?; + Ok(()) + } } impl CreateOutboard for PostOrderOutboard { @@ -217,6 +271,33 @@ impl CreateOutboard for PostOrderOutboard { this.sync()?; Ok(()) } + + fn create_sized_keyed( + data: impl Read, + size: u64, + block_size: BlockSize, + key: &[u8; 32], + ) -> io::Result + where + Self: Default + Sized, + { + let tree = BaoTree::new(size, block_size); + let mut res = Self { + tree, + ..Default::default() + }; + res.init_from_keyed(data, key)?; + res.sync()?; + Ok(res) + } + + fn init_from_keyed(&mut self, data: impl Read, key: &[u8; 32]) -> io::Result<()> { + let mut this = self; + let root = keyed_outboard(data, this.tree, &mut this, key)?; + this.root = root; + this.sync()?; + Ok(()) + } } impl OutboardMut for PostOrderOutboard { @@ -264,6 +345,7 @@ pub struct DecodeResponseIter<'a, R> { stack: SmallVec<[blake3::Hash; 10]>, encoded: R, buf: BytesMut, + key: Option<[u8; 32]>, } impl<'a, R: Read> DecodeResponseIter<'a, R> { @@ -295,9 +377,23 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { inner: ResponseIterRef::new(tree, ranges), encoded, buf, + key: None, } } + /// Create a new iterator to decode a keyed response. + pub fn new_keyed( + root: blake3::Hash, + tree: BaoTree, + encoded: R, + ranges: &'a ChunkRangesRef, + key: &[u8; 32], + ) -> Self { + let mut res = Self::new(root, tree, encoded, ranges); + res.key = Some(*key); + res + } + /// Get a reference to the buffer used for decoding. pub fn buffer(&self) -> &[u8] { &self.buf @@ -322,7 +418,7 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { let pair @ (l_hash, r_hash) = read_parent(&mut self.encoded) .map_err(|e| DecodeError::maybe_parent_not_found(e, node))?; let parent_hash = self.stack.pop().unwrap(); - let actual = parent_cv(&l_hash, &r_hash, is_root); + let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, self.key.as_ref()); if parent_hash != actual { return Err(DecodeError::ParentHashMismatch(node)); } @@ -344,7 +440,8 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { self.encoded .read_exact(&mut self.buf) .map_err(|e| DecodeError::maybe_leaf_not_found(e, start_chunk))?; - let actual = hash_subtree(start_chunk.0, &self.buf, is_root); + let actual = + hash_subtree_with_key(start_chunk.0, &self.buf, is_root, self.key.as_ref()); let leaf_hash = self.stack.pop().unwrap(); if leaf_hash != actual { return Err(DecodeError::LeafHashMismatch(start_chunk)); @@ -419,6 +516,27 @@ pub fn encode_ranges_validated( outboard: O, ranges: &ChunkRangesRef, encoded: W, +) -> result::Result<(), EncodeError> { + encode_ranges_validated_with_key(data, outboard, ranges, encoded, None) +} + +/// Encode ranges with BLAKE3 keyed hash validation. +pub fn keyed_encode_ranges_validated( + data: D, + outboard: O, + ranges: &ChunkRangesRef, + encoded: W, + key: &[u8; 32], +) -> result::Result<(), EncodeError> { + encode_ranges_validated_with_key(data, outboard, ranges, encoded, Some(key)) +} + +fn encode_ranges_validated_with_key( + data: D, + outboard: O, + ranges: &ChunkRangesRef, + encoded: W, + key: Option<&[u8; 32]>, ) -> result::Result<(), EncodeError> { if ranges.is_empty() { return Ok(()); @@ -442,7 +560,7 @@ pub fn encode_ranges_validated( .. } => { let (l_hash, r_hash) = outboard.load(node)?.unwrap(); - let actual = parent_cv(&l_hash, &r_hash, is_root); + let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, key); let expected = stack.pop().unwrap(); if actual != expected { return Err(EncodeError::ParentHashMismatch(node)); @@ -481,10 +599,11 @@ pub fn encode_ranges_validated( tree.block_size.to_u32(), true, &mut out_buf, + key, ); (actual, &out_buf[..]) } else { - let actual = hash_subtree(start_chunk.0, buf, is_root); + let actual = hash_subtree_with_key(start_chunk.0, buf, is_root, key); #[allow(clippy::redundant_slicing)] (actual, &buf[..]) }; @@ -503,17 +622,53 @@ pub fn encode_ranges_validated( /// If you do not want to update an outboard, use [super::outboard::EmptyOutboard] as /// the outboard. pub fn decode_ranges( + encoded: R, + ranges: &ChunkRangesRef, + target: W, + outboard: O, +) -> std::result::Result<(), DecodeError> +where + O: OutboardMut + Outboard, + R: Read, + W: WriteAt, +{ + decode_ranges_with_key(encoded, ranges, target, outboard, None) +} + +/// Decode a keyed response into a file while updating an outboard. +pub fn keyed_decode_ranges( + encoded: R, + ranges: &ChunkRangesRef, + target: W, + outboard: O, + key: &[u8; 32], +) -> std::result::Result<(), DecodeError> +where + O: OutboardMut + Outboard, + R: Read, + W: WriteAt, +{ + decode_ranges_with_key(encoded, ranges, target, outboard, Some(key)) +} + +fn decode_ranges_with_key( encoded: R, ranges: &ChunkRangesRef, mut target: W, mut outboard: O, + key: Option<&[u8; 32]>, ) -> std::result::Result<(), DecodeError> where O: OutboardMut + Outboard, R: Read, W: WriteAt, { - let iter = DecodeResponseIter::new(outboard.root(), outboard.tree(), encoded, ranges); + let iter = match key { + None => DecodeResponseIter::new(outboard.root(), outboard.tree(), encoded, ranges), + Some(key) => { + DecodeResponseIter::new_keyed(outboard.root(), outboard.tree(), encoded, ranges, key) + } + }; for item in iter { match item? { BaoContentItem::Parent(Parent { node, pair }) => { @@ -532,12 +687,31 @@ where /// Unlike [outboard_post_order], this will work with any outboard /// implementation, but it is not guaranteed that writes are sequential. pub fn outboard( + data: impl Read, + tree: BaoTree, + outboard: impl OutboardMut, +) -> io::Result { + outboard_with_key(data, tree, outboard, None) +} + +/// Compute the keyed outboard for the given data. +pub fn keyed_outboard( + data: impl Read, + tree: BaoTree, + outboard: impl OutboardMut, + key: &[u8; 32], +) -> io::Result { + outboard_with_key(data, tree, outboard, Some(key)) +} + +fn outboard_with_key( data: impl Read, tree: BaoTree, mut outboard: impl OutboardMut, + key: Option<&[u8; 32]>, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - let hash = outboard_impl(tree, data, &mut outboard, &mut buffer)?; + let hash = outboard_impl(tree, data, &mut outboard, &mut buffer, key)?; Ok(hash) } @@ -547,6 +721,7 @@ fn outboard_impl( mut data: impl Read, mut outboard: impl OutboardMut, buffer: &mut [u8], + key: Option<&[u8; 32]>, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -557,7 +732,7 @@ fn outboard_impl( let right_hash = stack.pop().unwrap(); let left_hash = stack.pop().unwrap(); outboard.save(node, &(left_hash, right_hash))?; - let parent = parent_cv(&left_hash, &right_hash, is_root); + let parent = parent_cv_with_key(&left_hash, &right_hash, is_root, key); stack.push(parent); } BaoChunk::Leaf { @@ -568,7 +743,7 @@ fn outboard_impl( } => { let buf = &mut buffer[..size]; data.read_exact(buf)?; - let hash = hash_subtree(start_chunk.0, buf, is_root); + let hash = hash_subtree_with_key(start_chunk.0, buf, is_root, key); stack.push(hash); } } @@ -585,12 +760,31 @@ fn outboard_impl( /// This will not add the size to the output. You need to store it somewhere else /// or append it yourself. pub fn outboard_post_order( + data: impl Read, + tree: BaoTree, + outboard: impl Write, +) -> io::Result { + outboard_post_order_with_key(data, tree, outboard, None) +} + +/// Compute the keyed post order outboard for the given data. +pub fn keyed_outboard_post_order( + data: impl Read, + tree: BaoTree, + outboard: impl Write, + key: &[u8; 32], +) -> io::Result { + outboard_post_order_with_key(data, tree, outboard, Some(key)) +} + +fn outboard_post_order_with_key( data: impl Read, tree: BaoTree, mut outboard: impl Write, + key: Option<&[u8; 32]>, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - let hash = outboard_post_order_impl(tree, data, &mut outboard, &mut buffer)?; + let hash = outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, key)?; Ok(hash) } @@ -600,6 +794,7 @@ fn outboard_post_order_impl( mut data: impl Read, mut outboard: impl Write, buffer: &mut [u8], + key: Option<&[u8; 32]>, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -611,7 +806,7 @@ fn outboard_post_order_impl( let left_hash = stack.pop().unwrap(); outboard.write_all(left_hash.as_bytes())?; outboard.write_all(right_hash.as_bytes())?; - let parent = parent_cv(&left_hash, &right_hash, is_root); + let parent = parent_cv_with_key(&left_hash, &right_hash, is_root, key); stack.push(parent); } BaoChunk::Leaf { @@ -622,7 +817,7 @@ fn outboard_post_order_impl( } => { let buf = &mut buffer[..size]; data.read_exact(buf)?; - let hash = hash_subtree(start_chunk.0, buf, is_root); + let hash = hash_subtree_with_key(start_chunk.0, buf, is_root, key); stack.push(hash); } } @@ -663,8 +858,8 @@ mod validate { use super::Outboard; use crate::{ - blake3, hash_subtree, io::LocalBoxFuture, parent_cv, rec::truncate_ranges, split, BaoTree, - ChunkNum, ChunkRangesRef, TreeNode, + blake3, hash_subtree_with_key, io::LocalBoxFuture, parent_cv_with_key, + rec::truncate_ranges, split, BaoTree, ChunkNum, ChunkRangesRef, TreeNode, }; /// Given a data file and an outboard, compute all valid ranges. @@ -677,12 +872,40 @@ mod validate { data: D, ranges: &'a ChunkRangesRef, ) -> impl IntoIterator>> + 'a + where + O: Outboard + 'a, + D: ReadAt + 'a, + { + valid_ranges_with_key(outboard, data, ranges, None) + } + + /// Given a data file and a keyed outboard, compute all valid ranges. + pub fn keyed_valid_ranges<'a, O, D>( + outboard: O, + data: D, + ranges: &'a ChunkRangesRef, + key: &'a [u8; 32], + ) -> impl IntoIterator>> + 'a + where + O: Outboard + 'a, + D: ReadAt + 'a, + { + valid_ranges_with_key(outboard, data, ranges, Some(key)) + } + + fn valid_ranges_with_key<'a, O, D>( + outboard: O, + data: D, + ranges: &'a ChunkRangesRef, + key: Option<&'a [u8; 32]>, + ) -> impl IntoIterator>> + 'a where O: Outboard + 'a, D: ReadAt + 'a, { Gen::new(move |co| async move { - if let Err(cause) = RecursiveDataValidator::validate(outboard, data, ranges, &co).await + if let Err(cause) = + RecursiveDataValidator::validate(outboard, data, ranges, &co, key).await { co.yield_(Err(cause)).await; } @@ -696,6 +919,7 @@ mod validate { data: D, buffer: Vec, co: &'a Co>>, + key: Option<&'a [u8; 32]>, } impl RecursiveDataValidator<'_, O, D> { @@ -704,6 +928,7 @@ mod validate { data: D, ranges: &ChunkRangesRef, co: &Co>>, + key: Option<&[u8; 32]>, ) -> io::Result<()> { let tree = outboard.tree(); let mut buffer = vec![0u8; tree.chunk_group_bytes()]; @@ -711,7 +936,7 @@ mod validate { // special case for a tree that fits in one block / chunk group let tmp = &mut buffer[..tree.size().try_into().unwrap()]; data.read_exact_at(0, tmp)?; - let actual = hash_subtree(0, tmp, true); + let actual = hash_subtree_with_key(0, tmp, true, key); if actual == outboard.root() { co.yield_(Ok(ChunkNum(0)..tree.chunks())).await; } @@ -727,6 +952,7 @@ mod validate { data, buffer, co, + key, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -743,7 +969,8 @@ mod validate { let tmp = &mut self.buffer[..len]; self.data.read_exact_at(range.start, tmp)?; // is_root is always false because the case of a single chunk group is handled before calling this function - let actual = hash_subtree(ChunkNum::full_chunks(range.start).0, tmp, is_root); + let actual = + hash_subtree_with_key(ChunkNum::full_chunks(range.start).0, tmp, is_root, self.key); if &actual == hash { // yield the left range self.co @@ -777,7 +1004,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = parent_cv(&l_hash, &r_hash, is_root); + let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, self.key); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -879,7 +1106,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = parent_cv(&l_hash, &r_hash, is_root); + let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, None); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -905,4 +1132,4 @@ mod validate { } } #[cfg(feature = "validate")] -pub use validate::{valid_outboard_ranges, valid_ranges}; +pub use validate::{keyed_valid_ranges, valid_outboard_ranges, valid_ranges}; diff --git a/src/lib.rs b/src/lib.rs index b0d8ec1..25592b3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -193,6 +193,13 @@ //! # } //! ``` //! +//! # Keyed hashing +//! +//! For domain-separated trees, use the `keyed_*` functions. They mirror the +//! standard API with an additional `key: &[u8; 32]` argument, like +//! [`blake3::keyed_hash`] mirrors [`blake3::hash`]. The key is out-of-band +//! metadata and is not included in the encoded stream. +//! //! # Compatibility with the [bao crate](https://crates.io/crates/bao) //! //! This crate will be compatible with the bao crate, provided you do the @@ -232,13 +239,40 @@ pub type ByteRanges = range_collections::RangeSet2; /// [ChunkRanges] implements [`AsRef`]. pub type ChunkRangesRef = range_collections::RangeSetRef; -fn hash_subtree(start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash { +pub(crate) fn hash_subtree(start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash { + hash_subtree_with_key(start_chunk, data, is_root, None) +} + +/// Compute the hash of a subtree using BLAKE3 keyed mode. +/// +/// See [keyed_parent_cv] for merging child hashes in keyed mode. +pub fn keyed_hash_subtree( + start_chunk: u64, + data: &[u8], + is_root: bool, + key: &[u8; 32], +) -> blake3::Hash { + hash_subtree_with_key(start_chunk, data, is_root, Some(key)) +} + +pub(crate) fn hash_subtree_with_key( + start_chunk: u64, + data: &[u8], + is_root: bool, + key: Option<&[u8; 32]>, +) -> blake3::Hash { use blake3::hazmat::{ChainingValue, HasherExt}; if is_root { debug_assert!(start_chunk == 0); - blake3::hash(data) + match key { + None => blake3::hash(data), + Some(key) => blake3::keyed_hash(key, data), + } } else { - let mut hasher = blake3::Hasher::new(); + let mut hasher = match key { + None => blake3::Hasher::new(), + Some(key) => blake3::Hasher::new_keyed(key), + }; hasher.set_input_offset(start_chunk * 1024); hasher.update(data); let non_root_hash: ChainingValue = hasher.finalize_non_root(); @@ -246,18 +280,41 @@ fn hash_subtree(start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash { } } -fn parent_cv(left_child: &blake3::Hash, right_child: &blake3::Hash, is_root: bool) -> blake3::Hash { +pub(crate) fn parent_cv( + left_child: &blake3::Hash, + right_child: &blake3::Hash, + is_root: bool, +) -> blake3::Hash { + parent_cv_with_key(left_child, right_child, is_root, None) +} + +/// Merge two child subtree hashes using BLAKE3 keyed mode. +pub fn keyed_parent_cv( + left_child: &blake3::Hash, + right_child: &blake3::Hash, + is_root: bool, + key: &[u8; 32], +) -> blake3::Hash { + parent_cv_with_key(left_child, right_child, is_root, Some(key)) +} + +pub(crate) fn parent_cv_with_key( + left_child: &blake3::Hash, + right_child: &blake3::Hash, + is_root: bool, + key: Option<&[u8; 32]>, +) -> blake3::Hash { use blake3::hazmat::{merge_subtrees_non_root, merge_subtrees_root, ChainingValue, Mode}; let left_child: ChainingValue = *left_child.as_bytes(); let right_child: ChainingValue = *right_child.as_bytes(); + let mode = match key { + None => Mode::Hash, + Some(key) => Mode::KeyedHash(key), + }; if is_root { - merge_subtrees_root(&left_child, &right_child, Mode::Hash) + merge_subtrees_root(&left_child, &right_child, mode) } else { - blake3::Hash::from(merge_subtrees_non_root( - &left_child, - &right_child, - Mode::Hash, - )) + blake3::Hash::from(merge_subtrees_non_root(&left_child, &right_child, mode)) } } diff --git a/src/rec.rs b/src/rec.rs index 4673991..65830c1 100644 --- a/src/rec.rs +++ b/src/rec.rs @@ -2,7 +2,7 @@ //! //! Encocding is used to compute hashes, decoding is only used in tests as a //! reference implementation. -use crate::{blake3, hash_subtree, parent_cv, split_inner, ChunkNum, ChunkRangesRef}; +use crate::{blake3, split_inner, ChunkNum, ChunkRangesRef}; /// Given a set of chunk ranges, adapt them for a tree of the given size. /// @@ -96,6 +96,7 @@ fn truncated_len(ranges: &ChunkRangesRef, size: u64) -> usize { /// This is used as a reference implementation in tests, but also to compute hashes /// below the chunk group size when creating responses for outboards with a chunk group /// size of >0. +#[allow(clippy::too_many_arguments)] // keyed mode adds `key`; splitting into a struct isn't worth it here pub(crate) fn encode_selected_rec( start_chunk: ChunkNum, data: &[u8], @@ -104,13 +105,14 @@ pub(crate) fn encode_selected_rec( min_level: u32, emit_data: bool, res: &mut Vec, + key: Option<&[u8; 32]>, ) -> blake3::Hash { use blake3::CHUNK_LEN; if data.len() <= CHUNK_LEN { if emit_data && !query.is_empty() { res.extend_from_slice(data); } - hash_subtree(start_chunk.0, data, is_root) + crate::hash_subtree_with_key(start_chunk.0, data, is_root, key) } else { let chunks = data.len() / CHUNK_LEN + (data.len() % CHUNK_LEN != 0) as usize; let chunks = chunks.next_power_of_two(); @@ -142,6 +144,7 @@ pub(crate) fn encode_selected_rec( min_level, emit_data, res, + key, ); let right = encode_selected_rec( mid_chunk, @@ -151,13 +154,14 @@ pub(crate) fn encode_selected_rec( min_level, emit_data, res, + key, ); // backfill the hashes if needed if let Some(o) = hash_offset { res[o..o + 32].copy_from_slice(left.as_bytes()); res[o + 32..o + 64].copy_from_slice(right.as_bytes()); } - parent_cv(&left, &right, is_root) + crate::parent_cv_with_key(&left, &right, is_root, key) } } @@ -275,6 +279,7 @@ mod test_support { 0, false, &mut res, + None, ); (res, hash) } @@ -290,6 +295,7 @@ mod test_support { 0, true, &mut res, + None, ); (res, hash) } @@ -430,10 +436,447 @@ mod test_support { block_size.to_u32(), true, &mut res, + None, ); (res, hash) } + use std::io::Cursor; + + use crate::io::outboard::{ + PostOrderMemOutboard, PostOrderOutboard, PreOrderMemOutboard, PreOrderOutboard, + }; + use crate::io::sync::{self, CreateOutboard, Outboard}; + + pub(crate) fn assert_post_order_outboard_matches_mem( + outboard: &PostOrderOutboard>, + data: &[u8], + block_size: BlockSize, + key: &[u8; 32], + ) { + let reference = PostOrderMemOutboard::create_keyed(data, block_size, key); + assert_eq!(outboard.root, reference.root); + let tree = outboard.tree; + let mut copied = PostOrderMemOutboard { + root: outboard.root, + tree, + data: vec![0; tree.outboard_hash_pairs() as usize * 64], + }; + sync::copy(outboard, &mut copied).unwrap(); + assert_eq!(copied.data, reference.data); + } + + pub(crate) fn assert_pre_order_outboard_matches_mem( + outboard: &PreOrderOutboard>, + data: &[u8], + block_size: BlockSize, + key: &[u8; 32], + ) { + let reference = PreOrderMemOutboard::create_keyed(data, block_size, key); + assert_eq!(outboard.root, reference.root); + let tree = outboard.tree; + let mut copied = PreOrderMemOutboard { + root: outboard.root, + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + sync::copy(outboard, &mut copied).unwrap(); + assert_eq!(copied.data, reference.data); + } + + fn assert_truncated_create_sized_keyed_post( + truncated: &PostOrderOutboard>, + data: &[u8], + truncated_size: u64, + block_size: BlockSize, + key: &[u8; 32], + ) { + assert_eq!(truncated.tree.size, truncated_size); + assert_eq!( + truncated.root(), + blake3::keyed_hash(key, &data[..truncated_size as usize]) + ); + assert_post_order_outboard_matches_mem( + truncated, + &data[..truncated_size as usize], + block_size, + key, + ); + } + + fn assert_truncated_create_sized_keyed_pre( + truncated: &PreOrderOutboard>, + data: &[u8], + truncated_size: u64, + block_size: BlockSize, + key: &[u8; 32], + ) { + assert_eq!(truncated.tree.size, truncated_size); + assert_eq!( + truncated.root(), + blake3::keyed_hash(key, &data[..truncated_size as usize]) + ); + assert_pre_order_outboard_matches_mem( + truncated, + &data[..truncated_size as usize], + block_size, + key, + ); + } + + pub(crate) fn keyed_create_sized_keyed_checks( + data: &[u8], + block_size: BlockSize, + key: &[u8; 32], + ) { + let size = data.len() as u64; + + let post: PostOrderOutboard> = + PostOrderOutboard::create_sized_keyed(Cursor::new(data), size, block_size, key) + .unwrap(); + assert_post_order_outboard_matches_mem(&post, data, block_size, key); + + let pre: PreOrderOutboard> = + PreOrderOutboard::create_sized_keyed(Cursor::new(data), size, block_size, key).unwrap(); + assert_pre_order_outboard_matches_mem(&pre, data, block_size, key); + + let truncated_size = 1024u64.min(size); + if truncated_size < size { + let truncated_post: PostOrderOutboard> = PostOrderOutboard::create_sized_keyed( + Cursor::new(data), + truncated_size, + BlockSize(0), + key, + ) + .unwrap(); + assert_truncated_create_sized_keyed_post( + &truncated_post, + data, + truncated_size, + BlockSize(0), + key, + ); + + let truncated_pre: PreOrderOutboard> = PreOrderOutboard::create_sized_keyed( + Cursor::new(data), + truncated_size, + BlockSize(0), + key, + ) + .unwrap(); + assert_truncated_create_sized_keyed_pre( + &truncated_pre, + data, + truncated_size, + BlockSize(0), + key, + ); + } + } + + pub(crate) fn keyed_init_from_keyed_checks(data: &[u8], block_size: BlockSize, key: &[u8; 32]) { + let tree = BaoTree::new(data.len() as u64, block_size); + let expected = blake3::keyed_hash(key, data); + + let mut post = PostOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + post.init_from_keyed(Cursor::new(data), key).unwrap(); + assert_eq!(post.root(), expected); + assert_post_order_outboard_matches_mem(&post, data, block_size, key); + + let mut pre = PreOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + pre.init_from_keyed(Cursor::new(data), key).unwrap(); + assert_eq!(pre.root(), expected); + assert_pre_order_outboard_matches_mem(&pre, data, block_size, key); + + let truncated_size = 1024u64.min(data.len() as u64); + if truncated_size < data.len() as u64 { + let truncated_tree = BaoTree::new(truncated_size, BlockSize(0)); + let truncated_expected = blake3::keyed_hash(key, &data[..truncated_size as usize]); + + let mut truncated_post = PostOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree: truncated_tree, + data: vec![0; truncated_tree.outboard_size().try_into().unwrap()], + }; + truncated_post + .init_from_keyed(Cursor::new(data), key) + .unwrap(); + assert_eq!(truncated_post.root(), truncated_expected); + assert_post_order_outboard_matches_mem( + &truncated_post, + &data[..truncated_size as usize], + BlockSize(0), + key, + ); + + let mut truncated_pre = PreOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree: truncated_tree, + data: vec![0; truncated_tree.outboard_size().try_into().unwrap()], + }; + truncated_pre + .init_from_keyed(Cursor::new(data), key) + .unwrap(); + assert_eq!(truncated_pre.root(), truncated_expected); + assert_pre_order_outboard_matches_mem( + &truncated_pre, + &data[..truncated_size as usize], + BlockSize(0), + key, + ); + } + } + + pub(crate) fn keyed_outboard_functions_checks( + data: &[u8], + block_size: BlockSize, + key: &[u8; 32], + ) { + let tree = BaoTree::new(data.len() as u64, block_size); + let expected = blake3::keyed_hash(key, data); + + let mut pre = PreOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + let root = sync::keyed_outboard(Cursor::new(data), tree, &mut pre, key).unwrap(); + pre.root = root; + assert_eq!(root, expected); + assert_pre_order_outboard_matches_mem(&pre, data, block_size, key); + + let mut post_buf = Vec::new(); + let root = + sync::keyed_outboard_post_order(Cursor::new(data), tree, &mut post_buf, key).unwrap(); + assert_eq!(root, expected); + assert_eq!(post_buf.len(), tree.outboard_size().try_into().unwrap()); + + let reference_post = PostOrderMemOutboard::create_keyed(data, block_size, key); + assert_eq!(post_buf, reference_post.data); + + let post_mem = PostOrderMemOutboard { + root, + tree, + data: post_buf, + }; + let pre_from_post = post_mem.flip(); + assert_eq!(pre_from_post.data, pre.data); + } + + #[cfg(feature = "tokio_fsm")] + pub(crate) async fn keyed_create_sized_keyed_checks_fsm( + data: &[u8], + block_size: BlockSize, + key: &[u8; 32], + ) { + use bytes::Bytes; + + let size = data.len() as u64; + + let post: PostOrderOutboard> = + > as crate::io::fsm::CreateOutboard>::create_sized_keyed( + Cursor::new(Bytes::from(data.to_vec())), + size, + block_size, + key, + ) + .await + .unwrap(); + assert_post_order_outboard_matches_mem(&post, data, block_size, key); + + let pre: PreOrderOutboard> = + > as crate::io::fsm::CreateOutboard>::create_sized_keyed( + Cursor::new(Bytes::from(data.to_vec())), + size, + block_size, + key, + ) + .await + .unwrap(); + assert_pre_order_outboard_matches_mem(&pre, data, block_size, key); + + let truncated_size = 1024u64.min(size); + if truncated_size < size { + let truncated_post: PostOrderOutboard> = + > as crate::io::fsm::CreateOutboard>::create_sized_keyed( + Cursor::new(Bytes::from(data.to_vec())), + truncated_size, + BlockSize(0), + key, + ) + .await + .unwrap(); + assert_truncated_create_sized_keyed_post( + &truncated_post, + data, + truncated_size, + BlockSize(0), + key, + ); + + let truncated_pre: PreOrderOutboard> = + > as crate::io::fsm::CreateOutboard>::create_sized_keyed( + Cursor::new(Bytes::from(data.to_vec())), + truncated_size, + BlockSize(0), + key, + ) + .await + .unwrap(); + assert_truncated_create_sized_keyed_pre( + &truncated_pre, + data, + truncated_size, + BlockSize(0), + key, + ); + } + } + + #[cfg(feature = "tokio_fsm")] + pub(crate) async fn keyed_outboard_functions_checks_fsm( + data: &[u8], + block_size: BlockSize, + key: &[u8; 32], + ) { + use crate::io::fsm::{keyed_outboard, keyed_outboard_post_order}; + use bytes::Bytes; + + let tree = BaoTree::new(data.len() as u64, block_size); + let expected = blake3::keyed_hash(key, data); + + let mut pre = PreOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + let root = keyed_outboard(Cursor::new(Bytes::from(data.to_vec())), tree, &mut pre, key) + .await + .unwrap(); + pre.root = root; + assert_eq!(root, expected); + assert_pre_order_outboard_matches_mem(&pre, data, block_size, key); + + let mut post_buf = Vec::new(); + let root = keyed_outboard_post_order( + Cursor::new(Bytes::from(data.to_vec())), + tree, + &mut post_buf, + key, + ) + .await + .unwrap(); + assert_eq!(root, expected); + assert_eq!(post_buf.len(), tree.outboard_size().try_into().unwrap()); + + let reference_post = PostOrderMemOutboard::create_keyed(data, block_size, key); + assert_eq!(post_buf, reference_post.data); + + let post_mem = PostOrderMemOutboard { + root, + tree, + data: post_buf, + }; + let pre_from_post = post_mem.flip(); + assert_eq!(pre_from_post.data, pre.data); + } + + #[cfg(feature = "tokio_fsm")] + pub(crate) async fn keyed_init_from_keyed_checks_fsm( + data: &[u8], + block_size: BlockSize, + key: &[u8; 32], + ) { + use bytes::Bytes; + + let tree = BaoTree::new(data.len() as u64, block_size); + let expected = blake3::keyed_hash(key, data); + + let mut post = PostOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + crate::io::fsm::CreateOutboard::init_from_keyed( + &mut post, + Cursor::new(Bytes::from(data.to_vec())), + key, + ) + .await + .unwrap(); + assert_eq!(post.root(), expected); + assert_post_order_outboard_matches_mem(&post, data, block_size, key); + + let mut pre = PreOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + crate::io::fsm::CreateOutboard::init_from_keyed( + &mut pre, + Cursor::new(Bytes::from(data.to_vec())), + key, + ) + .await + .unwrap(); + assert_eq!(pre.root(), expected); + assert_pre_order_outboard_matches_mem(&pre, data, block_size, key); + + let truncated_size = 1024u64.min(data.len() as u64); + if truncated_size < data.len() as u64 { + let truncated_tree = BaoTree::new(truncated_size, BlockSize(0)); + let truncated_expected = blake3::keyed_hash(key, &data[..truncated_size as usize]); + + let mut truncated_post = PostOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree: truncated_tree, + data: vec![0; truncated_tree.outboard_size().try_into().unwrap()], + }; + crate::io::fsm::CreateOutboard::init_from_keyed( + &mut truncated_post, + Cursor::new(Bytes::from(data.to_vec())), + key, + ) + .await + .unwrap(); + assert_eq!(truncated_post.root(), truncated_expected); + assert_post_order_outboard_matches_mem( + &truncated_post, + &data[..truncated_size as usize], + BlockSize(0), + key, + ); + + let mut truncated_pre = PreOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree: truncated_tree, + data: vec![0; truncated_tree.outboard_size().try_into().unwrap()], + }; + crate::io::fsm::CreateOutboard::init_from_keyed( + &mut truncated_pre, + Cursor::new(Bytes::from(data.to_vec())), + key, + ) + .await + .unwrap(); + assert_eq!(truncated_pre.root(), truncated_expected); + assert_pre_order_outboard_matches_mem( + &truncated_pre, + &data[..truncated_size as usize], + BlockSize(0), + key, + ); + } + } + /// Check that l and r of a 2-tuple are equal #[macro_export] macro_rules! assert_tuple_eq { diff --git a/src/tests.rs b/src/tests.rs index 4959f69..aba318d 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -19,17 +19,432 @@ use super::{ BaoTree, BlockSize, TreeNode, }; use crate::{ - assert_tuple_eq, blake3, - io::{full_chunk_groups, outboard::PreOrderMemOutboard, sync::Outboard, BaoContentItem, Leaf}, + assert_tuple_eq, blake3, hash_subtree, + io::{ + full_chunk_groups, + outboard::{PostOrderOutboard, PreOrderMemOutboard, PreOrderOutboard}, + sync::Outboard, + BaoContentItem, DecodeError, EncodeError, Leaf, + }, iter::{PostOrderChunkIter, PreOrderPartialIterRef, ResponseIterRef}, - prop_assert_tuple_eq, + keyed_hash_subtree, keyed_parent_cv, parent_cv, prop_assert_tuple_eq, rec::{ - encode_ranges_reference, encode_selected_rec, make_test_data, range_union, truncate_ranges, - ReferencePreOrderPartialChunkIterRef, + encode_ranges_reference, encode_selected_rec, keyed_create_sized_keyed_checks, + keyed_init_from_keyed_checks, keyed_outboard_functions_checks, make_test_data, range_union, + truncate_ranges, ReferencePreOrderPartialChunkIterRef, }, split, ChunkRanges, ChunkRangesRef, ResponseIter, }; +#[cfg(feature = "tokio_fsm")] +use crate::rec::{ + keyed_create_sized_keyed_checks_fsm, keyed_init_from_keyed_checks_fsm, + keyed_outboard_functions_checks_fsm, +}; + +fn keyed_encode_selected_reference( + data: &[u8], + block_size: BlockSize, + ranges: &ChunkRangesRef, + key: &[u8; 32], +) -> (blake3::Hash, Vec) { + let mut res = Vec::new(); + let max_skip_level = block_size.to_u32(); + let ranges = truncate_ranges(ranges, data.len() as u64); + let hash = encode_selected_rec( + ChunkNum(0), + data, + true, + ranges, + max_skip_level, + true, + &mut res, + Some(key), + ); + (hash, res) +} + +fn keyed_encode_decode_roundtrip_sync_impl(data: &[u8], block_size: BlockSize, key: &[u8; 32]) { + use crate::io::sync::{keyed_decode_ranges, keyed_encode_ranges_validated}; + + let outboard = PostOrderMemOutboard::create_keyed(data, block_size, key); + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + keyed_encode_ranges_validated(data, &outboard, &ranges, &mut encoded, key).unwrap(); + let size = outboard.tree.size; + let tree = BaoTree::new(size, block_size); + let mut decoded = Vec::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + keyed_decode_ranges( + Cursor::new(encoded), + &ranges, + &mut decoded, + &mut ob_res, + key, + ) + .unwrap(); + assert_eq!(decoded, data); + assert_eq!(ob_res.root(), outboard.root()); +} + +fn keyed_encode_decode_roundtrip_fsm_impl(data: Vec, block_size: BlockSize, key: &[u8; 32]) { + use crate::io::fsm::{keyed_decode_ranges, keyed_encode_ranges_validated}; + + let mut outboard = PostOrderMemOutboard::create_keyed(&data, block_size, key); + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + tokio::runtime::Runtime::new() + .unwrap() + .block_on(keyed_encode_ranges_validated( + Bytes::from(data.clone()), + &mut outboard, + &ranges, + &mut encoded, + key, + )) + .unwrap(); + let tree = outboard.tree(); + let mut decoded = bytes::BytesMut::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + tokio::runtime::Runtime::new() + .unwrap() + .block_on(keyed_decode_ranges( + Cursor::new(encoded.as_slice()), + ranges, + &mut decoded, + &mut ob_res, + key, + )) + .unwrap(); + assert_eq!(decoded.to_vec(), data); + assert_eq!(ob_res.root(), outboard.root()); +} + +/// Parent hash mismatch node for 10_000-byte payloads at block level 0. +fn keyed_multi_chunk_mismatch_node() -> TreeNode { + TreeNode(7) +} + +fn keyed_wrong_key_decode_sync_impl( + data: &[u8], + block_size: BlockSize, + expected_err: Option, +) { + use crate::io::sync::{keyed_decode_ranges, keyed_encode_ranges_validated}; + + let key_a = blake3::derive_key("bao-tree.test", b"key-a"); + let key_b = blake3::derive_key("bao-tree.test", b"key-b"); + let outboard = PostOrderMemOutboard::create_keyed(data, block_size, &key_a); + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + keyed_encode_ranges_validated(data, &outboard, &ranges, &mut encoded, &key_a).unwrap(); + let tree = outboard.tree(); + let mut decoded = Vec::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + let err = keyed_decode_ranges( + Cursor::new(encoded), + &ranges, + &mut decoded, + &mut ob_res, + &key_b, + ) + .unwrap_err(); + assert!(decoded.is_empty()); + match expected_err { + Some(expected) => assert_decode_error_eq(err, expected), + None => assert!(matches!( + err, + DecodeError::ParentHashMismatch(_) | DecodeError::LeafHashMismatch(_) + )), + } +} + +fn assert_decode_error_eq(got: DecodeError, expected: DecodeError) { + match (got, expected) { + (DecodeError::ParentHashMismatch(got), DecodeError::ParentHashMismatch(expected)) => { + assert_eq!(got, expected); + } + (DecodeError::LeafHashMismatch(got), DecodeError::LeafHashMismatch(expected)) => { + assert_eq!(got, expected); + } + (got, expected) => panic!("expected {expected:?}, got {got:?}"), + } +} + +fn assert_encode_error_eq(got: EncodeError, expected: EncodeError) { + match (got, expected) { + (EncodeError::ParentHashMismatch(got), EncodeError::ParentHashMismatch(expected)) => { + assert_eq!(got, expected); + } + (EncodeError::LeafHashMismatch(got), EncodeError::LeafHashMismatch(expected)) => { + assert_eq!(got, expected); + } + (got, expected) => panic!("expected {expected:?}, got {got:?}"), + } +} + +fn keyed_wrong_key_fails_encode_sync_impl( + data: &[u8], + block_size: BlockSize, + expected_err: EncodeError, +) { + use crate::io::sync::keyed_encode_ranges_validated; + + let key_a = blake3::derive_key("bao-tree.test", b"key-a"); + let key_b = blake3::derive_key("bao-tree.test", b"key-b"); + let outboard = PostOrderMemOutboard::create_keyed(data, block_size, &key_a); + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + let err = + keyed_encode_ranges_validated(data, &outboard, &ranges, &mut encoded, &key_b).unwrap_err(); + assert!(encoded.is_empty()); + assert_encode_error_eq(err, expected_err); +} + +fn unkeyed_encode_keyed_decode_fails_sync_impl( + data: &[u8], + block_size: BlockSize, + expected_err: DecodeError, +) { + use crate::io::sync::{encode_ranges_validated, keyed_decode_ranges}; + + let key = blake3::derive_key("bao-tree.test", b"keyed-decode"); + let outboard = PostOrderMemOutboard::create(data, block_size); + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + encode_ranges_validated(data, &outboard, &ranges, &mut encoded).unwrap(); + let tree = outboard.tree(); + let mut decoded = Vec::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + let err = keyed_decode_ranges( + Cursor::new(encoded), + &ranges, + &mut decoded, + &mut ob_res, + &key, + ) + .unwrap_err(); + assert!(decoded.is_empty()); + assert_decode_error_eq(err, expected_err); +} + +#[cfg(feature = "tokio_fsm")] +async fn keyed_wrong_key_decode_fsm_async_impl( + data: &[u8], + block_size: BlockSize, + expected_err: Option, +) { + use crate::io::fsm::{keyed_decode_ranges, keyed_encode_ranges_validated}; + + let key_a = blake3::derive_key("bao-tree.test", b"key-a"); + let key_b = blake3::derive_key("bao-tree.test", b"key-b"); + let mut outboard = PostOrderMemOutboard::create_keyed(data, block_size, &key_a); + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + keyed_encode_ranges_validated( + Bytes::from(data.to_vec()), + &mut outboard, + &ranges, + &mut encoded, + &key_a, + ) + .await + .unwrap(); + let tree = outboard.tree(); + let mut decoded = bytes::BytesMut::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + let err = keyed_decode_ranges( + Cursor::new(encoded.as_slice()), + ranges, + &mut decoded, + &mut ob_res, + &key_b, + ) + .await + .unwrap_err(); + assert!(decoded.is_empty()); + match expected_err { + Some(expected) => assert_decode_error_eq(err, expected), + None => assert!(matches!( + err, + DecodeError::ParentHashMismatch(_) | DecodeError::LeafHashMismatch(_) + )), + } +} + +#[cfg(feature = "tokio_fsm")] +async fn keyed_wrong_key_fails_encode_fsm_async_impl( + data: &[u8], + block_size: BlockSize, + expected_err: EncodeError, +) { + use crate::io::fsm::keyed_encode_ranges_validated; + + let key_a = blake3::derive_key("bao-tree.test", b"key-a"); + let key_b = blake3::derive_key("bao-tree.test", b"key-b"); + let mut outboard = PostOrderMemOutboard::create_keyed(data, block_size, &key_a); + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + let err = keyed_encode_ranges_validated( + Bytes::from(data.to_vec()), + &mut outboard, + &ranges, + &mut encoded, + &key_b, + ) + .await + .unwrap_err(); + assert!(encoded.is_empty()); + assert_encode_error_eq(err, expected_err); +} + +#[cfg(feature = "tokio_fsm")] +async fn unkeyed_encode_keyed_decode_fails_fsm_async_impl( + data: &[u8], + block_size: BlockSize, + expected_err: DecodeError, +) { + use crate::io::fsm::{encode_ranges_validated, keyed_decode_ranges}; + + let key = blake3::derive_key("bao-tree.test", b"keyed-decode-fsm"); + let mut outboard = PostOrderMemOutboard::create(data, block_size); + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + encode_ranges_validated( + Bytes::from(data.to_vec()), + &mut outboard, + &ranges, + &mut encoded, + ) + .await + .unwrap(); + let tree = outboard.tree(); + let mut decoded = bytes::BytesMut::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + let err = keyed_decode_ranges( + Cursor::new(encoded.as_slice()), + ranges, + &mut decoded, + &mut ob_res, + &key, + ) + .await + .unwrap_err(); + assert!(decoded.is_empty()); + assert_decode_error_eq(err, expected_err); +} + +fn keyed_bao_tree_slice_roundtrip_test( + data: Vec, + mut range: Range, + block_size: BlockSize, + key: &[u8; 32], +) { + use crate::io::sync::{keyed_encode_ranges_validated, DecodeResponseIter}; + + if range.start == range.end { + range.end.0 += 1; + } + let outboard = PostOrderMemOutboard::create_keyed(&data, block_size, key); + let ranges = ChunkRanges::from(range.clone()); + let mut encoded = Vec::new(); + keyed_encode_ranges_validated(&data, &outboard, &ranges, &mut encoded, key).unwrap(); + let expected = data.clone(); + let tree = outboard.tree(); + let iter = + DecodeResponseIter::new_keyed(outboard.root(), tree, Cursor::new(&encoded), &ranges, key); + let mut all_ranges: RangeSet2 = RangeSet2::empty(); + for item in iter { + match item.unwrap() { + BaoContentItem::Leaf(Leaf { offset, data }) => { + all_ranges |= RangeSet2::from(offset..offset + (data.len() as u64)); + let pos = offset.try_into().unwrap(); + assert_eq!(expected[pos..pos + data.len()], *data); + } + BaoContentItem::Parent(_) => {} + } + } + let byte_start = range.start.to_bytes(); + let byte_end = range.end.to_bytes().min(data.len() as u64); + let expected_coverage = RangeSet2::from(byte_start..byte_end); + assert_eq!(all_ranges, expected_coverage); +} + +#[cfg(feature = "tokio_fsm")] +async fn keyed_bao_tree_slice_roundtrip_fsm_test( + data: Vec, + mut range: Range, + block_size: BlockSize, + key: &[u8; 32], +) { + use crate::io::fsm::{keyed_encode_ranges_validated, ResponseDecoder, ResponseDecoderNext}; + + if range.start == range.end { + range.end.0 += 1; + } + let mut outboard = PostOrderMemOutboard::create_keyed(&data, block_size, key); + let ranges = ChunkRanges::from(range.clone()); + let mut encoded = Vec::new(); + keyed_encode_ranges_validated( + Bytes::from(data.clone()), + &mut outboard, + &ranges, + &mut encoded, + key, + ) + .await + .unwrap(); + let expected = data.clone(); + let tree = outboard.tree(); + let mut reading = ResponseDecoder::new_keyed( + outboard.root(), + ranges, + tree, + Cursor::new(encoded.as_slice()), + key, + ); + let mut all_ranges: RangeSet2 = RangeSet2::empty(); + while let ResponseDecoderNext::More((next, result)) = reading.next().await { + reading = next; + match result.unwrap() { + BaoContentItem::Leaf(Leaf { offset, data }) => { + all_ranges |= RangeSet2::from(offset..offset + (data.len() as u64)); + let pos = offset.try_into().unwrap(); + assert_eq!(expected[pos..pos + data.len()], *data); + } + BaoContentItem::Parent(_) => {} + } + } + let byte_start = range.start.to_bytes(); + let byte_end = range.end.to_bytes().min(data.len() as u64); + let expected_coverage = RangeSet2::from(byte_start..byte_end); + assert_eq!(all_ranges, expected_coverage); +} + /// Computes a reference pre order outboard using the bao crate (chunk_group_log = 0) and then flips it to a post-order outboard. fn post_order_outboard_bao(data: &[u8]) -> PostOrderMemOutboard { let mut outboard = Vec::new(); @@ -629,6 +1044,7 @@ fn encode_selected_rec_cases() { min_level, true, &mut actual_encoded, + None, ); actual_encoded.len() - data.len() }; @@ -654,6 +1070,7 @@ fn encode_selected_reference( max_skip_level, true, &mut res, + None, ); (hash, res) } @@ -742,6 +1159,651 @@ fn outboard_hash() { } } +#[test] +fn keyed_outboard_root_matches_blake3() { + let data = make_test_data(100_000); + let key = blake3::derive_key("bao-tree.test", b"format-1"); + for block_level in 0..=4u8 { + let outboard = PostOrderMemOutboard::create_keyed(&data, BlockSize(block_level), &key); + assert_eq!(outboard.root(), blake3::keyed_hash(&key, &data)); + } +} + +#[test] +fn keyed_domain_separation() { + let data = make_test_data(50_000); + let key1 = blake3::derive_key("bao-tree.test", b"format-1"); + let key2 = blake3::derive_key("bao-tree.test", b"format-2"); + let root1 = PostOrderMemOutboard::create_keyed(&data, BlockSize(2), &key1).root(); + let root2 = PostOrderMemOutboard::create_keyed(&data, BlockSize(2), &key2).root(); + assert_ne!(root1, root2); + assert_ne!(root1, blake3::hash(&data)); +} + +#[test] +fn keyed_encode_decode_roundtrip_sync() { + use crate::io::sync::{keyed_decode_ranges, keyed_encode_ranges_validated}; + + let data = make_test_data(50_000); + let key = blake3::derive_key("bao-tree.test", b"roundtrip"); + let block_size = BlockSize(2); + let outboard = PostOrderMemOutboard::create_keyed(&data, block_size, &key); + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + keyed_encode_ranges_validated(&data, &outboard, &ranges, &mut encoded, &key).unwrap(); + let size = outboard.tree.size; + let tree = BaoTree::new(size, block_size); + let mut decoded = Vec::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + keyed_decode_ranges( + Cursor::new(encoded), + &ranges, + &mut decoded, + &mut ob_res, + &key, + ) + .unwrap(); + assert_eq!(decoded, data); + assert_eq!(ob_res.root(), outboard.root()); +} + +#[test] +fn keyed_encode_decode_roundtrip_fsm() { + use crate::io::fsm::{keyed_decode_ranges, keyed_encode_ranges_validated}; + + let data = make_test_data(50_000); + let key = blake3::derive_key("bao-tree.test", b"roundtrip"); + let block_size = BlockSize(2); + let mut outboard = PostOrderMemOutboard::create_keyed(&data, block_size, &key); + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + tokio::runtime::Runtime::new() + .unwrap() + .block_on(keyed_encode_ranges_validated( + Bytes::from(data.clone()), + &mut outboard, + &ranges, + &mut encoded, + &key, + )) + .unwrap(); + let size = outboard.tree.size; + let tree = BaoTree::new(size, block_size); + let mut decoded = bytes::BytesMut::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + tokio::runtime::Runtime::new() + .unwrap() + .block_on(keyed_decode_ranges( + Cursor::new(encoded.as_slice()), + ranges, + &mut decoded, + &mut ob_res, + &key, + )) + .unwrap(); + assert_eq!(decoded.to_vec(), data); + assert_eq!(ob_res.root(), outboard.root()); +} + +#[test] +fn keyed_hash_subtree_differs_from_standard() { + use blake3::hazmat::HasherExt; + + let data = make_test_data(2048); + let key = blake3::derive_key("bao-tree.test", b"low-level-subtree"); + let standard = hash_subtree(0, &data, true); + let keyed = keyed_hash_subtree(0, &data, true, &key); + assert_ne!(standard, keyed); + assert_eq!(keyed, blake3::keyed_hash(&key, &data)); + let non_root_standard = hash_subtree(1, &data[..1024], false); + let non_root_keyed = keyed_hash_subtree(1, &data[..1024], false, &key); + assert_ne!(non_root_standard, non_root_keyed); + let mut hasher = blake3::Hasher::new_keyed(&key); + hasher.set_input_offset(1024); + hasher.update(&data[..1024]); + let expected_non_root = blake3::Hash::from(hasher.finalize_non_root()); + assert_eq!(non_root_keyed, expected_non_root); +} + +#[test] +fn keyed_parent_cv_differs_from_standard() { + use blake3::hazmat::{merge_subtrees_non_root, merge_subtrees_root, ChainingValue, Mode}; + + let left = blake3::hash(b"left"); + let right = blake3::hash(b"right"); + let key = blake3::derive_key("bao-tree.test", b"low-level-parent"); + let standard = parent_cv(&left, &right, true); + let keyed = keyed_parent_cv(&left, &right, true, &key); + assert_ne!(standard, keyed); + let standard_non_root = parent_cv(&left, &right, false); + let keyed_non_root = keyed_parent_cv(&left, &right, false, &key); + assert_ne!(standard_non_root, keyed_non_root); + let left_cv: ChainingValue = *left.as_bytes(); + let right_cv: ChainingValue = *right.as_bytes(); + let mode = Mode::KeyedHash(&key); + assert_eq!(keyed, merge_subtrees_root(&left_cv, &right_cv, mode)); + assert_eq!( + keyed_non_root, + blake3::Hash::from(merge_subtrees_non_root(&left_cv, &right_cv, mode)) + ); +} + +#[test] +fn keyed_pre_order_outboard_root_matches_blake3() { + let data = make_test_data(10_000); + let key = blake3::derive_key("bao-tree.test", b"pre-order"); + for block_level in 0..=4u8 { + let outboard = PreOrderMemOutboard::create_keyed(&data, BlockSize(block_level), &key); + assert_eq!(outboard.root(), blake3::keyed_hash(&key, &data)); + } +} + +#[test] +fn keyed_create_outboard_trait_sync() { + use crate::io::sync::CreateOutboard; + + let data = make_test_data(5000); + let key = blake3::derive_key("bao-tree.test", b"create-outboard"); + let block_size = BlockSize(2); + let post: PostOrderOutboard> = + PostOrderOutboard::create_keyed(Cursor::new(&data), block_size, &key).unwrap(); + assert_eq!(post.root(), blake3::keyed_hash(&key, &data)); + let pre: PreOrderOutboard> = + PreOrderOutboard::create_keyed(Cursor::new(&data), block_size, &key).unwrap(); + assert_eq!(pre.root(), blake3::keyed_hash(&key, &data)); +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn keyed_create_outboard_trait_fsm() { + use crate::io::fsm::CreateOutboard; + + let data = make_test_data(5000); + let key = blake3::derive_key("bao-tree.test", b"create-outboard-fsm"); + let block_size = BlockSize(2); + let post: PostOrderOutboard> = + PostOrderOutboard::create_keyed(Bytes::from(data.clone()), block_size, &key) + .await + .unwrap(); + assert_eq!(post.root(), blake3::keyed_hash(&key, &data)); + let pre: PreOrderOutboard> = + PreOrderOutboard::create_keyed(Bytes::from(data.clone()), block_size, &key) + .await + .unwrap(); + assert_eq!(pre.root(), blake3::keyed_hash(&key, &data)); +} + +#[test] +fn keyed_create_sized_keyed_sync() { + let data = make_test_data(5000); + let key = blake3::derive_key("bao-tree.test", b"create-sized-keyed"); + keyed_create_sized_keyed_checks(&data, BlockSize(2), &key); +} + +#[test] +fn keyed_create_sized_keyed_empty_sync() { + let data: Vec = vec![]; + let key = blake3::derive_key("bao-tree.test", b"create-sized-keyed-empty"); + keyed_create_sized_keyed_checks(&data, BlockSize(0), &key); +} + +#[test] +fn keyed_create_sized_keyed_oversize_sync() { + use crate::io::sync::CreateOutboard; + + let data = make_test_data(100); + let key = blake3::derive_key("bao-tree.test", b"create-sized-keyed-oversize"); + let oversize = data.len() as u64 + 100; + assert!(PostOrderOutboard::>::create_sized_keyed( + Cursor::new(&data), + oversize, + BlockSize(0), + &key + ) + .is_err()); + let tree = BaoTree::new(oversize, BlockSize(0)); + let mut post = PostOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + assert!(post.init_from_keyed(Cursor::new(&data), &key).is_err()); + assert!(PreOrderOutboard::>::create_sized_keyed( + Cursor::new(&data), + oversize, + BlockSize(0), + &key + ) + .is_err()); + let mut pre = PreOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + assert!(pre.init_from_keyed(Cursor::new(&data), &key).is_err()); +} + +#[test] +fn keyed_init_from_keyed_sync() { + let data = make_test_data(5000); + let key = blake3::derive_key("bao-tree.test", b"init-from-keyed"); + keyed_init_from_keyed_checks(&data, BlockSize(2), &key); +} + +#[test] +fn keyed_outboard_functions_sync() { + let data = make_test_data(5000); + let key = blake3::derive_key("bao-tree.test", b"keyed-outboard-fn"); + keyed_outboard_functions_checks(&data, BlockSize(2), &key); +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn keyed_create_sized_keyed_fsm() { + let data = make_test_data(5000); + let key = blake3::derive_key("bao-tree.test", b"create-sized-keyed-fsm"); + keyed_create_sized_keyed_checks_fsm(&data, BlockSize(2), &key).await; +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn keyed_create_sized_keyed_empty_fsm() { + let data: Vec = vec![]; + let key = blake3::derive_key("bao-tree.test", b"create-sized-keyed-empty-fsm"); + keyed_create_sized_keyed_checks_fsm(&data, BlockSize(0), &key).await; +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn keyed_create_sized_keyed_oversize_fsm() { + use crate::io::fsm::CreateOutboard; + + let data = make_test_data(100); + let key = blake3::derive_key("bao-tree.test", b"create-sized-keyed-oversize-fsm"); + let oversize = data.len() as u64 + 100; + assert!(PostOrderOutboard::>::create_sized_keyed( + Cursor::new(Bytes::from(data.clone())), + oversize, + BlockSize(0), + &key + ) + .await + .is_err()); + let tree = BaoTree::new(oversize, BlockSize(0)); + let mut post = PostOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + assert!(post + .init_from_keyed(Cursor::new(Bytes::from(data.clone())), &key) + .await + .is_err()); + assert!(PreOrderOutboard::>::create_sized_keyed( + Cursor::new(Bytes::from(data.clone())), + oversize, + BlockSize(0), + &key + ) + .await + .is_err()); + let mut pre = PreOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + assert!(pre + .init_from_keyed(Cursor::new(Bytes::from(data)), &key) + .await + .is_err()); +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn keyed_init_from_keyed_fsm() { + let data = make_test_data(5000); + let key = blake3::derive_key("bao-tree.test", b"init-from-keyed-fsm"); + keyed_init_from_keyed_checks_fsm(&data, BlockSize(2), &key).await; +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn keyed_outboard_functions_fsm() { + let data = make_test_data(5000); + let key = blake3::derive_key("bao-tree.test", b"keyed-outboard-fn-fsm"); + keyed_outboard_functions_checks_fsm(&data, BlockSize(2), &key).await; +} + +#[test] +fn keyed_wrong_key_fails_decode_sync() { + let data = make_test_data(10_000); + for block_level in 0..=4u8 { + keyed_wrong_key_decode_sync_impl(&data, BlockSize(block_level), None); + } +} + +#[test] +fn keyed_wrong_key_decode_error_variant_sync() { + let multi_chunk = make_test_data(10_000); + keyed_wrong_key_decode_sync_impl( + &multi_chunk, + BlockSize(0), + Some(DecodeError::ParentHashMismatch( + keyed_multi_chunk_mismatch_node(), + )), + ); + let single_byte = make_test_data(1); + keyed_wrong_key_decode_sync_impl( + &single_byte, + BlockSize(0), + Some(DecodeError::LeafHashMismatch(ChunkNum(0))), + ); +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn keyed_wrong_key_fails_decode_fsm() { + let data = make_test_data(10_000); + for block_level in 0..=4u8 { + keyed_wrong_key_decode_fsm_async_impl(&data, BlockSize(block_level), None).await; + } +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn keyed_wrong_key_decode_error_variant_fsm() { + let multi_chunk = make_test_data(10_000); + keyed_wrong_key_decode_fsm_async_impl( + &multi_chunk, + BlockSize(0), + Some(DecodeError::ParentHashMismatch( + keyed_multi_chunk_mismatch_node(), + )), + ) + .await; + let single_byte = make_test_data(1); + keyed_wrong_key_decode_fsm_async_impl( + &single_byte, + BlockSize(0), + Some(DecodeError::LeafHashMismatch(ChunkNum(0))), + ) + .await; +} + +#[test] +fn keyed_wrong_key_fails_encode_sync() { + let multi_chunk = make_test_data(10_000); + keyed_wrong_key_fails_encode_sync_impl( + &multi_chunk, + BlockSize(0), + EncodeError::ParentHashMismatch(keyed_multi_chunk_mismatch_node()), + ); + let single_byte = make_test_data(1); + keyed_wrong_key_fails_encode_sync_impl( + &single_byte, + BlockSize(0), + EncodeError::LeafHashMismatch(ChunkNum(0)), + ); +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn keyed_wrong_key_fails_encode_fsm() { + let multi_chunk = make_test_data(10_000); + keyed_wrong_key_fails_encode_fsm_async_impl( + &multi_chunk, + BlockSize(0), + EncodeError::ParentHashMismatch(keyed_multi_chunk_mismatch_node()), + ) + .await; + let single_byte = make_test_data(1); + keyed_wrong_key_fails_encode_fsm_async_impl( + &single_byte, + BlockSize(0), + EncodeError::LeafHashMismatch(ChunkNum(0)), + ) + .await; +} + +#[test] +fn keyed_outboard_unkeyed_decode_fails_sync() { + use crate::io::sync::{decode_ranges, keyed_encode_ranges_validated}; + + let multi_chunk = make_test_data(10_000); + let key = blake3::derive_key("bao-tree.test", b"unkeyed-decode"); + let outboard = PostOrderMemOutboard::create_keyed(&multi_chunk, BlockSize(0), &key); + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + keyed_encode_ranges_validated(&multi_chunk, &outboard, &ranges, &mut encoded, &key).unwrap(); + let tree = outboard.tree(); + let mut decoded = Vec::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + let err = decode_ranges(Cursor::new(encoded), &ranges, &mut decoded, &mut ob_res).unwrap_err(); + assert!(decoded.is_empty()); + assert_decode_error_eq( + err, + DecodeError::ParentHashMismatch(keyed_multi_chunk_mismatch_node()), + ); + + let single_byte = make_test_data(1); + let outboard = PostOrderMemOutboard::create_keyed(&single_byte, BlockSize(0), &key); + let mut encoded = Vec::new(); + keyed_encode_ranges_validated(&single_byte, &outboard, &ranges, &mut encoded, &key).unwrap(); + let tree = outboard.tree(); + let mut decoded = Vec::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + let err = decode_ranges(Cursor::new(encoded), &ranges, &mut decoded, &mut ob_res).unwrap_err(); + assert!(decoded.is_empty()); + assert_decode_error_eq(err, DecodeError::LeafHashMismatch(ChunkNum(0))); +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn keyed_outboard_unkeyed_decode_fails_fsm() { + use crate::io::fsm::{decode_ranges, keyed_encode_ranges_validated}; + + let multi_chunk = make_test_data(10_000); + let key = blake3::derive_key("bao-tree.test", b"unkeyed-decode-fsm"); + let mut outboard = PostOrderMemOutboard::create_keyed(&multi_chunk, BlockSize(0), &key); + let ranges = ChunkRanges::all(); + let ranges2 = ChunkRanges::all(); + let mut encoded = Vec::new(); + keyed_encode_ranges_validated( + Bytes::from(multi_chunk.clone()), + &mut outboard, + &ranges, + &mut encoded, + &key, + ) + .await + .unwrap(); + let tree = outboard.tree(); + let mut decoded = bytes::BytesMut::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + let err = decode_ranges( + Cursor::new(encoded.as_slice()), + ranges, + &mut decoded, + &mut ob_res, + ) + .await + .unwrap_err(); + assert!(decoded.is_empty()); + assert_decode_error_eq( + err, + DecodeError::ParentHashMismatch(keyed_multi_chunk_mismatch_node()), + ); + + let single_byte = make_test_data(1); + let mut outboard = PostOrderMemOutboard::create_keyed(&single_byte, BlockSize(0), &key); + let mut encoded = Vec::new(); + keyed_encode_ranges_validated( + Bytes::from(single_byte.clone()), + &mut outboard, + &ranges2, + &mut encoded, + &key, + ) + .await + .unwrap(); + let tree = outboard.tree(); + let mut decoded = bytes::BytesMut::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + let err = decode_ranges( + Cursor::new(encoded.as_slice()), + ranges2, + &mut decoded, + &mut ob_res, + ) + .await + .unwrap_err(); + assert!(decoded.is_empty()); + assert_decode_error_eq(err, DecodeError::LeafHashMismatch(ChunkNum(0))); +} + +#[test] +fn unkeyed_outboard_keyed_decode_fails_sync() { + let multi_chunk = make_test_data(10_000); + unkeyed_encode_keyed_decode_fails_sync_impl( + &multi_chunk, + BlockSize(0), + DecodeError::ParentHashMismatch(keyed_multi_chunk_mismatch_node()), + ); + let single_byte = make_test_data(1); + unkeyed_encode_keyed_decode_fails_sync_impl( + &single_byte, + BlockSize(0), + DecodeError::LeafHashMismatch(ChunkNum(0)), + ); +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn unkeyed_outboard_keyed_decode_fails_fsm() { + let multi_chunk = make_test_data(10_000); + unkeyed_encode_keyed_decode_fails_fsm_async_impl( + &multi_chunk, + BlockSize(0), + DecodeError::ParentHashMismatch(keyed_multi_chunk_mismatch_node()), + ) + .await; + let single_byte = make_test_data(1); + unkeyed_encode_keyed_decode_fails_fsm_async_impl( + &single_byte, + BlockSize(0), + DecodeError::LeafHashMismatch(ChunkNum(0)), + ) + .await; +} + +#[test] +fn keyed_encode_decode_edge_sizes_sync() { + use make_test_data as td; + + let key = blake3::derive_key("bao-tree.test", b"edge"); + let block_size = BlockSize(0); + for size in [0, 1, 1024, 1025] { + keyed_encode_decode_roundtrip_sync_impl(&td(size), block_size, &key); + } +} + +#[test] +fn keyed_encode_decode_edge_sizes_fsm() { + use make_test_data as td; + + let key = blake3::derive_key("bao-tree.test", b"edge"); + let block_size = BlockSize(0); + for size in [0, 1, 1024, 1025] { + keyed_encode_decode_roundtrip_fsm_impl(td(size), block_size, &key); + } +} + +fn keyed_bao_tree_slice_roundtrip_case_table(key: &[u8; 32]) { + use make_test_data as td; + + let cases = [ + (0, 0..1), + (1, 0..1), + (1023, 0..1), + (1024, 0..1), + (1025, 0..1), + (1025, 0..2), + (1025, 1..2), + (24 * 1024 + 1, 0..25), + ]; + for chunk_group_log in 0..4 { + let block_size = BlockSize(chunk_group_log); + for (count, range) in cases.clone() { + keyed_bao_tree_slice_roundtrip_test( + td(count), + ChunkNum(range.start)..ChunkNum(range.end), + block_size, + key, + ); + } + } +} + +#[test] +fn keyed_bao_tree_slice_roundtrip_cases() { + let key = blake3::derive_key("bao-tree.test", b"slice"); + keyed_bao_tree_slice_roundtrip_case_table(&key); +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn keyed_bao_tree_slice_roundtrip_fsm_cases() { + use make_test_data as td; + + let key = blake3::derive_key("bao-tree.test", b"slice-fsm"); + let cases = [ + (0, 0..1), + (1, 0..1), + (1023, 0..1), + (1024, 0..1), + (1025, 0..1), + (1025, 0..2), + (1025, 1..2), + (24 * 1024 + 1, 0..25), + ]; + for chunk_group_log in 0..4 { + let block_size = BlockSize(chunk_group_log); + for (count, range) in cases.clone() { + keyed_bao_tree_slice_roundtrip_fsm_test( + td(count), + ChunkNum(range.start)..ChunkNum(range.end), + block_size, + &key, + ) + .await; + } + } +} + #[test] fn select_last_chunk_0() { assert_tuple_eq!(select_last_chunk_impl(1, 0)); @@ -915,6 +1977,74 @@ proptest! { /// Checks that the simple recursive impl bao_encode_selected_recursive that /// does not need an outboard is the same as the more complex encode_ranges_validated /// that requires an outboard. + #[test] + fn keyed_encode_selected_reference_sync_proptest( + (size, ranges) in size_and_selection(1..100000, 2), + block_size in 0..5u8, + key_seed in proptest::collection::vec(any::(), 32), + ) { + let key: [u8; 32] = key_seed.try_into().unwrap(); + let data = make_test_data(size); + let expected_hash = blake3::keyed_hash(&key, &data); + let block_size = BlockSize(block_size); + let (actual_hash, actual_encoded) = + keyed_encode_selected_reference(&data, block_size, &ranges, &key); + let mut expected_encoded = Vec::new(); + let outboard = PostOrderMemOutboard::create_keyed(&data, block_size, &key); + crate::io::sync::keyed_encode_ranges_validated( + &data, + &outboard, + &ranges, + &mut expected_encoded, + &key, + ) + .unwrap(); + prop_assert_eq!(expected_hash, actual_hash); + prop_assert_eq!(hex::encode(expected_encoded), hex::encode(actual_encoded)); + } + + #[test] + fn keyed_encode_selected_reference_fsm_proptest( + (size, ranges) in size_and_selection(1..100000, 2), + block_size in 0..4u8, + key_seed in proptest::collection::vec(any::(), 32), + ) { + let key: [u8; 32] = key_seed.try_into().unwrap(); + let data = make_test_data(size); + let expected_hash = blake3::keyed_hash(&key, &data); + let block_size = BlockSize(block_size); + let (actual_hash, actual_encoded) = + keyed_encode_selected_reference(&data, block_size, &ranges, &key); + let mut expected_encoded = Vec::new(); + let outboard = PostOrderMemOutboard::create_keyed(&data, block_size, &key); + let data: Bytes = data.into(); + tokio::runtime::Runtime::new().unwrap().block_on( + crate::io::fsm::keyed_encode_ranges_validated( + data, + outboard, + &ranges, + &mut expected_encoded, + &key, + ), + ) + .unwrap(); + prop_assert_eq!(expected_hash, actual_hash); + prop_assert_eq!(expected_encoded, actual_encoded); + } + + #[test] + fn keyed_bao_tree_slice_roundtrip_proptest( + (len, start, size) in size_and_slice_overlapping(), + level in 0u8..6, + key_seed in proptest::collection::vec(any::(), 32), + ) { + let key: [u8; 32] = key_seed.try_into().unwrap(); + let level = BlockSize(level); + let data = make_test_data(len as usize); + let chunk_range = start .. start + size; + keyed_bao_tree_slice_roundtrip_test(data, chunk_range, level, &key); + } + #[test] fn encode_selected_reference_sync_proptest((size, ranges) in size_and_selection(1..100000, 2), block_size in 0..5u8) { let data = make_test_data(size); diff --git a/src/tests2.rs b/src/tests2.rs index 125364e..01d4c3f 100644 --- a/src/tests2.rs +++ b/src/tests2.rs @@ -25,15 +25,20 @@ use crate::{ BaoContentItem, Leaf, Parent, }, iter::{BaoChunk, PreOrderPartialChunkIterRef, ResponseIterRef}, - parent_cv, prop_assert_tuple_eq, + keyed_hash_subtree, keyed_parent_cv, parent_cv, prop_assert_tuple_eq, rec::{ - encode_selected_rec, get_leaf_ranges, make_test_data, partial_chunk_iter_reference, - range_union, response_iter_reference, select_nodes_rec, truncate_ranges, - ReferencePreOrderPartialChunkIterRef, + encode_selected_rec, get_leaf_ranges, keyed_create_sized_keyed_checks, + keyed_init_from_keyed_checks, keyed_outboard_functions_checks, make_test_data, + partial_chunk_iter_reference, range_union, response_iter_reference, select_nodes_rec, + truncate_ranges, ReferencePreOrderPartialChunkIterRef, }, BaoTree, BlockSize, ChunkNum, ChunkRanges, ChunkRangesRef, TreeNode, }; +fn keyed_test_key(context: &[u8]) -> [u8; 32] { + blake3::derive_key("bao-tree.test", context) +} + fn tree() -> impl Strategy { (0u64..100000, 0u8..5).prop_map(|(size, block_size)| { let block_size = BlockSize(block_size); @@ -141,6 +146,26 @@ fn post_traversal_chunks_iter_proptest(#[strategy(tree())] tree: BaoTree) { post_traversal_chunks_iter_impl(tree); } +/// Brute force test for a keyed outboard that computes expected hashes for each pair +fn keyed_outboard_test_sync(data: &[u8], outboard: impl crate::io::sync::Outboard, key: &[u8; 32]) { + let tree = outboard.tree(); + let nodes = tree + .pre_order_nodes_iter() + .enumerate() + .map(|(i, node)| (node, i == 0)) + .filter(|(node, _)| tree.is_relevant_for_outboard(*node)) + .collect::>(); + for (node, is_root) in nodes { + let (l_hash, r_hash) = outboard.load(node).unwrap().unwrap(); + let start_chunk = node.chunk_range().start; + let byte_range = tree.byte_range(node); + let data = &data[byte_range.start.try_into().unwrap()..byte_range.end.try_into().unwrap()]; + let expected = keyed_hash_subtree(start_chunk.0, data, is_root, key); + let actual = keyed_parent_cv(&l_hash, &r_hash, is_root, key); + assert_eq!(actual, expected); + } +} + /// Brute force test for an outboard that just computes the expected hash for each pair fn outboard_test_sync(data: &[u8], outboard: impl crate::io::sync::Outboard) { let tree = outboard.tree(); @@ -161,6 +186,30 @@ fn outboard_test_sync(data: &[u8], outboard: impl crate::io::sync::Outboard) { } } +/// Brute force test for a keyed outboard that computes expected hashes for each pair +async fn keyed_outboard_test_fsm( + data: &[u8], + mut outboard: impl crate::io::fsm::Outboard, + key: &[u8; 32], +) { + let tree = outboard.tree(); + let nodes = tree + .pre_order_nodes_iter() + .enumerate() + .map(|(i, node)| (node, i == 0)) + .filter(|(node, _)| tree.is_relevant_for_outboard(*node)) + .collect::>(); + for (node, is_root) in nodes { + let (l_hash, r_hash) = outboard.load(node).await.unwrap().unwrap(); + let start_chunk = node.chunk_range().start; + let byte_range = tree.byte_range(node); + let data = &data[byte_range.start.try_into().unwrap()..byte_range.end.try_into().unwrap()]; + let expected = keyed_hash_subtree(start_chunk.0, data, is_root, key); + let actual = keyed_parent_cv(&l_hash, &r_hash, is_root, key); + assert_eq!(actual, expected); + } +} + /// Brute force test for an outboard that just computes the expected hash for each pair async fn outboard_test_fsm(data: &[u8], mut outboard: impl crate::io::fsm::Outboard) { let tree = outboard.tree(); @@ -222,6 +271,119 @@ fn post_oder_outboard_fsm_proptest(#[strategy(tree())] tree: BaoTree) { post_oder_outboard_fsm_impl(tree); } +fn keyed_post_order_outboard_sync_impl(tree: BaoTree) { + let data = make_test_data(tree.size.try_into().unwrap()); + let key = keyed_test_key(&tree.size.to_le_bytes()); + let outboard = PostOrderMemOutboard::create_keyed(&data, tree.block_size, &key); + assert_eq!( + outboard.data.len() as u64, + outboard.tree().outboard_hash_pairs() * 64 + ); + keyed_outboard_test_sync(&data, outboard, &key); +} + +#[proptest] +fn keyed_post_order_outboard_sync_proptest(#[strategy(tree())] tree: BaoTree) { + keyed_post_order_outboard_sync_impl(tree); +} + +fn keyed_post_order_outboard_fsm_impl(tree: BaoTree) { + let data = make_test_data(tree.size.try_into().unwrap()); + let key = keyed_test_key(&tree.size.to_le_bytes()); + let outboard = PostOrderMemOutboard::create_keyed(&data, tree.block_size, &key); + assert_eq!( + outboard.data.len() as u64, + outboard.tree().outboard_hash_pairs() * 64 + ); + tokio::runtime::Runtime::new() + .unwrap() + .block_on(keyed_outboard_test_fsm(&data, outboard, &key)); +} + +#[proptest] +fn keyed_post_order_outboard_fsm_proptest(#[strategy(tree())] tree: BaoTree) { + keyed_post_order_outboard_fsm_impl(tree); +} + +fn keyed_pre_order_outboard_sync_impl(tree: BaoTree) { + let data = make_test_data(tree.size.try_into().unwrap()); + let key = keyed_test_key(&tree.size.to_le_bytes()); + let outboard = PreOrderMemOutboard::create_keyed(&data, tree.block_size, &key); + assert_eq!( + outboard.data.len(), + outboard.tree().outboard_size().try_into().unwrap() + ); + keyed_outboard_test_sync(&data, outboard, &key); +} + +#[proptest] +fn keyed_pre_order_outboard_sync_proptest(#[strategy(tree())] tree: BaoTree) { + keyed_pre_order_outboard_sync_impl(tree); +} + +fn keyed_pre_order_outboard_fsm_impl(tree: BaoTree) { + let data = make_test_data(tree.size.try_into().unwrap()); + let key = keyed_test_key(&tree.size.to_le_bytes()); + let outboard = PreOrderMemOutboard::create_keyed(&data, tree.block_size, &key); + assert_eq!( + outboard.data.len(), + outboard.tree().outboard_size().try_into().unwrap() + ); + tokio::runtime::Runtime::new() + .unwrap() + .block_on(keyed_outboard_test_fsm(&data, outboard, &key)); +} + +#[proptest] +fn keyed_pre_order_outboard_fsm_proptest(#[strategy(tree())] tree: BaoTree) { + keyed_pre_order_outboard_fsm_impl(tree); +} + +#[proptest] +fn keyed_create_sized_keyed_proptest(#[strategy(tree())] tree: BaoTree) { + let data = make_test_data(tree.size.try_into().unwrap()); + let key = keyed_test_key(&tree.size.to_le_bytes()); + keyed_create_sized_keyed_checks(&data, tree.block_size, &key); +} + +#[proptest] +fn keyed_init_from_keyed_proptest(#[strategy(tree())] tree: BaoTree) { + let data = make_test_data(tree.size.try_into().unwrap()); + let key = keyed_test_key(&tree.size.to_le_bytes()); + keyed_init_from_keyed_checks(&data, tree.block_size, &key); +} + +#[proptest] +fn keyed_outboard_functions_proptest(#[strategy(tree())] tree: BaoTree) { + let data = make_test_data(tree.size.try_into().unwrap()); + let key = keyed_test_key(&tree.size.to_le_bytes()); + keyed_outboard_functions_checks(&data, tree.block_size, &key); +} + +#[cfg(feature = "tokio_fsm")] +#[proptest] +fn keyed_create_sized_keyed_fsm_proptest(#[strategy(tree())] tree: BaoTree) { + let data = make_test_data(tree.size.try_into().unwrap()); + let key = keyed_test_key(&tree.size.to_le_bytes()); + run_blocking(crate::rec::keyed_create_sized_keyed_checks_fsm( + &data, + tree.block_size, + &key, + )); +} + +#[cfg(feature = "tokio_fsm")] +#[proptest] +fn keyed_init_from_keyed_fsm_proptest(#[strategy(tree())] tree: BaoTree) { + let data = make_test_data(tree.size.try_into().unwrap()); + let key = keyed_test_key(&tree.size.to_le_bytes()); + run_blocking(crate::rec::keyed_init_from_keyed_checks_fsm( + &data, + tree.block_size, + &key, + )); +} + fn mem_outboard_flip_impl(tree: BaoTree) { let data = make_test_data(tree.size.try_into().unwrap()); let post = PostOrderMemOutboard::create(&data, tree.block_size); @@ -279,6 +441,39 @@ mod validate { res } + fn keyed_valid_ranges_sync( + outboard: impl crate::io::sync::Outboard, + data: &[u8], + key: &[u8; 32], + ) -> ChunkRanges { + let ranges = ChunkRanges::all(); + let iter = crate::io::sync::keyed_valid_ranges(outboard, data, &ranges, key); + let mut res = ChunkRanges::empty(); + for item in iter { + let item = item.unwrap(); + res |= ChunkRanges::from(item); + } + res + } + + fn keyed_valid_ranges_fsm( + outboard: impl crate::io::fsm::Outboard, + data: Bytes, + key: &[u8; 32], + ) -> ChunkRanges { + run_blocking(async move { + let ranges = ChunkRanges::all(); + let mut stream = crate::io::fsm::keyed_valid_ranges(outboard, data, &ranges, key); + let mut res = ChunkRanges::empty(); + while let Some(item) = stream.next().await { + let item = item?; + res |= ChunkRanges::from(item); + } + std::io::Result::Ok(res) + }) + .unwrap() + } + /// range is a range of chunks. Just using u64 for convenience in tests fn valid_outboard_ranges_fsm(outboard: &mut PostOrderMemOutboard) -> ChunkRanges { run_blocking(async move { @@ -349,6 +544,123 @@ mod validate { } } + fn validate_keyed_pos_impl(tree: BaoTree) { + let size = tree.size.try_into().unwrap(); + let block_size = tree.block_size; + let data = make_test_data(size); + let key = blake3::derive_key("bao-tree.test", b"valid-ranges"); + let mut outboard = PostOrderMemOutboard::create_keyed(&data, block_size, &key); + let expected = ChunkRanges::from(..outboard.tree().chunks()); + let actual = keyed_valid_ranges_sync(&outboard, &data, &key); + assert_eq!(expected, actual); + let actual = keyed_valid_ranges_fsm(&mut outboard, data.into(), &key); + assert_eq!(expected, actual); + } + + #[proptest] + fn validate_keyed_pos_proptest(#[strategy(tree())] tree: BaoTree) { + validate_keyed_pos_impl(tree); + } + + #[test] + fn validate_keyed_pos_cases() { + let cases = [(0x401, 0), (0, 0), (1, 0), (1024, 0), (1025, 2)]; + for (size, block_level) in cases { + let tree = BaoTree::new(size, BlockSize(block_level)); + validate_keyed_pos_impl(tree); + } + } + + fn keyed_chunk_count(ranges: &ChunkRanges) -> u64 { + ranges + .boundaries() + .windows(2) + .map(|w| (w[1] - w[0]).0) + .sum() + } + + fn assert_keyed_valid_ranges_wrong_key( + outboard: &PostOrderMemOutboard, + data: &[u8], + wrong_key: &[u8; 32], + expected: &ChunkRanges, + ) { + let actual = keyed_valid_ranges_sync(outboard, data, wrong_key); + assert!(expected.is_superset(&actual)); + assert_ne!(actual, *expected); + let expected_chunks = keyed_chunk_count(expected); + let actual_chunks = keyed_chunk_count(&actual); + assert!(actual_chunks < expected_chunks); + let actual_fsm = keyed_valid_ranges_fsm(outboard.clone(), data.to_vec().into(), wrong_key); + assert!(expected.is_superset(&actual_fsm)); + assert_ne!(actual_fsm, *expected); + let actual_fsm_chunks = keyed_chunk_count(&actual_fsm); + assert!(actual_fsm_chunks < expected_chunks); + } + + fn validate_keyed_neg_impl(tree: BaoTree) { + let size = tree.size.try_into().unwrap(); + let block_size = tree.block_size; + let data = make_test_data(size); + let key = blake3::derive_key("bao-tree.test", b"valid-ranges"); + let wrong_key = blake3::derive_key("bao-tree.test", b"wrong-key"); + let outboard = PostOrderMemOutboard::create_keyed(&data, block_size, &key); + let expected = ChunkRanges::from(..outboard.tree().chunks()); + if size > 0 { + assert_keyed_valid_ranges_wrong_key(&outboard, &data, &wrong_key, &expected); + } + } + + #[test] + fn validate_keyed_neg_cases() { + let cases = [(0x2001, 0), (1025, 1)]; + for (size, block_level) in cases { + let tree = BaoTree::new(size, BlockSize(block_level)); + validate_keyed_neg_impl(tree); + } + } + + #[proptest] + fn validate_keyed_neg_proptest(#[strategy(tree())] tree: BaoTree) { + if tree.size > 0 { + validate_keyed_neg_impl(tree); + } + } + + /// Check that flipping a random bit in a keyed outboard makes at least one range invalid + fn validate_keyed_outboard_neg_impl(tree: BaoTree, rand: u32) { + let rand = rand as usize; + let size = tree.size.try_into().unwrap(); + let block_size = tree.block_size; + let data = make_test_data(size); + let key = blake3::derive_key("bao-tree.test", b"valid-ranges"); + let mut outboard = PostOrderMemOutboard::create_keyed(&data, block_size, &key); + let expected = ChunkRanges::from(..outboard.tree().chunks()); + if !outboard.data.is_empty() { + flip_bit(&mut outboard.data, rand); + let actual = keyed_valid_ranges_sync(&outboard, &data, &key); + assert_ne!(expected, actual); + let actual_fsm = keyed_valid_ranges_fsm(outboard.clone(), data.into(), &key); + assert_ne!(expected, actual_fsm); + } + } + + #[test] + fn validate_keyed_outboard_neg_cases() { + let cases = [((0x2001, 0), 2738363904)]; + for ((size, block_level), rand) in cases { + let tree = BaoTree::new(size, BlockSize(block_level)); + validate_keyed_outboard_neg_impl(tree, rand); + } + } + + #[proptest] + fn validate_keyed_outboard_neg_proptest(#[strategy(tree())] tree: BaoTree, rand: u32) { + if tree.size > 0 && tree.outboard_hash_pairs() > 0 { + validate_keyed_outboard_neg_impl(tree, rand); + } + } + fn flip_bit(data: &mut [u8], rand: usize) { // flip a random bit in the outboard // this is the post order outboard without the length suffix, @@ -535,6 +847,176 @@ async fn encode_decode_full_fsm_impl( ((data, outboard), (decoded.to_vec(), ob_res)) } +fn keyed_encode_decode_full_sync_impl( + data: &[u8], + outboard: PostOrderMemOutboard, + key: &[u8; 32], +) -> ( + (Vec, PostOrderMemOutboard), + (Vec, PostOrderMemOutboard), +) { + let ranges = ChunkRanges::all(); + let size = outboard.tree.size; + let mut encoded = Vec::new(); + crate::io::sync::keyed_encode_ranges_validated(data, &outboard, &ranges, &mut encoded, key) + .unwrap(); + let encoded_read = std::io::Cursor::new(encoded); + let tree = BaoTree::new(size, outboard.tree().block_size()); + let mut decoded = Vec::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + crate::io::sync::keyed_decode_ranges(encoded_read, &ranges, &mut decoded, &mut ob_res, key) + .unwrap(); + ((decoded, ob_res), (data.to_vec(), outboard)) +} + +async fn keyed_encode_decode_full_fsm_impl( + data: Vec, + outboard: PostOrderMemOutboard, + key: &[u8; 32], +) -> ( + (Vec, PostOrderMemOutboard), + (Vec, PostOrderMemOutboard), +) { + let size = outboard.tree.size; + let mut outboard = outboard; + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + crate::io::fsm::keyed_encode_ranges_validated( + Bytes::from(data.clone()), + &mut outboard, + &ranges, + &mut encoded, + key, + ) + .await + .unwrap(); + + let read_encoded = std::io::Cursor::new(encoded.as_slice()); + let mut ob_res = { + let tree = BaoTree::new(size, outboard.tree().block_size()); + let root = outboard.root(); + let outboard_size = usize::try_from(tree.outboard_hash_pairs() * 64).unwrap(); + let outboard_data = vec![0u8; outboard_size]; + PostOrderMemOutboard { + root, + tree, + data: outboard_data, + } + }; + let mut decoded = BytesMut::new(); + crate::io::fsm::keyed_decode_ranges(read_encoded, ranges, &mut decoded, &mut ob_res, key) + .await + .unwrap(); + ((data, outboard), (decoded.to_vec(), ob_res)) +} + +fn keyed_encode_decode_partial_sync_impl( + data: &[u8], + outboard: PostOrderMemOutboard, + ranges: &ChunkRangesRef, + key: &[u8; 32], +) -> bool { + let mut encoded = Vec::new(); + let size = outboard.tree.size; + crate::io::sync::keyed_encode_ranges_validated(data, &outboard, ranges, &mut encoded, key) + .unwrap(); + let expected_data = data; + let encoded_read = std::io::Cursor::new(encoded); + let tree = BaoTree::new(size, outboard.tree.block_size); + let iter = crate::io::sync::DecodeResponseIter::new_keyed( + outboard.root, + tree, + encoded_read, + ranges, + key, + ); + for item in iter { + let item = match item { + Ok(item) => item, + Err(_) => { + return false; + } + }; + match item { + BaoContentItem::Parent(Parent { node, pair }) => { + if let Some(expected_pair) = outboard.load(node).unwrap() { + if pair != expected_pair { + return false; + } + } + } + BaoContentItem::Leaf(Leaf { offset, data }) => { + let offset = offset.try_into().unwrap(); + if expected_data[offset..offset + data.len()] != data { + return false; + } + } + } + } + true +} + +async fn keyed_encode_decode_partial_fsm_impl( + data: &[u8], + outboard: PostOrderMemOutboard, + ranges: ChunkRanges, + key: &[u8; 32], +) -> bool { + let size = outboard.tree.size; + let mut encoded = Vec::new(); + let mut outboard = outboard; + crate::io::fsm::keyed_encode_ranges_validated( + Bytes::from(data.to_vec()), + &mut outboard, + &ranges, + &mut encoded, + key, + ) + .await + .unwrap(); + let expected_data = data; + let encoded_read = std::io::Cursor::new(encoded.as_slice()); + let mut reading = crate::io::fsm::ResponseDecoder::new_keyed( + outboard.root, + ranges, + BaoTree::new(size, outboard.tree.block_size), + encoded_read, + key, + ); + if size != outboard.tree.size { + return false; + } + while let ResponseDecoderNext::More((reading1, result)) = reading.next().await { + let item = match result { + Ok(item) => item, + Err(_) => { + return false; + } + }; + match item { + BaoContentItem::Leaf(Leaf { offset, data }) => { + let offset: usize = offset.try_into().unwrap(); + if expected_data[offset..offset + data.len()] != data { + return false; + } + } + BaoContentItem::Parent(Parent { node, pair }) => { + if let Some(expected_pair) = outboard.load(node).unwrap() { + if pair != expected_pair { + return false; + } + } + } + } + reading = reading1; + } + true +} + fn encode_decode_partial_sync_impl( data: &[u8], outboard: PostOrderMemOutboard, @@ -662,6 +1144,27 @@ fn encode_decode_partial_sync_proptest( prop_assert!(ok); } +#[proptest] +fn keyed_encode_decode_full_sync_proptest(#[strategy(tree())] tree: BaoTree) { + let data = make_test_data(tree.size.try_into().unwrap()); + let key = keyed_test_key(&tree.size.to_le_bytes()); + let outboard = PostOrderMemOutboard::create_keyed(&data, tree.block_size, &key); + prop_assert_tuple_eq!(keyed_encode_decode_full_sync_impl(&data, outboard, &key)); +} + +#[proptest] +fn keyed_encode_decode_partial_sync_proptest( + #[strategy(size_and_selection(0..100000, 2))] size_and_selection: (usize, ChunkRanges), + #[strategy(block_size())] block_size: BlockSize, +) { + let (size, selection) = size_and_selection; + let data = make_test_data(size); + let key = keyed_test_key(&(size as u64).to_le_bytes()); + let outboard = PostOrderMemOutboard::create_keyed(&data, block_size, &key); + let ok = keyed_encode_decode_partial_sync_impl(&data, outboard, &selection, &key); + prop_assert!(ok); +} + #[test] fn encode_decode_full_fsm_cases() { let cases = [BaoTree::new(0x1001, BlockSize(1))]; @@ -699,6 +1202,35 @@ fn encode_decode_partial_fsm_proptest( prop_assert!(ok); } +#[proptest] +fn keyed_encode_decode_full_fsm_proptest(#[strategy(tree())] tree: BaoTree) { + let data = make_test_data(tree.size.try_into().unwrap()); + let key = keyed_test_key(&tree.size.to_le_bytes()); + let outboard = PostOrderMemOutboard::create_keyed(&data, tree.block_size, &key); + let pair = tokio::runtime::Runtime::new() + .unwrap() + .block_on(keyed_encode_decode_full_fsm_impl(data, outboard, &key)); + prop_assert_tuple_eq!(pair); +} + +#[proptest] +fn keyed_encode_decode_partial_fsm_proptest( + #[strategy(size_and_selection(0..100000, 2))] size_and_selection: (usize, ChunkRanges), + #[strategy(block_size())] block_size: BlockSize, +) { + let (size, selection) = size_and_selection; + let data = make_test_data(size); + let key = keyed_test_key(&(size as u64).to_le_bytes()); + let outboard = PostOrderMemOutboard::create_keyed(&data, block_size, &key); + let ok = + tokio::runtime::Runtime::new() + .unwrap() + .block_on(keyed_encode_decode_partial_fsm_impl( + &data, outboard, selection, &key, + )); + prop_assert!(ok); +} + fn pre_order_nodes_iter_reference(tree: BaoTree, ranges: &ChunkRangesRef) -> Vec { let mut res = Vec::new(); select_nodes_rec( @@ -783,6 +1315,7 @@ fn encode_selected_reference( max_skip_level, true, &mut res, + None, ); (hash, res) } From 02916e784bb0afe0fd5a73c291c8c5335865e166 Mon Sep 17 00:00:00 2001 From: Hunter Beast Date: Mon, 6 Jul 2026 10:20:48 -0600 Subject: [PATCH 02/12] Replace Option key threading with BaoHashing trait Standard and keyed modes share one IO implementation through compile time strategy types instead of Option branches. Public keyed APIs are unchanged. Adds comments and hash_strategy naming for clarity. Fixes dead_code warnings in validate only builds. --- src/io/fsm.rs | 204 +++++++++++++++++++++++++++++-------------------- src/io/sync.rs | 197 ++++++++++++++++++++++++++++------------------- src/lib.rs | 123 +++++++++++++++++++++-------- src/rec.rs | 24 +++--- src/tests.rs | 9 ++- src/tests2.rs | 8 +- 6 files changed, 357 insertions(+), 208 deletions(-) diff --git a/src/io/fsm.rs b/src/io/fsm.rs index 00eba8d..1e5dba3 100644 --- a/src/io/fsm.rs +++ b/src/io/fsm.rs @@ -22,16 +22,15 @@ use smallvec::SmallVec; pub use super::BaoContentItem; use super::{combine_hash_pair, DecodeError}; use crate::{ - blake3, hash_subtree_with_key, + blake3, io::{ error::EncodeError, outboard::{PostOrderOutboard, PreOrderOutboard}, Leaf, Parent, }, iter::{BaoChunk, ResponseIter}, - parent_cv_with_key, rec::{encode_selected_rec, truncate_ranges, truncate_ranges_owned}, - BaoTree, BlockSize, ChunkRanges, ChunkRangesRef, TreeNode, + BaoHashing, BaoTree, BlockSize, ChunkRanges, ChunkRangesRef, Keyed, Standard, TreeNode, }; /// A binary merkle tree for blake3 hashes of a blob. @@ -402,25 +401,24 @@ pub(crate) fn parse_hash_pair(buf: Bytes) -> io::Result<(blake3::Hash, blake3::H Ok((l_hash, r_hash)) } +/// Generic over `H` so keyed and standard decoders share this implementation. #[derive(Debug)] -struct ResponseDecoderInner { +struct ResponseDecoderInner { iter: ResponseIter, stack: SmallVec<[blake3::Hash; 10]>, encoded: R, - key: Option<[u8; 32]>, + /// Compile time hashing strategy, either [Standard] or [Keyed]. + hash_strategy: H, } -impl ResponseDecoderInner { - fn new(tree: BaoTree, hash: blake3::Hash, ranges: ChunkRanges, encoded: R) -> Self { - Self::new_with_key(tree, hash, ranges, encoded, None) - } - - fn new_with_key( +impl ResponseDecoderInner { + /// Shared constructor used by [ResponseDecoder::new] and [ResponseDecoder::new_keyed]. + fn with_hash_strategy( tree: BaoTree, hash: blake3::Hash, ranges: ChunkRanges, encoded: R, - key: Option<[u8; 32]>, + hash_strategy: H, ) -> Self { // now that we know the size, we can canonicalize the ranges let ranges = truncate_ranges_owned(ranges, tree.size()); @@ -428,24 +426,27 @@ impl ResponseDecoderInner { iter: ResponseIter::new(tree, ranges), stack: SmallVec::new(), encoded, - key, + hash_strategy, }; res.stack.push(hash); res } } -/// Response decoder +/// Response decoder. +/// +/// Generic over `H` so keyed and standard decoders share this implementation. +/// Defaults to [Standard]. Use [Self::new_keyed] for keyed responses. #[derive(Debug)] -pub struct ResponseDecoder(Box>); +pub struct ResponseDecoder(Box>); /// Next type for ResponseDecoder. #[derive(Debug)] -pub enum ResponseDecoderNext { +pub enum ResponseDecoderNext { /// One more item, and you get back the state machine in the next state More( ( - ResponseDecoder, + ResponseDecoder, std::result::Result, ), ), @@ -458,8 +459,8 @@ impl ResponseDecoder { /// /// The size as well as the chunk size is given in the `tree` parameter. pub fn new(hash: blake3::Hash, ranges: ChunkRanges, tree: BaoTree, encoded: R) -> Self { - Self(Box::new(ResponseDecoderInner::new( - tree, hash, ranges, encoded, + Self(Box::new(ResponseDecoderInner::with_hash_strategy( + tree, hash, ranges, encoded, Standard, ))) } @@ -470,18 +471,37 @@ impl ResponseDecoder { tree: BaoTree, encoded: R, key: &[u8; 32], + ) -> ResponseDecoder { + ResponseDecoder(Box::new(ResponseDecoderInner::with_hash_strategy( + tree, + hash, + ranges, + encoded, + Keyed(*key), + ))) + } +} + +impl ResponseDecoder { + /// Shared constructor used by decode helpers and public new methods. + pub(crate) fn with_hash_strategy( + hash: blake3::Hash, + ranges: ChunkRanges, + tree: BaoTree, + encoded: R, + hash_strategy: H, ) -> Self { - Self(Box::new(ResponseDecoderInner::new_with_key( + Self(Box::new(ResponseDecoderInner::with_hash_strategy( tree, hash, ranges, encoded, - Some(*key), + hash_strategy, ))) } /// Proceed to the next state by reading the next chunk from the stream. - pub async fn next(mut self) -> ResponseDecoderNext { + pub async fn next(mut self) -> ResponseDecoderNext { if let Some(chunk) = self.0.iter.next() { let item = self.next0(chunk).await; ResponseDecoderNext::More((self, item)) @@ -522,7 +542,7 @@ impl ResponseDecoder { .map_err(|e| DecodeError::maybe_parent_not_found(e, node))?; let pair @ (l_hash, r_hash) = read_parent(&buf); let parent_hash = this.stack.pop().unwrap(); - let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, this.key.as_ref()); + let actual = this.hash_strategy.parent_cv(&l_hash, &r_hash, is_root); // Push the children in reverse order so they are popped in the correct order // only push right if the range intersects with the right child if right { @@ -552,8 +572,9 @@ impl ResponseDecoder { .await .map_err(|e| DecodeError::maybe_leaf_not_found(e, start_chunk))?; let leaf_hash = this.stack.pop().unwrap(); - let actual = - hash_subtree_with_key(start_chunk.0, &data, is_root, this.key.as_ref()); + let actual = this + .hash_strategy + .hash_subtree(start_chunk.0, &data, is_root); if leaf_hash != actual { return Err(DecodeError::LeafHashMismatch(start_chunk)); } @@ -630,7 +651,8 @@ where O: Outboard, W: AsyncStreamWriter, { - encode_ranges_validated_with_key(data, outboard, ranges, encoded, None).await + // Shared impl, standard BLAKE3 hash mode. + encode_ranges_validated_impl(data, outboard, ranges, encoded, Standard).await } /// Encode ranges with BLAKE3 keyed hash validation. @@ -646,15 +668,19 @@ where O: Outboard, W: AsyncStreamWriter, { - encode_ranges_validated_with_key(data, outboard, ranges, encoded, Some(key)).await + // Shared impl, keyed BLAKE3 mode. + encode_ranges_validated_impl(data, outboard, ranges, encoded, Keyed(*key)).await } -async fn encode_ranges_validated_with_key( +/// Shared encode path for standard and keyed APIs. +/// +/// `hash_strategy` selects which [BaoHashing] implementation validates hashes. +async fn encode_ranges_validated_impl( mut data: D, mut outboard: O, ranges: &ChunkRangesRef, encoded: W, - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> result::Result<(), EncodeError> where D: AsyncSliceReader, @@ -679,7 +705,7 @@ where .. } => { let (l_hash, r_hash) = outboard.load(node).await?.unwrap(); - let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, key); + let actual = hash_strategy.parent_cv(&l_hash, &r_hash, is_root); let expected = stack.pop().unwrap(); if actual != expected { return Err(EncodeError::ParentHashMismatch(node)); @@ -720,11 +746,11 @@ where tree.block_size.to_u32(), true, &mut out_buf, - key, + hash_strategy, ); (actual, out_buf.clone().into()) } else { - let actual = hash_subtree_with_key(start_chunk.0, &bytes, is_root, key); + let actual = hash_strategy.hash_subtree(start_chunk.0, &bytes, is_root); (actual, bytes) }; if actual != expected { @@ -755,7 +781,8 @@ where R: AsyncStreamReader, W: AsyncSliceWriter, { - decode_ranges_with_key(encoded, ranges, target, outboard, None).await + // Shared impl, standard BLAKE3 hash mode. + decode_ranges_impl(encoded, ranges, target, outboard, Standard).await } /// Decode a keyed response into a file while updating an outboard. @@ -771,27 +798,32 @@ where R: AsyncStreamReader, W: AsyncSliceWriter, { - decode_ranges_with_key(encoded, ranges, target, outboard, Some(key)).await + // Shared impl, keyed BLAKE3 mode. + decode_ranges_impl(encoded, ranges, target, outboard, Keyed(*key)).await } -async fn decode_ranges_with_key( +/// Shared decode path for standard and keyed APIs. +/// +/// `hash_strategy` selects which [BaoHashing] implementation verifies hashes. +async fn decode_ranges_impl( encoded: R, ranges: ChunkRanges, mut target: W, mut outboard: O, - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> std::result::Result<(), DecodeError> where O: OutboardMut + Outboard, R: AsyncStreamReader, W: AsyncSliceWriter, { - let mut reading = match key { - None => ResponseDecoder::new(outboard.root(), ranges, outboard.tree(), encoded), - Some(key) => { - ResponseDecoder::new_keyed(outboard.root(), ranges, outboard.tree(), encoded, key) - } - }; + let mut reading = ResponseDecoder::with_hash_strategy( + outboard.root(), + ranges, + outboard.tree(), + encoded, + hash_strategy, + ); loop { let item = match reading.next().await { ResponseDecoderNext::Done(_reader) => break, @@ -826,7 +858,8 @@ pub async fn outboard( tree: BaoTree, outboard: impl OutboardMut, ) -> io::Result { - outboard_with_key(data, tree, outboard, None).await + // Shared impl, standard BLAKE3 hash mode. + outboard_with_hash_strategy(data, tree, outboard, Standard).await } /// Compute the keyed outboard for the given data. @@ -836,27 +869,30 @@ pub async fn keyed_outboard( outboard: impl OutboardMut, key: &[u8; 32], ) -> io::Result { - outboard_with_key(data, tree, outboard, Some(key)).await + // Shared impl, keyed BLAKE3 mode. + outboard_with_hash_strategy(data, tree, outboard, Keyed(*key)).await } -async fn outboard_with_key( +/// Allocates a chunk group buffer and delegates to [outboard_impl]. +async fn outboard_with_hash_strategy( data: impl AsyncStreamReader, tree: BaoTree, mut outboard: impl OutboardMut, - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - let hash = outboard_impl(tree, data, &mut outboard, &mut buffer, key).await?; - Ok(hash) + outboard_impl(tree, data, &mut outboard, &mut buffer, hash_strategy).await } -/// Internal helper for [outboard_post_order]. This takes a buffer of the chunk group size. -async fn outboard_impl( +/// Shared outboard traversal for standard and keyed APIs. +/// +/// `hash_strategy` selects which [BaoHashing] implementation computes hashes. +async fn outboard_impl( tree: BaoTree, mut data: impl AsyncStreamReader, mut outboard: impl OutboardMut, buffer: &mut [u8], - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -867,7 +903,7 @@ async fn outboard_impl( let right_hash = stack.pop().unwrap(); let left_hash = stack.pop().unwrap(); outboard.save(node, &(left_hash, right_hash)).await?; - let parent = parent_cv_with_key(&left_hash, &right_hash, is_root, key); + let parent = hash_strategy.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -877,7 +913,7 @@ async fn outboard_impl( .. } => { let buf = data.read_bytes_exact(size).await?; - let hash = hash_subtree_with_key(start_chunk.0, &buf, is_root, key); + let hash = hash_strategy.hash_subtree(start_chunk.0, &buf, is_root); stack.push(hash); } } @@ -898,7 +934,8 @@ pub async fn outboard_post_order( tree: BaoTree, outboard: impl AsyncStreamWriter, ) -> io::Result { - outboard_post_order_with_key(data, tree, outboard, None).await + // Shared impl, standard BLAKE3 hash mode. + outboard_post_order_with_hash_strategy(data, tree, outboard, Standard).await } /// Compute the keyed post order outboard for the given data. @@ -908,27 +945,30 @@ pub async fn keyed_outboard_post_order( outboard: impl AsyncStreamWriter, key: &[u8; 32], ) -> io::Result { - outboard_post_order_with_key(data, tree, outboard, Some(key)).await + // Shared impl, keyed BLAKE3 mode. + outboard_post_order_with_hash_strategy(data, tree, outboard, Keyed(*key)).await } -async fn outboard_post_order_with_key( +/// Allocates a chunk group buffer and delegates to [outboard_post_order_impl]. +async fn outboard_post_order_with_hash_strategy( data: impl AsyncStreamReader, tree: BaoTree, mut outboard: impl AsyncStreamWriter, - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - let hash = outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, key).await?; - Ok(hash) + outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, hash_strategy).await } -/// Internal helper for [outboard_post_order]. This takes a buffer of the chunk group size. -async fn outboard_post_order_impl( +/// Shared post order outboard traversal for standard and keyed APIs. +/// +/// `hash_strategy` selects which [BaoHashing] implementation computes hashes. +async fn outboard_post_order_impl( tree: BaoTree, mut data: impl AsyncStreamReader, mut outboard: impl AsyncStreamWriter, buffer: &mut [u8], - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -940,7 +980,7 @@ async fn outboard_post_order_impl( let left_hash = stack.pop().unwrap(); outboard.write(left_hash.as_bytes()).await?; outboard.write(right_hash.as_bytes()).await?; - let parent = parent_cv_with_key(&left_hash, &right_hash, is_root, key); + let parent = hash_strategy.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -950,7 +990,7 @@ async fn outboard_post_order_impl( .. } => { let buf = data.read_bytes_exact(size).await?; - let hash = hash_subtree_with_key(start_chunk.0, &buf, is_root, key); + let hash = hash_strategy.hash_subtree(start_chunk.0, &buf, is_root); stack.push(hash); } } @@ -984,8 +1024,8 @@ mod validate { use super::Outboard; use crate::{ - blake3, hash_subtree_with_key, io::LocalBoxFuture, parent_cv_with_key, - rec::truncate_ranges, split, BaoTree, ChunkNum, ChunkRangesRef, TreeNode, + blake3, io::LocalBoxFuture, rec::truncate_ranges, split, BaoHashing, BaoTree, ChunkNum, + ChunkRangesRef, Keyed, Standard, TreeNode, }; /// Given a data file and an outboard, compute all valid ranges. @@ -1002,7 +1042,8 @@ mod validate { O: Outboard + 'a, D: AsyncSliceReader + 'a, { - valid_ranges_with_key(outboard, data, ranges, None) + // Shared impl, standard BLAKE3 hash mode. + valid_ranges_impl(outboard, data, ranges, Standard) } /// Given a data file and a keyed outboard, compute all valid ranges. @@ -1016,14 +1057,16 @@ mod validate { O: Outboard + 'a, D: AsyncSliceReader + 'a, { - valid_ranges_with_key(outboard, data, ranges, Some(key)) + // Shared impl, keyed BLAKE3 mode. + valid_ranges_impl(outboard, data, ranges, Keyed(*key)) } - fn valid_ranges_with_key<'a, O, D>( + /// Shared validation path for standard and keyed APIs. + fn valid_ranges_impl<'a, O, D, H: BaoHashing + Copy + 'a>( outboard: O, data: D, ranges: &'a ChunkRangesRef, - key: Option<&'a [u8; 32]>, + hash_strategy: H, ) -> impl Stream>> + 'a where O: Outboard + 'a, @@ -1031,29 +1074,29 @@ mod validate { { Gen::new(move |co| async move { if let Err(cause) = - RecursiveDataValidator::validate(outboard, data, ranges, &co, key).await + RecursiveDataValidator::validate(outboard, data, ranges, &co, hash_strategy).await { co.yield_(Err(cause)).await; } }) } - struct RecursiveDataValidator<'a, O: Outboard, D: AsyncSliceReader> { + struct RecursiveDataValidator<'a, O: Outboard, D: AsyncSliceReader, H: BaoHashing + Copy> { tree: BaoTree, shifted_filled_size: TreeNode, outboard: O, data: D, co: &'a Co>>, - key: Option<&'a [u8; 32]>, + hash_strategy: H, } - impl RecursiveDataValidator<'_, O, D> { + impl RecursiveDataValidator<'_, O, D, H> { async fn validate( outboard: O, data: D, ranges: &ChunkRangesRef, co: &Co>>, - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> io::Result<()> { let tree = outboard.tree(); if tree.blocks() == 1 { @@ -1062,7 +1105,7 @@ mod validate { let data = data .read_exact_at(0, tree.size().try_into().unwrap()) .await?; - let actual = hash_subtree_with_key(0, &data, true, key); + let actual = hash_strategy.hash_subtree(0, &data, true); if actual == outboard.root() { co.yield_(Ok(ChunkNum(0)..tree.chunks())).await; } @@ -1077,7 +1120,7 @@ mod validate { outboard, data, co, - key, + hash_strategy, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -1093,11 +1136,10 @@ mod validate { let len = (range.end - range.start).try_into().unwrap(); let data = self.data.read_exact_at(range.start, len).await?; // is_root is always false because the case of a single chunk group is handled before calling this function - let actual = hash_subtree_with_key( + let actual = self.hash_strategy.hash_subtree( ChunkNum::full_chunks(range.start).0, &data, is_root, - self.key, ); if &actual == hash { // yield the left range @@ -1132,7 +1174,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, self.key); + let actual = self.hash_strategy.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -1235,7 +1277,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, None); + let actual = Standard.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); diff --git a/src/io/sync.rs b/src/io/sync.rs index 853bcfe..697ee53 100644 --- a/src/io/sync.rs +++ b/src/io/sync.rs @@ -14,16 +14,15 @@ use smallvec::SmallVec; use super::{combine_hash_pair, BaoContentItem, DecodeError}; pub use crate::rec::truncate_ranges; use crate::{ - blake3, hash_subtree_with_key, + blake3, io::{ error::EncodeError, outboard::{parse_hash_pair, PostOrderOutboard, PreOrderOutboard}, Leaf, Parent, }, iter::{BaoChunk, ResponseIterRef}, - parent_cv_with_key, rec::encode_selected_rec, - BaoTree, BlockSize, ChunkRangesRef, TreeNode, + BaoHashing, BaoTree, BlockSize, ChunkRangesRef, Keyed, Standard, TreeNode, }; /// A binary merkle tree for blake3 hashes of a blob. @@ -338,14 +337,18 @@ impl Outboard for PostOrderOutboard { } } -/// Iterator that can be used to decode a response to a range request +/// Iterator that can be used to decode a response to a range request. +/// +/// Generic over `H` so keyed and standard decoders share this implementation. +/// Defaults to [Standard]. Use [Self::new_keyed] for keyed responses. #[derive(Debug)] -pub struct DecodeResponseIter<'a, R> { +pub struct DecodeResponseIter<'a, R, H: BaoHashing + Copy = Standard> { inner: ResponseIterRef<'a>, stack: SmallVec<[blake3::Hash; 10]>, encoded: R, buf: BytesMut, - key: Option<[u8; 32]>, + /// Compile time hashing strategy, either [Standard] or [Keyed]. + hash_strategy: H, } impl<'a, R: Read> DecodeResponseIter<'a, R> { @@ -369,16 +372,7 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { ranges: &'a ChunkRangesRef, buf: BytesMut, ) -> Self { - let ranges = truncate_ranges(ranges, tree.size()); - let mut stack = SmallVec::new(); - stack.push(root); - Self { - stack, - inner: ResponseIterRef::new(tree, ranges), - encoded, - buf, - key: None, - } + DecodeResponseIter::with_hash_strategy(root, tree, encoded, ranges, buf, Standard) } /// Create a new iterator to decode a keyed response. @@ -388,10 +382,32 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { encoded: R, ranges: &'a ChunkRangesRef, key: &[u8; 32], + ) -> DecodeResponseIter<'a, R, Keyed> { + let buf = BytesMut::with_capacity(tree.block_size().bytes()); + DecodeResponseIter::with_hash_strategy(root, tree, encoded, ranges, buf, Keyed(*key)) + } +} + +impl<'a, R: Read, H: BaoHashing + Copy> DecodeResponseIter<'a, R, H> { + /// Shared constructor used by [Self::new] and [Self::new_keyed]. + pub(crate) fn with_hash_strategy( + root: blake3::Hash, + tree: BaoTree, + encoded: R, + ranges: &'a ChunkRangesRef, + buf: BytesMut, + hash_strategy: H, ) -> Self { - let mut res = Self::new(root, tree, encoded, ranges); - res.key = Some(*key); - res + let ranges = truncate_ranges(ranges, tree.size()); + let mut stack = SmallVec::new(); + stack.push(root); + Self { + stack, + inner: ResponseIterRef::new(tree, ranges), + encoded, + buf, + hash_strategy, + } } /// Get a reference to the buffer used for decoding. @@ -418,7 +434,7 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { let pair @ (l_hash, r_hash) = read_parent(&mut self.encoded) .map_err(|e| DecodeError::maybe_parent_not_found(e, node))?; let parent_hash = self.stack.pop().unwrap(); - let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, self.key.as_ref()); + let actual = self.hash_strategy.parent_cv(&l_hash, &r_hash, is_root); if parent_hash != actual { return Err(DecodeError::ParentHashMismatch(node)); } @@ -440,8 +456,9 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { self.encoded .read_exact(&mut self.buf) .map_err(|e| DecodeError::maybe_leaf_not_found(e, start_chunk))?; - let actual = - hash_subtree_with_key(start_chunk.0, &self.buf, is_root, self.key.as_ref()); + let actual = self + .hash_strategy + .hash_subtree(start_chunk.0, &self.buf, is_root); let leaf_hash = self.stack.pop().unwrap(); if leaf_hash != actual { return Err(DecodeError::LeafHashMismatch(start_chunk)); @@ -459,7 +476,7 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { } } -impl Iterator for DecodeResponseIter<'_, R> { +impl Iterator for DecodeResponseIter<'_, R, H> { type Item = result::Result; fn next(&mut self) -> Option { @@ -517,7 +534,8 @@ pub fn encode_ranges_validated( ranges: &ChunkRangesRef, encoded: W, ) -> result::Result<(), EncodeError> { - encode_ranges_validated_with_key(data, outboard, ranges, encoded, None) + // Shared impl, standard BLAKE3 hash mode. + encode_ranges_validated_impl(data, outboard, ranges, encoded, Standard) } /// Encode ranges with BLAKE3 keyed hash validation. @@ -528,15 +546,19 @@ pub fn keyed_encode_ranges_validated( encoded: W, key: &[u8; 32], ) -> result::Result<(), EncodeError> { - encode_ranges_validated_with_key(data, outboard, ranges, encoded, Some(key)) + // Shared impl, keyed BLAKE3 mode. + encode_ranges_validated_impl(data, outboard, ranges, encoded, Keyed(*key)) } -fn encode_ranges_validated_with_key( +/// Shared encode path for standard and keyed APIs. +/// +/// `hash_strategy` selects which [BaoHashing] implementation validates hashes. +fn encode_ranges_validated_impl( data: D, outboard: O, ranges: &ChunkRangesRef, encoded: W, - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> result::Result<(), EncodeError> { if ranges.is_empty() { return Ok(()); @@ -560,7 +582,7 @@ fn encode_ranges_validated_with_key( .. } => { let (l_hash, r_hash) = outboard.load(node)?.unwrap(); - let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, key); + let actual = hash_strategy.parent_cv(&l_hash, &r_hash, is_root); let expected = stack.pop().unwrap(); if actual != expected { return Err(EncodeError::ParentHashMismatch(node)); @@ -599,11 +621,11 @@ fn encode_ranges_validated_with_key( tree.block_size.to_u32(), true, &mut out_buf, - key, + hash_strategy, ); (actual, &out_buf[..]) } else { - let actual = hash_subtree_with_key(start_chunk.0, buf, is_root, key); + let actual = hash_strategy.hash_subtree(start_chunk.0, buf, is_root); #[allow(clippy::redundant_slicing)] (actual, &buf[..]) }; @@ -632,7 +654,8 @@ where R: Read, W: WriteAt, { - decode_ranges_with_key(encoded, ranges, target, outboard, None) + // Shared impl, standard BLAKE3 hash mode. + decode_ranges_impl(encoded, ranges, target, outboard, Standard) } /// Decode a keyed response into a file while updating an outboard. @@ -648,27 +671,33 @@ where R: Read, W: WriteAt, { - decode_ranges_with_key(encoded, ranges, target, outboard, Some(key)) + // Shared impl, keyed BLAKE3 mode. + decode_ranges_impl(encoded, ranges, target, outboard, Keyed(*key)) } -fn decode_ranges_with_key( +/// Shared decode path for standard and keyed APIs. +/// +/// `hash_strategy` selects which [BaoHashing] implementation verifies hashes. +fn decode_ranges_impl( encoded: R, ranges: &ChunkRangesRef, mut target: W, mut outboard: O, - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> std::result::Result<(), DecodeError> where O: OutboardMut + Outboard, R: Read, W: WriteAt, { - let iter = match key { - None => DecodeResponseIter::new(outboard.root(), outboard.tree(), encoded, ranges), - Some(key) => { - DecodeResponseIter::new_keyed(outboard.root(), outboard.tree(), encoded, ranges, key) - } - }; + let iter = DecodeResponseIter::with_hash_strategy( + outboard.root(), + outboard.tree(), + encoded, + ranges, + BytesMut::with_capacity(outboard.tree().block_size().bytes()), + hash_strategy, + ); for item in iter { match item? { BaoContentItem::Parent(Parent { node, pair }) => { @@ -691,7 +720,8 @@ pub fn outboard( tree: BaoTree, outboard: impl OutboardMut, ) -> io::Result { - outboard_with_key(data, tree, outboard, None) + // Shared impl, standard BLAKE3 hash mode. + outboard_with_hash_strategy(data, tree, outboard, Standard) } /// Compute the keyed outboard for the given data. @@ -701,27 +731,30 @@ pub fn keyed_outboard( outboard: impl OutboardMut, key: &[u8; 32], ) -> io::Result { - outboard_with_key(data, tree, outboard, Some(key)) + // Shared impl, keyed BLAKE3 mode. + outboard_with_hash_strategy(data, tree, outboard, Keyed(*key)) } -fn outboard_with_key( +/// Allocates a chunk group buffer and delegates to [outboard_impl]. +fn outboard_with_hash_strategy( data: impl Read, tree: BaoTree, mut outboard: impl OutboardMut, - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - let hash = outboard_impl(tree, data, &mut outboard, &mut buffer, key)?; - Ok(hash) + outboard_impl(tree, data, &mut outboard, &mut buffer, hash_strategy) } -/// Internal helper for [outboard_post_order]. This takes a buffer of the chunk group size. -fn outboard_impl( +/// Shared outboard traversal for standard and keyed APIs. +/// +/// `hash_strategy` selects which [BaoHashing] implementation computes hashes. +fn outboard_impl( tree: BaoTree, mut data: impl Read, mut outboard: impl OutboardMut, buffer: &mut [u8], - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -732,7 +765,7 @@ fn outboard_impl( let right_hash = stack.pop().unwrap(); let left_hash = stack.pop().unwrap(); outboard.save(node, &(left_hash, right_hash))?; - let parent = parent_cv_with_key(&left_hash, &right_hash, is_root, key); + let parent = hash_strategy.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -743,7 +776,7 @@ fn outboard_impl( } => { let buf = &mut buffer[..size]; data.read_exact(buf)?; - let hash = hash_subtree_with_key(start_chunk.0, buf, is_root, key); + let hash = hash_strategy.hash_subtree(start_chunk.0, buf, is_root); stack.push(hash); } } @@ -764,7 +797,8 @@ pub fn outboard_post_order( tree: BaoTree, outboard: impl Write, ) -> io::Result { - outboard_post_order_with_key(data, tree, outboard, None) + // Shared impl, standard BLAKE3 hash mode. + outboard_post_order_with_hash_strategy(data, tree, outboard, Standard) } /// Compute the keyed post order outboard for the given data. @@ -774,27 +808,30 @@ pub fn keyed_outboard_post_order( outboard: impl Write, key: &[u8; 32], ) -> io::Result { - outboard_post_order_with_key(data, tree, outboard, Some(key)) + // Shared impl, keyed BLAKE3 mode. + outboard_post_order_with_hash_strategy(data, tree, outboard, Keyed(*key)) } -fn outboard_post_order_with_key( +/// Allocates a chunk group buffer and delegates to [outboard_post_order_impl]. +fn outboard_post_order_with_hash_strategy( data: impl Read, tree: BaoTree, mut outboard: impl Write, - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - let hash = outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, key)?; - Ok(hash) + outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, hash_strategy) } -/// Internal helper for [outboard_post_order]. This takes a buffer of the chunk group size. -fn outboard_post_order_impl( +/// Shared post order outboard traversal for standard and keyed APIs. +/// +/// `hash_strategy` selects which [BaoHashing] implementation computes hashes. +fn outboard_post_order_impl( tree: BaoTree, mut data: impl Read, mut outboard: impl Write, buffer: &mut [u8], - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -806,7 +843,7 @@ fn outboard_post_order_impl( let left_hash = stack.pop().unwrap(); outboard.write_all(left_hash.as_bytes())?; outboard.write_all(right_hash.as_bytes())?; - let parent = parent_cv_with_key(&left_hash, &right_hash, is_root, key); + let parent = hash_strategy.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -817,7 +854,7 @@ fn outboard_post_order_impl( } => { let buf = &mut buffer[..size]; data.read_exact(buf)?; - let hash = hash_subtree_with_key(start_chunk.0, buf, is_root, key); + let hash = hash_strategy.hash_subtree(start_chunk.0, buf, is_root); stack.push(hash); } } @@ -858,8 +895,8 @@ mod validate { use super::Outboard; use crate::{ - blake3, hash_subtree_with_key, io::LocalBoxFuture, parent_cv_with_key, - rec::truncate_ranges, split, BaoTree, ChunkNum, ChunkRangesRef, TreeNode, + blake3, io::LocalBoxFuture, rec::truncate_ranges, split, BaoHashing, BaoTree, ChunkNum, + ChunkRangesRef, Keyed, Standard, TreeNode, }; /// Given a data file and an outboard, compute all valid ranges. @@ -876,7 +913,8 @@ mod validate { O: Outboard + 'a, D: ReadAt + 'a, { - valid_ranges_with_key(outboard, data, ranges, None) + // Shared impl, standard BLAKE3 hash mode. + valid_ranges_impl(outboard, data, ranges, Standard) } /// Given a data file and a keyed outboard, compute all valid ranges. @@ -890,14 +928,16 @@ mod validate { O: Outboard + 'a, D: ReadAt + 'a, { - valid_ranges_with_key(outboard, data, ranges, Some(key)) + // Shared impl, keyed BLAKE3 mode. + valid_ranges_impl(outboard, data, ranges, Keyed(*key)) } - fn valid_ranges_with_key<'a, O, D>( + /// Shared validation path for standard and keyed APIs. + fn valid_ranges_impl<'a, O, D, H: BaoHashing + Copy + 'a>( outboard: O, data: D, ranges: &'a ChunkRangesRef, - key: Option<&'a [u8; 32]>, + hash_strategy: H, ) -> impl IntoIterator>> + 'a where O: Outboard + 'a, @@ -905,30 +945,30 @@ mod validate { { Gen::new(move |co| async move { if let Err(cause) = - RecursiveDataValidator::validate(outboard, data, ranges, &co, key).await + RecursiveDataValidator::validate(outboard, data, ranges, &co, hash_strategy).await { co.yield_(Err(cause)).await; } }) } - struct RecursiveDataValidator<'a, O: Outboard, D: ReadAt> { + struct RecursiveDataValidator<'a, O: Outboard, D: ReadAt, H: BaoHashing + Copy> { tree: BaoTree, shifted_filled_size: TreeNode, outboard: O, data: D, buffer: Vec, co: &'a Co>>, - key: Option<&'a [u8; 32]>, + hash_strategy: H, } - impl RecursiveDataValidator<'_, O, D> { + impl RecursiveDataValidator<'_, O, D, H> { async fn validate( outboard: O, data: D, ranges: &ChunkRangesRef, co: &Co>>, - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> io::Result<()> { let tree = outboard.tree(); let mut buffer = vec![0u8; tree.chunk_group_bytes()]; @@ -936,7 +976,7 @@ mod validate { // special case for a tree that fits in one block / chunk group let tmp = &mut buffer[..tree.size().try_into().unwrap()]; data.read_exact_at(0, tmp)?; - let actual = hash_subtree_with_key(0, tmp, true, key); + let actual = hash_strategy.hash_subtree(0, tmp, true); if actual == outboard.root() { co.yield_(Ok(ChunkNum(0)..tree.chunks())).await; } @@ -952,7 +992,7 @@ mod validate { data, buffer, co, - key, + hash_strategy, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -970,7 +1010,8 @@ mod validate { self.data.read_exact_at(range.start, tmp)?; // is_root is always false because the case of a single chunk group is handled before calling this function let actual = - hash_subtree_with_key(ChunkNum::full_chunks(range.start).0, tmp, is_root, self.key); + self.hash_strategy + .hash_subtree(ChunkNum::full_chunks(range.start).0, tmp, is_root); if &actual == hash { // yield the left range self.co @@ -1004,7 +1045,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, self.key); + let actual = self.hash_strategy.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -1106,7 +1147,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = parent_cv_with_key(&l_hash, &r_hash, is_root, None); + let actual = Standard.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); diff --git a/src/lib.rs b/src/lib.rs index 25592b3..4d79779 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -239,8 +239,83 @@ pub type ByteRanges = range_collections::RangeSet2; /// [ChunkRanges] implements [`AsRef`]. pub type ChunkRangesRef = range_collections::RangeSetRef; +/// Hashing strategy for shared encode and decode paths. +/// +/// Internal IO is generic over this trait so standard and keyed modes share one +/// implementation without runtime branches or duplicated bodies. +/// +/// Pass [Standard] for BLAKE3 hash mode or [Keyed] for BLAKE3 keyed mode. The +/// strategy is a compile time value, so each call site monomorphizes with zero +/// extra cost. +/// +/// Public entry points like `encode_ranges_validated` and +/// `keyed_encode_ranges_validated` pick the strategy at the API boundary and +/// delegate to a single shared function. +pub trait BaoHashing: Copy { + /// Hash a subtree of chunk data. + fn hash_subtree(&self, start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash; + /// Combine two child chaining values into a parent chaining value. + fn parent_cv( + &self, + left_child: &blake3::Hash, + right_child: &blake3::Hash, + is_root: bool, + ) -> blake3::Hash; +} + +/// BLAKE3 hash mode. Default [BaoHashing] strategy for unkeyed APIs. +/// +/// Routes through the crate internal [hash_subtree] and [parent_cv] helpers so +/// validate only builds still link those symbols. +#[derive(Clone, Copy)] +pub struct Standard; + +impl BaoHashing for Standard { + fn hash_subtree(&self, start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash { + hash_subtree(start_chunk, data, is_root) + } + + fn parent_cv( + &self, + left_child: &blake3::Hash, + right_child: &blake3::Hash, + is_root: bool, + ) -> blake3::Hash { + parent_cv(left_child, right_child, is_root) + } +} + +/// BLAKE3 keyed mode. Wraps a 32 byte key for domain separated hashing. +#[derive(Clone, Copy)] +pub struct Keyed(pub [u8; 32]); + +impl BaoHashing for Keyed { + fn hash_subtree(&self, start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash { + keyed_hash_subtree(start_chunk, data, is_root, &self.0) + } + + fn parent_cv( + &self, + left_child: &blake3::Hash, + right_child: &blake3::Hash, + is_root: bool, + ) -> blake3::Hash { + keyed_parent_cv(left_child, right_child, is_root, &self.0) + } +} + pub(crate) fn hash_subtree(start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash { - hash_subtree_with_key(start_chunk, data, is_root, None) + use blake3::hazmat::{ChainingValue, HasherExt}; + if is_root { + debug_assert!(start_chunk == 0); + blake3::hash(data) + } else { + let mut hasher = blake3::Hasher::new(); + hasher.set_input_offset(start_chunk * 1024); + hasher.update(data); + let non_root_hash: ChainingValue = hasher.finalize_non_root(); + blake3::Hash::from(non_root_hash) + } } /// Compute the hash of a subtree using BLAKE3 keyed mode. @@ -251,28 +326,13 @@ pub fn keyed_hash_subtree( data: &[u8], is_root: bool, key: &[u8; 32], -) -> blake3::Hash { - hash_subtree_with_key(start_chunk, data, is_root, Some(key)) -} - -pub(crate) fn hash_subtree_with_key( - start_chunk: u64, - data: &[u8], - is_root: bool, - key: Option<&[u8; 32]>, ) -> blake3::Hash { use blake3::hazmat::{ChainingValue, HasherExt}; if is_root { debug_assert!(start_chunk == 0); - match key { - None => blake3::hash(data), - Some(key) => blake3::keyed_hash(key, data), - } + blake3::keyed_hash(key, data) } else { - let mut hasher = match key { - None => blake3::Hasher::new(), - Some(key) => blake3::Hasher::new_keyed(key), - }; + let mut hasher = blake3::Hasher::new_keyed(key); hasher.set_input_offset(start_chunk * 1024); hasher.update(data); let non_root_hash: ChainingValue = hasher.finalize_non_root(); @@ -285,7 +345,18 @@ pub(crate) fn parent_cv( right_child: &blake3::Hash, is_root: bool, ) -> blake3::Hash { - parent_cv_with_key(left_child, right_child, is_root, None) + use blake3::hazmat::{merge_subtrees_non_root, merge_subtrees_root, ChainingValue, Mode}; + let left_child: ChainingValue = *left_child.as_bytes(); + let right_child: ChainingValue = *right_child.as_bytes(); + if is_root { + merge_subtrees_root(&left_child, &right_child, Mode::Hash) + } else { + blake3::Hash::from(merge_subtrees_non_root( + &left_child, + &right_child, + Mode::Hash, + )) + } } /// Merge two child subtree hashes using BLAKE3 keyed mode. @@ -294,23 +365,11 @@ pub fn keyed_parent_cv( right_child: &blake3::Hash, is_root: bool, key: &[u8; 32], -) -> blake3::Hash { - parent_cv_with_key(left_child, right_child, is_root, Some(key)) -} - -pub(crate) fn parent_cv_with_key( - left_child: &blake3::Hash, - right_child: &blake3::Hash, - is_root: bool, - key: Option<&[u8; 32]>, ) -> blake3::Hash { use blake3::hazmat::{merge_subtrees_non_root, merge_subtrees_root, ChainingValue, Mode}; let left_child: ChainingValue = *left_child.as_bytes(); let right_child: ChainingValue = *right_child.as_bytes(); - let mode = match key { - None => Mode::Hash, - Some(key) => Mode::KeyedHash(key), - }; + let mode = Mode::KeyedHash(key); if is_root { merge_subtrees_root(&left_child, &right_child, mode) } else { diff --git a/src/rec.rs b/src/rec.rs index 65830c1..ec6c6ac 100644 --- a/src/rec.rs +++ b/src/rec.rs @@ -96,8 +96,12 @@ fn truncated_len(ranges: &ChunkRangesRef, size: u64) -> usize { /// This is used as a reference implementation in tests, but also to compute hashes /// below the chunk group size when creating responses for outboards with a chunk group /// size of >0. -#[allow(clippy::too_many_arguments)] // keyed mode adds `key`; splitting into a struct isn't worth it here -pub(crate) fn encode_selected_rec( +/// Recursive reference encoder shared by standard and keyed paths. +/// +/// `hash_strategy` selects which [crate::BaoHashing] implementation to use when +/// computing subtree and parent hashes. +#[allow(clippy::too_many_arguments)] // keyed mode adds `hash_strategy`; splitting into a struct isn't worth it here +pub(crate) fn encode_selected_rec( start_chunk: ChunkNum, data: &[u8], is_root: bool, @@ -105,14 +109,14 @@ pub(crate) fn encode_selected_rec( min_level: u32, emit_data: bool, res: &mut Vec, - key: Option<&[u8; 32]>, + hash_strategy: H, ) -> blake3::Hash { use blake3::CHUNK_LEN; if data.len() <= CHUNK_LEN { if emit_data && !query.is_empty() { res.extend_from_slice(data); } - crate::hash_subtree_with_key(start_chunk.0, data, is_root, key) + hash_strategy.hash_subtree(start_chunk.0, data, is_root) } else { let chunks = data.len() / CHUNK_LEN + (data.len() % CHUNK_LEN != 0) as usize; let chunks = chunks.next_power_of_two(); @@ -144,7 +148,7 @@ pub(crate) fn encode_selected_rec( min_level, emit_data, res, - key, + hash_strategy, ); let right = encode_selected_rec( mid_chunk, @@ -154,14 +158,14 @@ pub(crate) fn encode_selected_rec( min_level, emit_data, res, - key, + hash_strategy, ); // backfill the hashes if needed if let Some(o) = hash_offset { res[o..o + 32].copy_from_slice(left.as_bytes()); res[o + 32..o + 64].copy_from_slice(right.as_bytes()); } - crate::parent_cv_with_key(&left, &right, is_root, key) + hash_strategy.parent_cv(&left, &right, is_root) } } @@ -279,7 +283,7 @@ mod test_support { 0, false, &mut res, - None, + crate::Standard, ); (res, hash) } @@ -295,7 +299,7 @@ mod test_support { 0, true, &mut res, - None, + crate::Standard, ); (res, hash) } @@ -436,7 +440,7 @@ mod test_support { block_size.to_u32(), true, &mut res, - None, + crate::Standard, ); (res, hash) } diff --git a/src/tests.rs b/src/tests.rs index aba318d..08faab9 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -33,7 +33,7 @@ use crate::{ keyed_init_from_keyed_checks, keyed_outboard_functions_checks, make_test_data, range_union, truncate_ranges, ReferencePreOrderPartialChunkIterRef, }, - split, ChunkRanges, ChunkRangesRef, ResponseIter, + split, ChunkRanges, ChunkRangesRef, Keyed, ResponseIter, Standard, }; #[cfg(feature = "tokio_fsm")] @@ -42,6 +42,7 @@ use crate::rec::{ keyed_outboard_functions_checks_fsm, }; +/// Reference encoder using the [Keyed] hashing strategy. fn keyed_encode_selected_reference( data: &[u8], block_size: BlockSize, @@ -59,7 +60,7 @@ fn keyed_encode_selected_reference( max_skip_level, true, &mut res, - Some(key), + Keyed(*key), ); (hash, res) } @@ -1044,7 +1045,7 @@ fn encode_selected_rec_cases() { min_level, true, &mut actual_encoded, - None, + Standard, ); actual_encoded.len() - data.len() }; @@ -1070,7 +1071,7 @@ fn encode_selected_reference( max_skip_level, true, &mut res, - None, + Standard, ); (hash, res) } diff --git a/src/tests2.rs b/src/tests2.rs index 01d4c3f..616d9b8 100644 --- a/src/tests2.rs +++ b/src/tests2.rs @@ -32,7 +32,7 @@ use crate::{ partial_chunk_iter_reference, range_union, response_iter_reference, select_nodes_rec, truncate_ranges, ReferencePreOrderPartialChunkIterRef, }, - BaoTree, BlockSize, ChunkNum, ChunkRanges, ChunkRangesRef, TreeNode, + BaoTree, BlockSize, ChunkNum, ChunkRanges, ChunkRangesRef, Standard, TreeNode, }; fn keyed_test_key(context: &[u8]) -> [u8; 32] { @@ -1298,7 +1298,9 @@ fn selection_reference_comparison_proptest( } } -/// Reference implementation of encode_ranges_validated that uses the simple recursive impl +/// Reference implementation of encode_ranges_validated that uses the simple recursive impl. +/// +/// Uses the [Standard] hashing strategy for unkeyed BLAKE3 mode. fn encode_selected_reference( data: &[u8], block_size: BlockSize, @@ -1315,7 +1317,7 @@ fn encode_selected_reference( max_skip_level, true, &mut res, - None, + Standard, ); (hash, res) } From 81974e2f32c5f6825c9a271c042954e73b602cc5 Mon Sep 17 00:00:00 2001 From: Hunter Beast Date: Mon, 6 Jul 2026 11:52:08 -0600 Subject: [PATCH 03/12] Polish BaoHashing refactor for upstream review Trim comments to bao-tree style, hide strategy types from rustdoc, and add Keyed decode type aliases. Public keyed APIs unchanged. --- src/io/fsm.rs | 50 +++++++++++++++++++------------------------------- src/io/sync.rs | 40 ++++++++++++---------------------------- src/lib.rs | 23 +++++++---------------- src/rec.rs | 4 +--- 4 files changed, 39 insertions(+), 78 deletions(-) diff --git a/src/io/fsm.rs b/src/io/fsm.rs index 1e5dba3..21fc1a0 100644 --- a/src/io/fsm.rs +++ b/src/io/fsm.rs @@ -401,18 +401,15 @@ pub(crate) fn parse_hash_pair(buf: Bytes) -> io::Result<(blake3::Hash, blake3::H Ok((l_hash, r_hash)) } -/// Generic over `H` so keyed and standard decoders share this implementation. #[derive(Debug)] struct ResponseDecoderInner { iter: ResponseIter, stack: SmallVec<[blake3::Hash; 10]>, encoded: R, - /// Compile time hashing strategy, either [Standard] or [Keyed]. hash_strategy: H, } impl ResponseDecoderInner { - /// Shared constructor used by [ResponseDecoder::new] and [ResponseDecoder::new_keyed]. fn with_hash_strategy( tree: BaoTree, hash: blake3::Hash, @@ -435,12 +432,16 @@ impl ResponseDecoderInner { /// Response decoder. /// -/// Generic over `H` so keyed and standard decoders share this implementation. -/// Defaults to [Standard]. Use [Self::new_keyed] for keyed responses. +/// Keyed callers should use [KeyedResponseDecoder] via [Self::new_keyed]. #[derive(Debug)] pub struct ResponseDecoder(Box>); -/// Next type for ResponseDecoder. +/// Keyed response decoder. +/// +/// See [ResponseDecoder::new_keyed]. +pub type KeyedResponseDecoder = ResponseDecoder; + +/// Next type for [ResponseDecoder]. #[derive(Debug)] pub enum ResponseDecoderNext { /// One more item, and you get back the state machine in the next state @@ -454,6 +455,11 @@ pub enum ResponseDecoderNext { Done(R), } +/// Next type for [KeyedResponseDecoder]. +/// +/// See [ResponseDecoder::new_keyed]. +pub type KeyedResponseDecoderNext = ResponseDecoderNext; + impl ResponseDecoder { /// Create a new response decoder state machine, when you have already read the size. /// @@ -471,7 +477,7 @@ impl ResponseDecoder { tree: BaoTree, encoded: R, key: &[u8; 32], - ) -> ResponseDecoder { + ) -> KeyedResponseDecoder { ResponseDecoder(Box::new(ResponseDecoderInner::with_hash_strategy( tree, hash, @@ -483,7 +489,6 @@ impl ResponseDecoder { } impl ResponseDecoder { - /// Shared constructor used by decode helpers and public new methods. pub(crate) fn with_hash_strategy( hash: blake3::Hash, ranges: ChunkRanges, @@ -651,7 +656,6 @@ where O: Outboard, W: AsyncStreamWriter, { - // Shared impl, standard BLAKE3 hash mode. encode_ranges_validated_impl(data, outboard, ranges, encoded, Standard).await } @@ -668,13 +672,10 @@ where O: Outboard, W: AsyncStreamWriter, { - // Shared impl, keyed BLAKE3 mode. encode_ranges_validated_impl(data, outboard, ranges, encoded, Keyed(*key)).await } -/// Shared encode path for standard and keyed APIs. -/// -/// `hash_strategy` selects which [BaoHashing] implementation validates hashes. +/// Generic encode body monomorphized over the compile time hashing strategy. async fn encode_ranges_validated_impl( mut data: D, mut outboard: O, @@ -781,7 +782,6 @@ where R: AsyncStreamReader, W: AsyncSliceWriter, { - // Shared impl, standard BLAKE3 hash mode. decode_ranges_impl(encoded, ranges, target, outboard, Standard).await } @@ -798,13 +798,10 @@ where R: AsyncStreamReader, W: AsyncSliceWriter, { - // Shared impl, keyed BLAKE3 mode. decode_ranges_impl(encoded, ranges, target, outboard, Keyed(*key)).await } -/// Shared decode path for standard and keyed APIs. -/// -/// `hash_strategy` selects which [BaoHashing] implementation verifies hashes. +/// Generic decode body monomorphized over the compile time hashing strategy. async fn decode_ranges_impl( encoded: R, ranges: ChunkRanges, @@ -843,6 +840,7 @@ where } Ok(()) } + fn read_parent(buf: &[u8]) -> (blake3::Hash, blake3::Hash) { let l_hash = blake3::Hash::from(<[u8; 32]>::try_from(&buf[..32]).unwrap()); let r_hash = blake3::Hash::from(<[u8; 32]>::try_from(&buf[32..64]).unwrap()); @@ -858,7 +856,6 @@ pub async fn outboard( tree: BaoTree, outboard: impl OutboardMut, ) -> io::Result { - // Shared impl, standard BLAKE3 hash mode. outboard_with_hash_strategy(data, tree, outboard, Standard).await } @@ -869,7 +866,6 @@ pub async fn keyed_outboard( outboard: impl OutboardMut, key: &[u8; 32], ) -> io::Result { - // Shared impl, keyed BLAKE3 mode. outboard_with_hash_strategy(data, tree, outboard, Keyed(*key)).await } @@ -884,9 +880,7 @@ async fn outboard_with_hash_strategy( outboard_impl(tree, data, &mut outboard, &mut buffer, hash_strategy).await } -/// Shared outboard traversal for standard and keyed APIs. -/// -/// `hash_strategy` selects which [BaoHashing] implementation computes hashes. +/// Generic outboard traversal monomorphized over the compile time hashing strategy. async fn outboard_impl( tree: BaoTree, mut data: impl AsyncStreamReader, @@ -934,7 +928,6 @@ pub async fn outboard_post_order( tree: BaoTree, outboard: impl AsyncStreamWriter, ) -> io::Result { - // Shared impl, standard BLAKE3 hash mode. outboard_post_order_with_hash_strategy(data, tree, outboard, Standard).await } @@ -945,7 +938,6 @@ pub async fn keyed_outboard_post_order( outboard: impl AsyncStreamWriter, key: &[u8; 32], ) -> io::Result { - // Shared impl, keyed BLAKE3 mode. outboard_post_order_with_hash_strategy(data, tree, outboard, Keyed(*key)).await } @@ -960,9 +952,7 @@ async fn outboard_post_order_with_hash_strategy( outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, hash_strategy).await } -/// Shared post order outboard traversal for standard and keyed APIs. -/// -/// `hash_strategy` selects which [BaoHashing] implementation computes hashes. +/// Generic post order outboard traversal monomorphized over the compile time hashing strategy. async fn outboard_post_order_impl( tree: BaoTree, mut data: impl AsyncStreamReader, @@ -1042,7 +1032,6 @@ mod validate { O: Outboard + 'a, D: AsyncSliceReader + 'a, { - // Shared impl, standard BLAKE3 hash mode. valid_ranges_impl(outboard, data, ranges, Standard) } @@ -1057,11 +1046,10 @@ mod validate { O: Outboard + 'a, D: AsyncSliceReader + 'a, { - // Shared impl, keyed BLAKE3 mode. valid_ranges_impl(outboard, data, ranges, Keyed(*key)) } - /// Shared validation path for standard and keyed APIs. + /// Generic validation body monomorphized over the compile time hashing strategy. fn valid_ranges_impl<'a, O, D, H: BaoHashing + Copy + 'a>( outboard: O, data: D, diff --git a/src/io/sync.rs b/src/io/sync.rs index 697ee53..2058b58 100644 --- a/src/io/sync.rs +++ b/src/io/sync.rs @@ -339,18 +339,21 @@ impl Outboard for PostOrderOutboard { /// Iterator that can be used to decode a response to a range request. /// -/// Generic over `H` so keyed and standard decoders share this implementation. -/// Defaults to [Standard]. Use [Self::new_keyed] for keyed responses. +/// Keyed callers should use [KeyedDecodeResponseIter] via [Self::new_keyed]. #[derive(Debug)] pub struct DecodeResponseIter<'a, R, H: BaoHashing + Copy = Standard> { inner: ResponseIterRef<'a>, stack: SmallVec<[blake3::Hash; 10]>, encoded: R, buf: BytesMut, - /// Compile time hashing strategy, either [Standard] or [Keyed]. hash_strategy: H, } +/// Keyed response decoder iterator. +/// +/// See [DecodeResponseIter::new_keyed]. +pub type KeyedDecodeResponseIter<'a, R> = DecodeResponseIter<'a, R, Keyed>; + impl<'a, R: Read> DecodeResponseIter<'a, R> { /// Create a new iterator to decode a response. /// @@ -382,14 +385,13 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { encoded: R, ranges: &'a ChunkRangesRef, key: &[u8; 32], - ) -> DecodeResponseIter<'a, R, Keyed> { + ) -> KeyedDecodeResponseIter<'a, R> { let buf = BytesMut::with_capacity(tree.block_size().bytes()); DecodeResponseIter::with_hash_strategy(root, tree, encoded, ranges, buf, Keyed(*key)) } } impl<'a, R: Read, H: BaoHashing + Copy> DecodeResponseIter<'a, R, H> { - /// Shared constructor used by [Self::new] and [Self::new_keyed]. pub(crate) fn with_hash_strategy( root: blake3::Hash, tree: BaoTree, @@ -534,7 +536,6 @@ pub fn encode_ranges_validated( ranges: &ChunkRangesRef, encoded: W, ) -> result::Result<(), EncodeError> { - // Shared impl, standard BLAKE3 hash mode. encode_ranges_validated_impl(data, outboard, ranges, encoded, Standard) } @@ -546,13 +547,10 @@ pub fn keyed_encode_ranges_validated( encoded: W, key: &[u8; 32], ) -> result::Result<(), EncodeError> { - // Shared impl, keyed BLAKE3 mode. encode_ranges_validated_impl(data, outboard, ranges, encoded, Keyed(*key)) } -/// Shared encode path for standard and keyed APIs. -/// -/// `hash_strategy` selects which [BaoHashing] implementation validates hashes. +/// Generic encode body monomorphized over the compile time hashing strategy. fn encode_ranges_validated_impl( data: D, outboard: O, @@ -654,7 +652,6 @@ where R: Read, W: WriteAt, { - // Shared impl, standard BLAKE3 hash mode. decode_ranges_impl(encoded, ranges, target, outboard, Standard) } @@ -671,13 +668,10 @@ where R: Read, W: WriteAt, { - // Shared impl, keyed BLAKE3 mode. decode_ranges_impl(encoded, ranges, target, outboard, Keyed(*key)) } -/// Shared decode path for standard and keyed APIs. -/// -/// `hash_strategy` selects which [BaoHashing] implementation verifies hashes. +/// Generic decode body monomorphized over the compile time hashing strategy. fn decode_ranges_impl( encoded: R, ranges: &ChunkRangesRef, @@ -720,7 +714,6 @@ pub fn outboard( tree: BaoTree, outboard: impl OutboardMut, ) -> io::Result { - // Shared impl, standard BLAKE3 hash mode. outboard_with_hash_strategy(data, tree, outboard, Standard) } @@ -731,7 +724,6 @@ pub fn keyed_outboard( outboard: impl OutboardMut, key: &[u8; 32], ) -> io::Result { - // Shared impl, keyed BLAKE3 mode. outboard_with_hash_strategy(data, tree, outboard, Keyed(*key)) } @@ -746,9 +738,7 @@ fn outboard_with_hash_strategy( outboard_impl(tree, data, &mut outboard, &mut buffer, hash_strategy) } -/// Shared outboard traversal for standard and keyed APIs. -/// -/// `hash_strategy` selects which [BaoHashing] implementation computes hashes. +/// Generic outboard traversal monomorphized over the compile time hashing strategy. fn outboard_impl( tree: BaoTree, mut data: impl Read, @@ -797,7 +787,6 @@ pub fn outboard_post_order( tree: BaoTree, outboard: impl Write, ) -> io::Result { - // Shared impl, standard BLAKE3 hash mode. outboard_post_order_with_hash_strategy(data, tree, outboard, Standard) } @@ -808,7 +797,6 @@ pub fn keyed_outboard_post_order( outboard: impl Write, key: &[u8; 32], ) -> io::Result { - // Shared impl, keyed BLAKE3 mode. outboard_post_order_with_hash_strategy(data, tree, outboard, Keyed(*key)) } @@ -823,9 +811,7 @@ fn outboard_post_order_with_hash_strategy( outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, hash_strategy) } -/// Shared post order outboard traversal for standard and keyed APIs. -/// -/// `hash_strategy` selects which [BaoHashing] implementation computes hashes. +/// Generic post order outboard traversal monomorphized over the compile time hashing strategy. fn outboard_post_order_impl( tree: BaoTree, mut data: impl Read, @@ -913,7 +899,6 @@ mod validate { O: Outboard + 'a, D: ReadAt + 'a, { - // Shared impl, standard BLAKE3 hash mode. valid_ranges_impl(outboard, data, ranges, Standard) } @@ -928,11 +913,10 @@ mod validate { O: Outboard + 'a, D: ReadAt + 'a, { - // Shared impl, keyed BLAKE3 mode. valid_ranges_impl(outboard, data, ranges, Keyed(*key)) } - /// Shared validation path for standard and keyed APIs. + /// Generic validation body monomorphized over the compile time hashing strategy. fn valid_ranges_impl<'a, O, D, H: BaoHashing + Copy + 'a>( outboard: O, data: D, diff --git a/src/lib.rs b/src/lib.rs index 4d79779..134cdcb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -239,18 +239,10 @@ pub type ByteRanges = range_collections::RangeSet2; /// [ChunkRanges] implements [`AsRef`]. pub type ChunkRangesRef = range_collections::RangeSetRef; -/// Hashing strategy for shared encode and decode paths. +/// Compile time hashing strategy for shared encode and decode paths. /// -/// Internal IO is generic over this trait so standard and keyed modes share one -/// implementation without runtime branches or duplicated bodies. -/// -/// Pass [Standard] for BLAKE3 hash mode or [Keyed] for BLAKE3 keyed mode. The -/// strategy is a compile time value, so each call site monomorphizes with zero -/// extra cost. -/// -/// Public entry points like `encode_ranges_validated` and -/// `keyed_encode_ranges_validated` pick the strategy at the API boundary and -/// delegate to a single shared function. +/// Use the standard or `keyed_*` public APIs rather than this trait directly. +#[doc(hidden)] pub trait BaoHashing: Copy { /// Hash a subtree of chunk data. fn hash_subtree(&self, start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash; @@ -263,10 +255,8 @@ pub trait BaoHashing: Copy { ) -> blake3::Hash; } -/// BLAKE3 hash mode. Default [BaoHashing] strategy for unkeyed APIs. -/// -/// Routes through the crate internal [hash_subtree] and [parent_cv] helpers so -/// validate only builds still link those symbols. +/// BLAKE3 hash mode strategy for unkeyed APIs. +#[doc(hidden)] #[derive(Clone, Copy)] pub struct Standard; @@ -285,7 +275,8 @@ impl BaoHashing for Standard { } } -/// BLAKE3 keyed mode. Wraps a 32 byte key for domain separated hashing. +/// BLAKE3 keyed mode strategy. Wraps a 32 byte key for domain separated hashing. +#[doc(hidden)] #[derive(Clone, Copy)] pub struct Keyed(pub [u8; 32]); diff --git a/src/rec.rs b/src/rec.rs index ec6c6ac..fc1c0a6 100644 --- a/src/rec.rs +++ b/src/rec.rs @@ -96,10 +96,8 @@ fn truncated_len(ranges: &ChunkRangesRef, size: u64) -> usize { /// This is used as a reference implementation in tests, but also to compute hashes /// below the chunk group size when creating responses for outboards with a chunk group /// size of >0. -/// Recursive reference encoder shared by standard and keyed paths. /// -/// `hash_strategy` selects which [crate::BaoHashing] implementation to use when -/// computing subtree and parent hashes. +/// `hash_strategy` is the compile time hashing mode for subtree and parent hashes. #[allow(clippy::too_many_arguments)] // keyed mode adds `hash_strategy`; splitting into a struct isn't worth it here pub(crate) fn encode_selected_rec( start_chunk: ChunkNum, From b2dc6306d976c28d319092b215d7a19bd03fa6d0 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Aug 2026 10:33:31 +0200 Subject: [PATCH 04/12] Add back missing Debug instance for Standard and add Debug instance for Keyed. --- src/lib.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/lib.rs b/src/lib.rs index 134cdcb..53ea769 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -257,7 +257,7 @@ pub trait BaoHashing: Copy { /// BLAKE3 hash mode strategy for unkeyed APIs. #[doc(hidden)] -#[derive(Clone, Copy)] +#[derive(Debug, Clone, Copy)] pub struct Standard; impl BaoHashing for Standard { @@ -280,6 +280,12 @@ impl BaoHashing for Standard { #[derive(Clone, Copy)] pub struct Keyed(pub [u8; 32]); +impl std::fmt::Debug for Keyed { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Keyed").finish_non_exhaustive() + } +} + impl BaoHashing for Keyed { fn hash_subtree(&self, start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash { keyed_hash_subtree(start_chunk, data, is_root, &self.0) From 4f6c865652e159044517f72b33e37c7911faa21e Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Aug 2026 10:50:32 +0200 Subject: [PATCH 05/12] Implement keyed hashing for the mixed io --- src/io/mixed.rs | 300 ++++++++++++++++++++++++++++++++---------------- 1 file changed, 204 insertions(+), 96 deletions(-) diff --git a/src/io/mixed.rs b/src/io/mixed.rs index 0f9f623..4bdf64e 100644 --- a/src/io/mixed.rs +++ b/src/io/mixed.rs @@ -8,8 +8,8 @@ use smallvec::SmallVec; use super::{sync::Outboard, EncodeError, Leaf, Parent}; use crate::{ - hash_subtree, iter::BaoChunk, parent_cv, rec::truncate_ranges, split_inner, ChunkNum, - ChunkRangesRef, TreeNode, + iter::BaoChunk, rec::truncate_ranges, split_inner, BaoHashing, ChunkNum, ChunkRangesRef, Keyed, + Standard, TreeNode, }; /// A content item for the bao streaming protocol. @@ -84,122 +84,144 @@ where O: Outboard, F: Sender, { - send.send(EncodedItem::Size(outboard.tree().size())).await?; - let res = match traverse_ranges_validated_impl(data, outboard, ranges, send).await { - Ok(Ok(())) => EncodedItem::Done, - Err(cause) => EncodedItem::Error(cause), - Ok(Err(err)) => return Err(err), - }; - send.send(res).await + traverse_ranges_validated_impl(data, outboard, ranges, send, Standard).await } -/// Encode ranges relevant to a query from a reader and outboard to a writer +/// Traverse ranges relevant to a query from a reader and keyed outboard to a stream /// -/// This function validates the data before writing. +/// This function validates the data before writing, using BLAKE3 keyed hashing. /// /// It is possible to encode ranges from a partial file and outboard. /// This will either succeed if the requested ranges are all present, or fail /// as soon as a range is missing. -async fn traverse_ranges_validated_impl( +pub async fn keyed_traverse_ranges_validated( data: D, outboard: O, ranges: &ChunkRangesRef, send: &mut F, -) -> result::Result, EncodeError> + key: &[u8; 32], +) -> std::result::Result<(), F::Error> where D: ReadBytesAt, O: Outboard, F: Sender, { - if ranges.is_empty() { - return Ok(Ok(())); - } - let mut stack: SmallVec<[_; 10]> = SmallVec::<[blake3::Hash; 10]>::new(); - stack.push(outboard.root()); - let data = data; - let tree = outboard.tree(); - // canonicalize ranges - let ranges = truncate_ranges(ranges, tree.size()); - for item in tree.ranges_pre_order_chunks_iter_ref(ranges, 0) { - match item { - BaoChunk::Parent { - is_root, - left, - right, - node, - .. - } => { - let (l_hash, r_hash) = outboard.load(node)?.unwrap(); - let actual = parent_cv(&l_hash, &r_hash, is_root); - let expected = stack.pop().unwrap(); - if actual != expected { - return Err(EncodeError::ParentHashMismatch(node)); - } - if right { - stack.push(r_hash); - } - if left { - stack.push(l_hash); - } - let item = Parent { + traverse_ranges_validated_impl(data, outboard, ranges, send, Keyed(*key)).await +} + +async fn traverse_ranges_validated_impl( + data: D, + outboard: O, + ranges: &ChunkRangesRef, + send: &mut F, + hash_strategy: H, +) -> std::result::Result<(), F::Error> +where + D: ReadBytesAt, + O: Outboard, + F: Sender, + H: BaoHashing, +{ + send.send(EncodedItem::Size(outboard.tree().size())).await?; + let res: result::Result, EncodeError> = async { + if ranges.is_empty() { + return Ok(Ok(())); + } + let mut stack: SmallVec<[_; 10]> = SmallVec::<[blake3::Hash; 10]>::new(); + stack.push(outboard.root()); + let data = data; + let tree = outboard.tree(); + // canonicalize ranges + let ranges = truncate_ranges(ranges, tree.size()); + for item in tree.ranges_pre_order_chunks_iter_ref(ranges, 0) { + match item { + BaoChunk::Parent { + is_root, + left, + right, node, - pair: (l_hash, r_hash), - }; - if let Err(e) = send.send(item.into()).await { - return Ok(Err(e)); - } - } - BaoChunk::Leaf { - start_chunk, - size, - is_root, - ranges, - .. - } => { - let expected = stack.pop().unwrap(); - let start = start_chunk.to_bytes(); - let buffer = data.read_bytes_at(start, size)?; - if !ranges.is_all() { - // we need to encode just a part of the data - // - // write into an out buffer to ensure we detect mismatches - // before writing to the output. - let mut out_buf = Vec::new(); - let actual = traverse_selected_rec( - start_chunk, - buffer, - is_root, - ranges, - tree.block_size.to_u32(), - true, - &mut out_buf, - ); + .. + } => { + let (l_hash, r_hash) = outboard.load(node)?.unwrap(); + let actual = hash_strategy.parent_cv(&l_hash, &r_hash, is_root); + let expected = stack.pop().unwrap(); if actual != expected { - return Err(EncodeError::LeafHashMismatch(start_chunk)); + return Err(EncodeError::ParentHashMismatch(node)); } - for item in out_buf.into_iter() { - if let Err(e) = send.send(item).await { - return Ok(Err(e)); - } + if right { + stack.push(r_hash); } - } else { - let actual = hash_subtree(start_chunk.0, &buffer, is_root); - #[allow(clippy::redundant_slicing)] - if actual != expected { - return Err(EncodeError::LeafHashMismatch(start_chunk)); + if left { + stack.push(l_hash); } - let item = Leaf { - data: buffer, - offset: start_chunk.to_bytes(), + let item = Parent { + node, + pair: (l_hash, r_hash), }; if let Err(e) = send.send(item.into()).await { return Ok(Err(e)); } - }; + } + BaoChunk::Leaf { + start_chunk, + size, + is_root, + ranges, + .. + } => { + let expected = stack.pop().unwrap(); + let start = start_chunk.to_bytes(); + let buffer = data.read_bytes_at(start, size)?; + if !ranges.is_all() { + // we need to encode just a part of the data + // + // write into an out buffer to ensure we detect mismatches + // before writing to the output. + let mut out_buf = Vec::new(); + let actual = traverse_selected_rec_impl( + start_chunk, + buffer, + is_root, + ranges, + tree.block_size.to_u32(), + true, + &mut out_buf, + hash_strategy, + ); + if actual != expected { + return Err(EncodeError::LeafHashMismatch(start_chunk)); + } + for item in out_buf.into_iter() { + if let Err(e) = send.send(item).await { + return Ok(Err(e)); + } + } + } else { + let actual = hash_strategy.hash_subtree(start_chunk.0, &buffer, is_root); + #[allow(clippy::redundant_slicing)] + if actual != expected { + return Err(EncodeError::LeafHashMismatch(start_chunk)); + } + let item = Leaf { + data: buffer, + offset: start_chunk.to_bytes(), + }; + if let Err(e) = send.send(item.into()).await { + return Ok(Err(e)); + } + }; + } } } + Ok(Ok(())) } - Ok(Ok(())) + .await; + let res = match res { + Ok(Ok(())) => EncodedItem::Done, + Err(cause) => EncodedItem::Error(cause), + Ok(Err(err)) => return Err(err), + }; + send.send(res).await } /// Encode ranges relevant to a query from a slice and outboard to a buffer. @@ -207,7 +229,7 @@ where /// This will compute the root hash, so it will have to traverse the entire tree. /// The `ranges` parameter just controls which parts of the data are written. /// -/// Except for writing to a buffer, this is the same as [hash_subtree]. +/// Except for writing to a buffer, this is the same as [crate::hash_subtree]. /// The `min_level` parameter controls the minimum level that will be emitted as a leaf. /// Set this to 0 to disable chunk groups entirely. /// The `emit_data` parameter controls whether the data is written to the buffer. @@ -228,6 +250,53 @@ pub fn traverse_selected_rec( min_level: u32, emit_data: bool, res: &mut Vec, +) -> blake3::Hash { + traverse_selected_rec_impl( + start_chunk, + data, + is_root, + query, + min_level, + emit_data, + res, + Standard, + ) +} + +/// Keyed version of [traverse_selected_rec], using BLAKE3 keyed hashing. +#[allow(clippy::too_many_arguments)] +pub fn keyed_traverse_selected_rec( + start_chunk: ChunkNum, + data: Bytes, + is_root: bool, + query: &ChunkRangesRef, + min_level: u32, + emit_data: bool, + res: &mut Vec, + key: &[u8; 32], +) -> blake3::Hash { + traverse_selected_rec_impl( + start_chunk, + data, + is_root, + query, + min_level, + emit_data, + res, + Keyed(*key), + ) +} + +#[allow(clippy::too_many_arguments)] +fn traverse_selected_rec_impl( + start_chunk: ChunkNum, + data: Bytes, + is_root: bool, + query: &ChunkRangesRef, + min_level: u32, + emit_data: bool, + res: &mut Vec, + hash_strategy: H, ) -> blake3::Hash { use blake3::CHUNK_LEN; if data.len() <= CHUNK_LEN { @@ -240,7 +309,7 @@ pub fn traverse_selected_rec( .into(), ); } - hash_subtree(start_chunk.0, &data, is_root) + hash_strategy.hash_subtree(start_chunk.0, &data, is_root) } else { let chunks = data.len() / CHUNK_LEN + (data.len() % CHUNK_LEN != 0) as usize; let chunks = chunks.next_power_of_two(); @@ -268,7 +337,7 @@ pub fn traverse_selected_rec( None }; // recurse to the left and right to compute the hashes and emit data - let left = traverse_selected_rec( + let left = traverse_selected_rec_impl( start_chunk, data.slice(..mid_bytes), false, @@ -276,8 +345,9 @@ pub fn traverse_selected_rec( min_level, emit_data, res, + hash_strategy, ); - let right = traverse_selected_rec( + let right = traverse_selected_rec_impl( mid_chunk, data.slice(mid_bytes..), false, @@ -285,6 +355,7 @@ pub fn traverse_selected_rec( min_level, emit_data, res, + hash_strategy, ); // backfill the hashes if needed if let Some(o) = hash_offset { @@ -296,7 +367,7 @@ pub fn traverse_selected_rec( } .into(); } - parent_cv(&left, &right, is_root) + hash_strategy.parent_cv(&left, &right, is_root) } } @@ -304,7 +375,10 @@ pub fn traverse_selected_rec( mod tests { use super::*; use crate::{ - io::{outboard::PreOrderMemOutboard, sync::encode_ranges_validated}, + io::{ + outboard::PreOrderMemOutboard, + sync::{encode_ranges_validated, keyed_encode_ranges_validated}, + }, BlockSize, ChunkRanges, }; @@ -343,6 +417,40 @@ mod tests { let encoded2 = flatten(res); assert_eq!(encoded, encoded2); } + + #[tokio::test] + async fn keyed_smoke() { + let data = [0u8; 100000]; + let key = blake3::derive_key("bao-tree.test", b"mixed"); + let outboard = PreOrderMemOutboard::create_keyed(data, BlockSize::from_chunk_log(4), &key); + let (mut tx, mut rx) = tokio::sync::mpsc::channel(10); + let mut encoded = Vec::new(); + keyed_encode_ranges_validated( + &data[..], + &outboard, + &ChunkRanges::all(), + &mut encoded, + &key, + ) + .unwrap(); + tokio::spawn(async move { + keyed_traverse_ranges_validated( + &data[..], + &outboard, + &ChunkRanges::all(), + &mut tx, + &key, + ) + .await + .unwrap(); + }); + let mut res = Vec::new(); + while let Some(item) = rx.recv().await { + res.push(item); + } + let encoded2 = flatten(res); + assert_eq!(encoded, encoded2); + } } /// Trait identical to `ReadAt` but returning `Bytes` instead of reading into a buffer. From 31f4feedf92ee56add20abf16395178bd9a21b6d Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Aug 2026 11:00:31 +0200 Subject: [PATCH 06/12] Replace open BaoHashing strategy pattern with an enum This way we don't have a public API that we have to hide. Also less generics noise. I don't have an use case for somebody actually supplying another strategy. In fact I think it is easy to get this wrong so I would prefer people to fork in that case. --- src/io/fsm.rs | 173 +++++++++++++++++++++--------------------------- src/io/mixed.rs | 34 +++++----- src/io/sync.rs | 133 ++++++++++++++++++------------------- src/lib.rs | 70 +++++++------------- src/rec.rs | 22 +++--- src/tests.rs | 10 +-- src/tests2.rs | 6 +- 7 files changed, 200 insertions(+), 248 deletions(-) diff --git a/src/io/fsm.rs b/src/io/fsm.rs index 21fc1a0..05c3d1f 100644 --- a/src/io/fsm.rs +++ b/src/io/fsm.rs @@ -30,7 +30,7 @@ use crate::{ }, iter::{BaoChunk, ResponseIter}, rec::{encode_selected_rec, truncate_ranges, truncate_ranges_owned}, - BaoHashing, BaoTree, BlockSize, ChunkRanges, ChunkRangesRef, Keyed, Standard, TreeNode, + BaoTree, BlockSize, ChunkRanges, ChunkRangesRef, HashMode, TreeNode, }; /// A binary merkle tree for blake3 hashes of a blob. @@ -402,20 +402,20 @@ pub(crate) fn parse_hash_pair(buf: Bytes) -> io::Result<(blake3::Hash, blake3::H } #[derive(Debug)] -struct ResponseDecoderInner { +struct ResponseDecoderInner { iter: ResponseIter, stack: SmallVec<[blake3::Hash; 10]>, encoded: R, - hash_strategy: H, + mode: HashMode, } -impl ResponseDecoderInner { - fn with_hash_strategy( +impl ResponseDecoderInner { + fn with_mode( tree: BaoTree, hash: blake3::Hash, ranges: ChunkRanges, encoded: R, - hash_strategy: H, + mode: HashMode, ) -> Self { // now that we know the size, we can canonicalize the ranges let ranges = truncate_ranges_owned(ranges, tree.size()); @@ -423,7 +423,7 @@ impl ResponseDecoderInner { iter: ResponseIter::new(tree, ranges), stack: SmallVec::new(), encoded, - hash_strategy, + mode, }; res.stack.push(hash); res @@ -432,22 +432,17 @@ impl ResponseDecoderInner { /// Response decoder. /// -/// Keyed callers should use [KeyedResponseDecoder] via [Self::new_keyed]. +/// Keyed callers should use [Self::new_keyed]. #[derive(Debug)] -pub struct ResponseDecoder(Box>); - -/// Keyed response decoder. -/// -/// See [ResponseDecoder::new_keyed]. -pub type KeyedResponseDecoder = ResponseDecoder; +pub struct ResponseDecoder(Box>); /// Next type for [ResponseDecoder]. #[derive(Debug)] -pub enum ResponseDecoderNext { +pub enum ResponseDecoderNext { /// One more item, and you get back the state machine in the next state More( ( - ResponseDecoder, + ResponseDecoder, std::result::Result, ), ), @@ -455,18 +450,17 @@ pub enum ResponseDecoderNext { Done(R), } -/// Next type for [KeyedResponseDecoder]. -/// -/// See [ResponseDecoder::new_keyed]. -pub type KeyedResponseDecoderNext = ResponseDecoderNext; - impl ResponseDecoder { /// Create a new response decoder state machine, when you have already read the size. /// /// The size as well as the chunk size is given in the `tree` parameter. pub fn new(hash: blake3::Hash, ranges: ChunkRanges, tree: BaoTree, encoded: R) -> Self { - Self(Box::new(ResponseDecoderInner::with_hash_strategy( - tree, hash, ranges, encoded, Standard, + Self(Box::new(ResponseDecoderInner::with_mode( + tree, + hash, + ranges, + encoded, + HashMode::Standard, ))) } @@ -477,36 +471,32 @@ impl ResponseDecoder { tree: BaoTree, encoded: R, key: &[u8; 32], - ) -> KeyedResponseDecoder { - ResponseDecoder(Box::new(ResponseDecoderInner::with_hash_strategy( + ) -> Self { + Self(Box::new(ResponseDecoderInner::with_mode( tree, hash, ranges, encoded, - Keyed(*key), + HashMode::Keyed(*key), ))) } } -impl ResponseDecoder { - pub(crate) fn with_hash_strategy( +impl ResponseDecoder { + pub(crate) fn with_mode( hash: blake3::Hash, ranges: ChunkRanges, tree: BaoTree, encoded: R, - hash_strategy: H, + mode: HashMode, ) -> Self { - Self(Box::new(ResponseDecoderInner::with_hash_strategy( - tree, - hash, - ranges, - encoded, - hash_strategy, + Self(Box::new(ResponseDecoderInner::with_mode( + tree, hash, ranges, encoded, mode, ))) } /// Proceed to the next state by reading the next chunk from the stream. - pub async fn next(mut self) -> ResponseDecoderNext { + pub async fn next(mut self) -> ResponseDecoderNext { if let Some(chunk) = self.0.iter.next() { let item = self.next0(chunk).await; ResponseDecoderNext::More((self, item)) @@ -547,7 +537,7 @@ impl ResponseDecoder { .map_err(|e| DecodeError::maybe_parent_not_found(e, node))?; let pair @ (l_hash, r_hash) = read_parent(&buf); let parent_hash = this.stack.pop().unwrap(); - let actual = this.hash_strategy.parent_cv(&l_hash, &r_hash, is_root); + let actual = this.mode.parent_cv(&l_hash, &r_hash, is_root); // Push the children in reverse order so they are popped in the correct order // only push right if the range intersects with the right child if right { @@ -577,9 +567,7 @@ impl ResponseDecoder { .await .map_err(|e| DecodeError::maybe_leaf_not_found(e, start_chunk))?; let leaf_hash = this.stack.pop().unwrap(); - let actual = this - .hash_strategy - .hash_subtree(start_chunk.0, &data, is_root); + let actual = this.mode.hash_subtree(start_chunk.0, &data, is_root); if leaf_hash != actual { return Err(DecodeError::LeafHashMismatch(start_chunk)); } @@ -656,7 +644,7 @@ where O: Outboard, W: AsyncStreamWriter, { - encode_ranges_validated_impl(data, outboard, ranges, encoded, Standard).await + encode_ranges_validated_impl(data, outboard, ranges, encoded, HashMode::Standard).await } /// Encode ranges with BLAKE3 keyed hash validation. @@ -672,16 +660,16 @@ where O: Outboard, W: AsyncStreamWriter, { - encode_ranges_validated_impl(data, outboard, ranges, encoded, Keyed(*key)).await + encode_ranges_validated_impl(data, outboard, ranges, encoded, HashMode::Keyed(*key)).await } /// Generic encode body monomorphized over the compile time hashing strategy. -async fn encode_ranges_validated_impl( +async fn encode_ranges_validated_impl( mut data: D, mut outboard: O, ranges: &ChunkRangesRef, encoded: W, - hash_strategy: H, + mode: HashMode, ) -> result::Result<(), EncodeError> where D: AsyncSliceReader, @@ -706,7 +694,7 @@ where .. } => { let (l_hash, r_hash) = outboard.load(node).await?.unwrap(); - let actual = hash_strategy.parent_cv(&l_hash, &r_hash, is_root); + let actual = mode.parent_cv(&l_hash, &r_hash, is_root); let expected = stack.pop().unwrap(); if actual != expected { return Err(EncodeError::ParentHashMismatch(node)); @@ -747,11 +735,11 @@ where tree.block_size.to_u32(), true, &mut out_buf, - hash_strategy, + mode, ); (actual, out_buf.clone().into()) } else { - let actual = hash_strategy.hash_subtree(start_chunk.0, &bytes, is_root); + let actual = mode.hash_subtree(start_chunk.0, &bytes, is_root); (actual, bytes) }; if actual != expected { @@ -782,7 +770,7 @@ where R: AsyncStreamReader, W: AsyncSliceWriter, { - decode_ranges_impl(encoded, ranges, target, outboard, Standard).await + decode_ranges_impl(encoded, ranges, target, outboard, HashMode::Standard).await } /// Decode a keyed response into a file while updating an outboard. @@ -798,29 +786,24 @@ where R: AsyncStreamReader, W: AsyncSliceWriter, { - decode_ranges_impl(encoded, ranges, target, outboard, Keyed(*key)).await + decode_ranges_impl(encoded, ranges, target, outboard, HashMode::Keyed(*key)).await } /// Generic decode body monomorphized over the compile time hashing strategy. -async fn decode_ranges_impl( +async fn decode_ranges_impl( encoded: R, ranges: ChunkRanges, mut target: W, mut outboard: O, - hash_strategy: H, + mode: HashMode, ) -> std::result::Result<(), DecodeError> where O: OutboardMut + Outboard, R: AsyncStreamReader, W: AsyncSliceWriter, { - let mut reading = ResponseDecoder::with_hash_strategy( - outboard.root(), - ranges, - outboard.tree(), - encoded, - hash_strategy, - ); + let mut reading = + ResponseDecoder::with_mode(outboard.root(), ranges, outboard.tree(), encoded, mode); loop { let item = match reading.next().await { ResponseDecoderNext::Done(_reader) => break, @@ -856,7 +839,7 @@ pub async fn outboard( tree: BaoTree, outboard: impl OutboardMut, ) -> io::Result { - outboard_with_hash_strategy(data, tree, outboard, Standard).await + outboard_with_mode(data, tree, outboard, HashMode::Standard).await } /// Compute the keyed outboard for the given data. @@ -866,27 +849,27 @@ pub async fn keyed_outboard( outboard: impl OutboardMut, key: &[u8; 32], ) -> io::Result { - outboard_with_hash_strategy(data, tree, outboard, Keyed(*key)).await + outboard_with_mode(data, tree, outboard, HashMode::Keyed(*key)).await } /// Allocates a chunk group buffer and delegates to [outboard_impl]. -async fn outboard_with_hash_strategy( +async fn outboard_with_mode( data: impl AsyncStreamReader, tree: BaoTree, mut outboard: impl OutboardMut, - hash_strategy: H, + mode: HashMode, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - outboard_impl(tree, data, &mut outboard, &mut buffer, hash_strategy).await + outboard_impl(tree, data, &mut outboard, &mut buffer, mode).await } /// Generic outboard traversal monomorphized over the compile time hashing strategy. -async fn outboard_impl( +async fn outboard_impl( tree: BaoTree, mut data: impl AsyncStreamReader, mut outboard: impl OutboardMut, buffer: &mut [u8], - hash_strategy: H, + mode: HashMode, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -897,7 +880,7 @@ async fn outboard_impl( let right_hash = stack.pop().unwrap(); let left_hash = stack.pop().unwrap(); outboard.save(node, &(left_hash, right_hash)).await?; - let parent = hash_strategy.parent_cv(&left_hash, &right_hash, is_root); + let parent = mode.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -907,7 +890,7 @@ async fn outboard_impl( .. } => { let buf = data.read_bytes_exact(size).await?; - let hash = hash_strategy.hash_subtree(start_chunk.0, &buf, is_root); + let hash = mode.hash_subtree(start_chunk.0, &buf, is_root); stack.push(hash); } } @@ -928,7 +911,7 @@ pub async fn outboard_post_order( tree: BaoTree, outboard: impl AsyncStreamWriter, ) -> io::Result { - outboard_post_order_with_hash_strategy(data, tree, outboard, Standard).await + outboard_post_order_with_mode(data, tree, outboard, HashMode::Standard).await } /// Compute the keyed post order outboard for the given data. @@ -938,27 +921,27 @@ pub async fn keyed_outboard_post_order( outboard: impl AsyncStreamWriter, key: &[u8; 32], ) -> io::Result { - outboard_post_order_with_hash_strategy(data, tree, outboard, Keyed(*key)).await + outboard_post_order_with_mode(data, tree, outboard, HashMode::Keyed(*key)).await } /// Allocates a chunk group buffer and delegates to [outboard_post_order_impl]. -async fn outboard_post_order_with_hash_strategy( +async fn outboard_post_order_with_mode( data: impl AsyncStreamReader, tree: BaoTree, mut outboard: impl AsyncStreamWriter, - hash_strategy: H, + mode: HashMode, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, hash_strategy).await + outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, mode).await } /// Generic post order outboard traversal monomorphized over the compile time hashing strategy. -async fn outboard_post_order_impl( +async fn outboard_post_order_impl( tree: BaoTree, mut data: impl AsyncStreamReader, mut outboard: impl AsyncStreamWriter, buffer: &mut [u8], - hash_strategy: H, + mode: HashMode, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -970,7 +953,7 @@ async fn outboard_post_order_impl( let left_hash = stack.pop().unwrap(); outboard.write(left_hash.as_bytes()).await?; outboard.write(right_hash.as_bytes()).await?; - let parent = hash_strategy.parent_cv(&left_hash, &right_hash, is_root); + let parent = mode.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -980,7 +963,7 @@ async fn outboard_post_order_impl( .. } => { let buf = data.read_bytes_exact(size).await?; - let hash = hash_strategy.hash_subtree(start_chunk.0, &buf, is_root); + let hash = mode.hash_subtree(start_chunk.0, &buf, is_root); stack.push(hash); } } @@ -1014,8 +997,8 @@ mod validate { use super::Outboard; use crate::{ - blake3, io::LocalBoxFuture, rec::truncate_ranges, split, BaoHashing, BaoTree, ChunkNum, - ChunkRangesRef, Keyed, Standard, TreeNode, + blake3, io::LocalBoxFuture, rec::truncate_ranges, split, BaoTree, ChunkNum, ChunkRangesRef, + HashMode, TreeNode, }; /// Given a data file and an outboard, compute all valid ranges. @@ -1032,7 +1015,7 @@ mod validate { O: Outboard + 'a, D: AsyncSliceReader + 'a, { - valid_ranges_impl(outboard, data, ranges, Standard) + valid_ranges_impl(outboard, data, ranges, HashMode::Standard) } /// Given a data file and a keyed outboard, compute all valid ranges. @@ -1046,15 +1029,15 @@ mod validate { O: Outboard + 'a, D: AsyncSliceReader + 'a, { - valid_ranges_impl(outboard, data, ranges, Keyed(*key)) + valid_ranges_impl(outboard, data, ranges, HashMode::Keyed(*key)) } /// Generic validation body monomorphized over the compile time hashing strategy. - fn valid_ranges_impl<'a, O, D, H: BaoHashing + Copy + 'a>( + fn valid_ranges_impl<'a, O, D>( outboard: O, data: D, ranges: &'a ChunkRangesRef, - hash_strategy: H, + mode: HashMode, ) -> impl Stream>> + 'a where O: Outboard + 'a, @@ -1062,29 +1045,29 @@ mod validate { { Gen::new(move |co| async move { if let Err(cause) = - RecursiveDataValidator::validate(outboard, data, ranges, &co, hash_strategy).await + RecursiveDataValidator::validate(outboard, data, ranges, &co, mode).await { co.yield_(Err(cause)).await; } }) } - struct RecursiveDataValidator<'a, O: Outboard, D: AsyncSliceReader, H: BaoHashing + Copy> { + struct RecursiveDataValidator<'a, O: Outboard, D: AsyncSliceReader> { tree: BaoTree, shifted_filled_size: TreeNode, outboard: O, data: D, co: &'a Co>>, - hash_strategy: H, + mode: HashMode, } - impl RecursiveDataValidator<'_, O, D, H> { + impl RecursiveDataValidator<'_, O, D> { async fn validate( outboard: O, data: D, ranges: &ChunkRangesRef, co: &Co>>, - hash_strategy: H, + mode: HashMode, ) -> io::Result<()> { let tree = outboard.tree(); if tree.blocks() == 1 { @@ -1093,7 +1076,7 @@ mod validate { let data = data .read_exact_at(0, tree.size().try_into().unwrap()) .await?; - let actual = hash_strategy.hash_subtree(0, &data, true); + let actual = mode.hash_subtree(0, &data, true); if actual == outboard.root() { co.yield_(Ok(ChunkNum(0)..tree.chunks())).await; } @@ -1108,7 +1091,7 @@ mod validate { outboard, data, co, - hash_strategy, + mode, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -1124,11 +1107,9 @@ mod validate { let len = (range.end - range.start).try_into().unwrap(); let data = self.data.read_exact_at(range.start, len).await?; // is_root is always false because the case of a single chunk group is handled before calling this function - let actual = self.hash_strategy.hash_subtree( - ChunkNum::full_chunks(range.start).0, - &data, - is_root, - ); + let actual = + self.mode + .hash_subtree(ChunkNum::full_chunks(range.start).0, &data, is_root); if &actual == hash { // yield the left range self.co @@ -1162,7 +1143,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = self.hash_strategy.parent_cv(&l_hash, &r_hash, is_root); + let actual = self.mode.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -1265,7 +1246,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = Standard.parent_cv(&l_hash, &r_hash, is_root); + let actual = HashMode::Standard.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); diff --git a/src/io/mixed.rs b/src/io/mixed.rs index 4bdf64e..363237a 100644 --- a/src/io/mixed.rs +++ b/src/io/mixed.rs @@ -8,8 +8,7 @@ use smallvec::SmallVec; use super::{sync::Outboard, EncodeError, Leaf, Parent}; use crate::{ - iter::BaoChunk, rec::truncate_ranges, split_inner, BaoHashing, ChunkNum, ChunkRangesRef, Keyed, - Standard, TreeNode, + iter::BaoChunk, rec::truncate_ranges, split_inner, ChunkNum, ChunkRangesRef, HashMode, TreeNode, }; /// A content item for the bao streaming protocol. @@ -84,7 +83,7 @@ where O: Outboard, F: Sender, { - traverse_ranges_validated_impl(data, outboard, ranges, send, Standard).await + traverse_ranges_validated_impl(data, outboard, ranges, send, HashMode::Standard).await } /// Traverse ranges relevant to a query from a reader and keyed outboard to a stream @@ -106,21 +105,20 @@ where O: Outboard, F: Sender, { - traverse_ranges_validated_impl(data, outboard, ranges, send, Keyed(*key)).await + traverse_ranges_validated_impl(data, outboard, ranges, send, HashMode::Keyed(*key)).await } -async fn traverse_ranges_validated_impl( +async fn traverse_ranges_validated_impl( data: D, outboard: O, ranges: &ChunkRangesRef, send: &mut F, - hash_strategy: H, + mode: HashMode, ) -> std::result::Result<(), F::Error> where D: ReadBytesAt, O: Outboard, F: Sender, - H: BaoHashing, { send.send(EncodedItem::Size(outboard.tree().size())).await?; let res: result::Result, EncodeError> = async { @@ -143,7 +141,7 @@ where .. } => { let (l_hash, r_hash) = outboard.load(node)?.unwrap(); - let actual = hash_strategy.parent_cv(&l_hash, &r_hash, is_root); + let actual = mode.parent_cv(&l_hash, &r_hash, is_root); let expected = stack.pop().unwrap(); if actual != expected { return Err(EncodeError::ParentHashMismatch(node)); @@ -186,7 +184,7 @@ where tree.block_size.to_u32(), true, &mut out_buf, - hash_strategy, + mode, ); if actual != expected { return Err(EncodeError::LeafHashMismatch(start_chunk)); @@ -197,7 +195,7 @@ where } } } else { - let actual = hash_strategy.hash_subtree(start_chunk.0, &buffer, is_root); + let actual = mode.hash_subtree(start_chunk.0, &buffer, is_root); #[allow(clippy::redundant_slicing)] if actual != expected { return Err(EncodeError::LeafHashMismatch(start_chunk)); @@ -259,7 +257,7 @@ pub fn traverse_selected_rec( min_level, emit_data, res, - Standard, + HashMode::Standard, ) } @@ -283,12 +281,12 @@ pub fn keyed_traverse_selected_rec( min_level, emit_data, res, - Keyed(*key), + HashMode::Keyed(*key), ) } #[allow(clippy::too_many_arguments)] -fn traverse_selected_rec_impl( +fn traverse_selected_rec_impl( start_chunk: ChunkNum, data: Bytes, is_root: bool, @@ -296,7 +294,7 @@ fn traverse_selected_rec_impl( min_level: u32, emit_data: bool, res: &mut Vec, - hash_strategy: H, + mode: HashMode, ) -> blake3::Hash { use blake3::CHUNK_LEN; if data.len() <= CHUNK_LEN { @@ -309,7 +307,7 @@ fn traverse_selected_rec_impl( .into(), ); } - hash_strategy.hash_subtree(start_chunk.0, &data, is_root) + mode.hash_subtree(start_chunk.0, &data, is_root) } else { let chunks = data.len() / CHUNK_LEN + (data.len() % CHUNK_LEN != 0) as usize; let chunks = chunks.next_power_of_two(); @@ -345,7 +343,7 @@ fn traverse_selected_rec_impl( min_level, emit_data, res, - hash_strategy, + mode, ); let right = traverse_selected_rec_impl( mid_chunk, @@ -355,7 +353,7 @@ fn traverse_selected_rec_impl( min_level, emit_data, res, - hash_strategy, + mode, ); // backfill the hashes if needed if let Some(o) = hash_offset { @@ -367,7 +365,7 @@ fn traverse_selected_rec_impl( } .into(); } - hash_strategy.parent_cv(&left, &right, is_root) + mode.parent_cv(&left, &right, is_root) } } diff --git a/src/io/sync.rs b/src/io/sync.rs index 2058b58..6897332 100644 --- a/src/io/sync.rs +++ b/src/io/sync.rs @@ -22,7 +22,7 @@ use crate::{ }, iter::{BaoChunk, ResponseIterRef}, rec::encode_selected_rec, - BaoHashing, BaoTree, BlockSize, ChunkRangesRef, Keyed, Standard, TreeNode, + BaoTree, BlockSize, ChunkRangesRef, HashMode, TreeNode, }; /// A binary merkle tree for blake3 hashes of a blob. @@ -339,21 +339,16 @@ impl Outboard for PostOrderOutboard { /// Iterator that can be used to decode a response to a range request. /// -/// Keyed callers should use [KeyedDecodeResponseIter] via [Self::new_keyed]. +/// Keyed callers should use [Self::new_keyed]. #[derive(Debug)] -pub struct DecodeResponseIter<'a, R, H: BaoHashing + Copy = Standard> { +pub struct DecodeResponseIter<'a, R> { inner: ResponseIterRef<'a>, stack: SmallVec<[blake3::Hash; 10]>, encoded: R, buf: BytesMut, - hash_strategy: H, + mode: HashMode, } -/// Keyed response decoder iterator. -/// -/// See [DecodeResponseIter::new_keyed]. -pub type KeyedDecodeResponseIter<'a, R> = DecodeResponseIter<'a, R, Keyed>; - impl<'a, R: Read> DecodeResponseIter<'a, R> { /// Create a new iterator to decode a response. /// @@ -375,7 +370,7 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { ranges: &'a ChunkRangesRef, buf: BytesMut, ) -> Self { - DecodeResponseIter::with_hash_strategy(root, tree, encoded, ranges, buf, Standard) + DecodeResponseIter::with_mode(root, tree, encoded, ranges, buf, HashMode::Standard) } /// Create a new iterator to decode a keyed response. @@ -385,20 +380,20 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { encoded: R, ranges: &'a ChunkRangesRef, key: &[u8; 32], - ) -> KeyedDecodeResponseIter<'a, R> { + ) -> Self { let buf = BytesMut::with_capacity(tree.block_size().bytes()); - DecodeResponseIter::with_hash_strategy(root, tree, encoded, ranges, buf, Keyed(*key)) + DecodeResponseIter::with_mode(root, tree, encoded, ranges, buf, HashMode::Keyed(*key)) } } -impl<'a, R: Read, H: BaoHashing + Copy> DecodeResponseIter<'a, R, H> { - pub(crate) fn with_hash_strategy( +impl<'a, R: Read> DecodeResponseIter<'a, R> { + pub(crate) fn with_mode( root: blake3::Hash, tree: BaoTree, encoded: R, ranges: &'a ChunkRangesRef, buf: BytesMut, - hash_strategy: H, + mode: HashMode, ) -> Self { let ranges = truncate_ranges(ranges, tree.size()); let mut stack = SmallVec::new(); @@ -408,7 +403,7 @@ impl<'a, R: Read, H: BaoHashing + Copy> DecodeResponseIter<'a, R, H> { inner: ResponseIterRef::new(tree, ranges), encoded, buf, - hash_strategy, + mode, } } @@ -436,7 +431,7 @@ impl<'a, R: Read, H: BaoHashing + Copy> DecodeResponseIter<'a, R, H> { let pair @ (l_hash, r_hash) = read_parent(&mut self.encoded) .map_err(|e| DecodeError::maybe_parent_not_found(e, node))?; let parent_hash = self.stack.pop().unwrap(); - let actual = self.hash_strategy.parent_cv(&l_hash, &r_hash, is_root); + let actual = self.mode.parent_cv(&l_hash, &r_hash, is_root); if parent_hash != actual { return Err(DecodeError::ParentHashMismatch(node)); } @@ -458,9 +453,7 @@ impl<'a, R: Read, H: BaoHashing + Copy> DecodeResponseIter<'a, R, H> { self.encoded .read_exact(&mut self.buf) .map_err(|e| DecodeError::maybe_leaf_not_found(e, start_chunk))?; - let actual = self - .hash_strategy - .hash_subtree(start_chunk.0, &self.buf, is_root); + let actual = self.mode.hash_subtree(start_chunk.0, &self.buf, is_root); let leaf_hash = self.stack.pop().unwrap(); if leaf_hash != actual { return Err(DecodeError::LeafHashMismatch(start_chunk)); @@ -478,7 +471,7 @@ impl<'a, R: Read, H: BaoHashing + Copy> DecodeResponseIter<'a, R, H> { } } -impl Iterator for DecodeResponseIter<'_, R, H> { +impl Iterator for DecodeResponseIter<'_, R> { type Item = result::Result; fn next(&mut self) -> Option { @@ -536,7 +529,7 @@ pub fn encode_ranges_validated( ranges: &ChunkRangesRef, encoded: W, ) -> result::Result<(), EncodeError> { - encode_ranges_validated_impl(data, outboard, ranges, encoded, Standard) + encode_ranges_validated_impl(data, outboard, ranges, encoded, HashMode::Standard) } /// Encode ranges with BLAKE3 keyed hash validation. @@ -547,16 +540,16 @@ pub fn keyed_encode_ranges_validated( encoded: W, key: &[u8; 32], ) -> result::Result<(), EncodeError> { - encode_ranges_validated_impl(data, outboard, ranges, encoded, Keyed(*key)) + encode_ranges_validated_impl(data, outboard, ranges, encoded, HashMode::Keyed(*key)) } /// Generic encode body monomorphized over the compile time hashing strategy. -fn encode_ranges_validated_impl( +fn encode_ranges_validated_impl( data: D, outboard: O, ranges: &ChunkRangesRef, encoded: W, - hash_strategy: H, + mode: HashMode, ) -> result::Result<(), EncodeError> { if ranges.is_empty() { return Ok(()); @@ -580,7 +573,7 @@ fn encode_ranges_validated_impl .. } => { let (l_hash, r_hash) = outboard.load(node)?.unwrap(); - let actual = hash_strategy.parent_cv(&l_hash, &r_hash, is_root); + let actual = mode.parent_cv(&l_hash, &r_hash, is_root); let expected = stack.pop().unwrap(); if actual != expected { return Err(EncodeError::ParentHashMismatch(node)); @@ -619,11 +612,11 @@ fn encode_ranges_validated_impl tree.block_size.to_u32(), true, &mut out_buf, - hash_strategy, + mode, ); (actual, &out_buf[..]) } else { - let actual = hash_strategy.hash_subtree(start_chunk.0, buf, is_root); + let actual = mode.hash_subtree(start_chunk.0, buf, is_root); #[allow(clippy::redundant_slicing)] (actual, &buf[..]) }; @@ -652,7 +645,7 @@ where R: Read, W: WriteAt, { - decode_ranges_impl(encoded, ranges, target, outboard, Standard) + decode_ranges_impl(encoded, ranges, target, outboard, HashMode::Standard) } /// Decode a keyed response into a file while updating an outboard. @@ -668,29 +661,29 @@ where R: Read, W: WriteAt, { - decode_ranges_impl(encoded, ranges, target, outboard, Keyed(*key)) + decode_ranges_impl(encoded, ranges, target, outboard, HashMode::Keyed(*key)) } /// Generic decode body monomorphized over the compile time hashing strategy. -fn decode_ranges_impl( +fn decode_ranges_impl( encoded: R, ranges: &ChunkRangesRef, mut target: W, mut outboard: O, - hash_strategy: H, + mode: HashMode, ) -> std::result::Result<(), DecodeError> where O: OutboardMut + Outboard, R: Read, W: WriteAt, { - let iter = DecodeResponseIter::with_hash_strategy( + let iter = DecodeResponseIter::with_mode( outboard.root(), outboard.tree(), encoded, ranges, BytesMut::with_capacity(outboard.tree().block_size().bytes()), - hash_strategy, + mode, ); for item in iter { match item? { @@ -714,7 +707,7 @@ pub fn outboard( tree: BaoTree, outboard: impl OutboardMut, ) -> io::Result { - outboard_with_hash_strategy(data, tree, outboard, Standard) + outboard_with_mode(data, tree, outboard, HashMode::Standard) } /// Compute the keyed outboard for the given data. @@ -724,27 +717,27 @@ pub fn keyed_outboard( outboard: impl OutboardMut, key: &[u8; 32], ) -> io::Result { - outboard_with_hash_strategy(data, tree, outboard, Keyed(*key)) + outboard_with_mode(data, tree, outboard, HashMode::Keyed(*key)) } /// Allocates a chunk group buffer and delegates to [outboard_impl]. -fn outboard_with_hash_strategy( +fn outboard_with_mode( data: impl Read, tree: BaoTree, mut outboard: impl OutboardMut, - hash_strategy: H, + mode: HashMode, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - outboard_impl(tree, data, &mut outboard, &mut buffer, hash_strategy) + outboard_impl(tree, data, &mut outboard, &mut buffer, mode) } /// Generic outboard traversal monomorphized over the compile time hashing strategy. -fn outboard_impl( +fn outboard_impl( tree: BaoTree, mut data: impl Read, mut outboard: impl OutboardMut, buffer: &mut [u8], - hash_strategy: H, + mode: HashMode, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -755,7 +748,7 @@ fn outboard_impl( let right_hash = stack.pop().unwrap(); let left_hash = stack.pop().unwrap(); outboard.save(node, &(left_hash, right_hash))?; - let parent = hash_strategy.parent_cv(&left_hash, &right_hash, is_root); + let parent = mode.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -766,7 +759,7 @@ fn outboard_impl( } => { let buf = &mut buffer[..size]; data.read_exact(buf)?; - let hash = hash_strategy.hash_subtree(start_chunk.0, buf, is_root); + let hash = mode.hash_subtree(start_chunk.0, buf, is_root); stack.push(hash); } } @@ -787,7 +780,7 @@ pub fn outboard_post_order( tree: BaoTree, outboard: impl Write, ) -> io::Result { - outboard_post_order_with_hash_strategy(data, tree, outboard, Standard) + outboard_post_order_with_mode(data, tree, outboard, HashMode::Standard) } /// Compute the keyed post order outboard for the given data. @@ -797,27 +790,27 @@ pub fn keyed_outboard_post_order( outboard: impl Write, key: &[u8; 32], ) -> io::Result { - outboard_post_order_with_hash_strategy(data, tree, outboard, Keyed(*key)) + outboard_post_order_with_mode(data, tree, outboard, HashMode::Keyed(*key)) } /// Allocates a chunk group buffer and delegates to [outboard_post_order_impl]. -fn outboard_post_order_with_hash_strategy( +fn outboard_post_order_with_mode( data: impl Read, tree: BaoTree, mut outboard: impl Write, - hash_strategy: H, + mode: HashMode, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, hash_strategy) + outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, mode) } /// Generic post order outboard traversal monomorphized over the compile time hashing strategy. -fn outboard_post_order_impl( +fn outboard_post_order_impl( tree: BaoTree, mut data: impl Read, mut outboard: impl Write, buffer: &mut [u8], - hash_strategy: H, + mode: HashMode, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -829,7 +822,7 @@ fn outboard_post_order_impl( let left_hash = stack.pop().unwrap(); outboard.write_all(left_hash.as_bytes())?; outboard.write_all(right_hash.as_bytes())?; - let parent = hash_strategy.parent_cv(&left_hash, &right_hash, is_root); + let parent = mode.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -840,7 +833,7 @@ fn outboard_post_order_impl( } => { let buf = &mut buffer[..size]; data.read_exact(buf)?; - let hash = hash_strategy.hash_subtree(start_chunk.0, buf, is_root); + let hash = mode.hash_subtree(start_chunk.0, buf, is_root); stack.push(hash); } } @@ -881,8 +874,8 @@ mod validate { use super::Outboard; use crate::{ - blake3, io::LocalBoxFuture, rec::truncate_ranges, split, BaoHashing, BaoTree, ChunkNum, - ChunkRangesRef, Keyed, Standard, TreeNode, + blake3, io::LocalBoxFuture, rec::truncate_ranges, split, BaoTree, ChunkNum, ChunkRangesRef, + HashMode, TreeNode, }; /// Given a data file and an outboard, compute all valid ranges. @@ -899,7 +892,7 @@ mod validate { O: Outboard + 'a, D: ReadAt + 'a, { - valid_ranges_impl(outboard, data, ranges, Standard) + valid_ranges_impl(outboard, data, ranges, HashMode::Standard) } /// Given a data file and a keyed outboard, compute all valid ranges. @@ -913,15 +906,15 @@ mod validate { O: Outboard + 'a, D: ReadAt + 'a, { - valid_ranges_impl(outboard, data, ranges, Keyed(*key)) + valid_ranges_impl(outboard, data, ranges, HashMode::Keyed(*key)) } /// Generic validation body monomorphized over the compile time hashing strategy. - fn valid_ranges_impl<'a, O, D, H: BaoHashing + Copy + 'a>( + fn valid_ranges_impl<'a, O, D>( outboard: O, data: D, ranges: &'a ChunkRangesRef, - hash_strategy: H, + mode: HashMode, ) -> impl IntoIterator>> + 'a where O: Outboard + 'a, @@ -929,30 +922,30 @@ mod validate { { Gen::new(move |co| async move { if let Err(cause) = - RecursiveDataValidator::validate(outboard, data, ranges, &co, hash_strategy).await + RecursiveDataValidator::validate(outboard, data, ranges, &co, mode).await { co.yield_(Err(cause)).await; } }) } - struct RecursiveDataValidator<'a, O: Outboard, D: ReadAt, H: BaoHashing + Copy> { + struct RecursiveDataValidator<'a, O: Outboard, D: ReadAt> { tree: BaoTree, shifted_filled_size: TreeNode, outboard: O, data: D, buffer: Vec, co: &'a Co>>, - hash_strategy: H, + mode: HashMode, } - impl RecursiveDataValidator<'_, O, D, H> { + impl RecursiveDataValidator<'_, O, D> { async fn validate( outboard: O, data: D, ranges: &ChunkRangesRef, co: &Co>>, - hash_strategy: H, + mode: HashMode, ) -> io::Result<()> { let tree = outboard.tree(); let mut buffer = vec![0u8; tree.chunk_group_bytes()]; @@ -960,7 +953,7 @@ mod validate { // special case for a tree that fits in one block / chunk group let tmp = &mut buffer[..tree.size().try_into().unwrap()]; data.read_exact_at(0, tmp)?; - let actual = hash_strategy.hash_subtree(0, tmp, true); + let actual = mode.hash_subtree(0, tmp, true); if actual == outboard.root() { co.yield_(Ok(ChunkNum(0)..tree.chunks())).await; } @@ -976,7 +969,7 @@ mod validate { data, buffer, co, - hash_strategy, + mode, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -993,9 +986,9 @@ mod validate { let tmp = &mut self.buffer[..len]; self.data.read_exact_at(range.start, tmp)?; // is_root is always false because the case of a single chunk group is handled before calling this function - let actual = - self.hash_strategy - .hash_subtree(ChunkNum::full_chunks(range.start).0, tmp, is_root); + let actual = self + .mode + .hash_subtree(ChunkNum::full_chunks(range.start).0, tmp, is_root); if &actual == hash { // yield the left range self.co @@ -1029,7 +1022,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = self.hash_strategy.parent_cv(&l_hash, &r_hash, is_root); + let actual = self.mode.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -1131,7 +1124,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = Standard.parent_cv(&l_hash, &r_hash, is_root); + let actual = HashMode::Standard.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); diff --git a/src/lib.rs b/src/lib.rs index 53ea769..57ccfe7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -239,65 +239,45 @@ pub type ByteRanges = range_collections::RangeSet2; /// [ChunkRanges] implements [`AsRef`]. pub type ChunkRangesRef = range_collections::RangeSetRef; -/// Compile time hashing strategy for shared encode and decode paths. -/// -/// Use the standard or `keyed_*` public APIs rather than this trait directly. -#[doc(hidden)] -pub trait BaoHashing: Copy { - /// Hash a subtree of chunk data. - fn hash_subtree(&self, start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash; - /// Combine two child chaining values into a parent chaining value. - fn parent_cv( - &self, - left_child: &blake3::Hash, - right_child: &blake3::Hash, - is_root: bool, - ) -> blake3::Hash; +/// Hashing mode for shared encode and decode paths, either standard or keyed BLAKE3. +#[derive(Clone, Copy)] +pub(crate) enum HashMode { + Standard, + Keyed([u8; 32]), } -/// BLAKE3 hash mode strategy for unkeyed APIs. -#[doc(hidden)] -#[derive(Debug, Clone, Copy)] -pub struct Standard; - -impl BaoHashing for Standard { - fn hash_subtree(&self, start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash { - hash_subtree(start_chunk, data, is_root) +impl std::fmt::Debug for HashMode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + HashMode::Standard => f.write_str("Standard"), + HashMode::Keyed(_) => f.debug_struct("Keyed").finish_non_exhaustive(), + } } +} - fn parent_cv( +impl HashMode { + pub(crate) fn hash_subtree( &self, - left_child: &blake3::Hash, - right_child: &blake3::Hash, + start_chunk: u64, + data: &[u8], is_root: bool, ) -> blake3::Hash { - parent_cv(left_child, right_child, is_root) - } -} - -/// BLAKE3 keyed mode strategy. Wraps a 32 byte key for domain separated hashing. -#[doc(hidden)] -#[derive(Clone, Copy)] -pub struct Keyed(pub [u8; 32]); - -impl std::fmt::Debug for Keyed { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Keyed").finish_non_exhaustive() - } -} - -impl BaoHashing for Keyed { - fn hash_subtree(&self, start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash { - keyed_hash_subtree(start_chunk, data, is_root, &self.0) + match self { + HashMode::Standard => hash_subtree(start_chunk, data, is_root), + HashMode::Keyed(key) => keyed_hash_subtree(start_chunk, data, is_root, key), + } } - fn parent_cv( + pub(crate) fn parent_cv( &self, left_child: &blake3::Hash, right_child: &blake3::Hash, is_root: bool, ) -> blake3::Hash { - keyed_parent_cv(left_child, right_child, is_root, &self.0) + match self { + HashMode::Standard => parent_cv(left_child, right_child, is_root), + HashMode::Keyed(key) => keyed_parent_cv(left_child, right_child, is_root, key), + } } } diff --git a/src/rec.rs b/src/rec.rs index fc1c0a6..a5fa8e4 100644 --- a/src/rec.rs +++ b/src/rec.rs @@ -97,9 +97,9 @@ fn truncated_len(ranges: &ChunkRangesRef, size: u64) -> usize { /// below the chunk group size when creating responses for outboards with a chunk group /// size of >0. /// -/// `hash_strategy` is the compile time hashing mode for subtree and parent hashes. -#[allow(clippy::too_many_arguments)] // keyed mode adds `hash_strategy`; splitting into a struct isn't worth it here -pub(crate) fn encode_selected_rec( +/// `mode` is the hashing mode for subtree and parent hashes. +#[allow(clippy::too_many_arguments)] +pub(crate) fn encode_selected_rec( start_chunk: ChunkNum, data: &[u8], is_root: bool, @@ -107,14 +107,14 @@ pub(crate) fn encode_selected_rec( min_level: u32, emit_data: bool, res: &mut Vec, - hash_strategy: H, + mode: crate::HashMode, ) -> blake3::Hash { use blake3::CHUNK_LEN; if data.len() <= CHUNK_LEN { if emit_data && !query.is_empty() { res.extend_from_slice(data); } - hash_strategy.hash_subtree(start_chunk.0, data, is_root) + mode.hash_subtree(start_chunk.0, data, is_root) } else { let chunks = data.len() / CHUNK_LEN + (data.len() % CHUNK_LEN != 0) as usize; let chunks = chunks.next_power_of_two(); @@ -146,7 +146,7 @@ pub(crate) fn encode_selected_rec( min_level, emit_data, res, - hash_strategy, + mode, ); let right = encode_selected_rec( mid_chunk, @@ -156,14 +156,14 @@ pub(crate) fn encode_selected_rec( min_level, emit_data, res, - hash_strategy, + mode, ); // backfill the hashes if needed if let Some(o) = hash_offset { res[o..o + 32].copy_from_slice(left.as_bytes()); res[o + 32..o + 64].copy_from_slice(right.as_bytes()); } - hash_strategy.parent_cv(&left, &right, is_root) + mode.parent_cv(&left, &right, is_root) } } @@ -281,7 +281,7 @@ mod test_support { 0, false, &mut res, - crate::Standard, + crate::HashMode::Standard, ); (res, hash) } @@ -297,7 +297,7 @@ mod test_support { 0, true, &mut res, - crate::Standard, + crate::HashMode::Standard, ); (res, hash) } @@ -438,7 +438,7 @@ mod test_support { block_size.to_u32(), true, &mut res, - crate::Standard, + crate::HashMode::Standard, ); (res, hash) } diff --git a/src/tests.rs b/src/tests.rs index 08faab9..333b34d 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -33,7 +33,7 @@ use crate::{ keyed_init_from_keyed_checks, keyed_outboard_functions_checks, make_test_data, range_union, truncate_ranges, ReferencePreOrderPartialChunkIterRef, }, - split, ChunkRanges, ChunkRangesRef, Keyed, ResponseIter, Standard, + split, ChunkRanges, ChunkRangesRef, HashMode, ResponseIter, }; #[cfg(feature = "tokio_fsm")] @@ -42,7 +42,7 @@ use crate::rec::{ keyed_outboard_functions_checks_fsm, }; -/// Reference encoder using the [Keyed] hashing strategy. +/// Reference encoder using BLAKE3 keyed mode. fn keyed_encode_selected_reference( data: &[u8], block_size: BlockSize, @@ -60,7 +60,7 @@ fn keyed_encode_selected_reference( max_skip_level, true, &mut res, - Keyed(*key), + HashMode::Keyed(*key), ); (hash, res) } @@ -1045,7 +1045,7 @@ fn encode_selected_rec_cases() { min_level, true, &mut actual_encoded, - Standard, + HashMode::Standard, ); actual_encoded.len() - data.len() }; @@ -1071,7 +1071,7 @@ fn encode_selected_reference( max_skip_level, true, &mut res, - Standard, + HashMode::Standard, ); (hash, res) } diff --git a/src/tests2.rs b/src/tests2.rs index 616d9b8..701b07e 100644 --- a/src/tests2.rs +++ b/src/tests2.rs @@ -32,7 +32,7 @@ use crate::{ partial_chunk_iter_reference, range_union, response_iter_reference, select_nodes_rec, truncate_ranges, ReferencePreOrderPartialChunkIterRef, }, - BaoTree, BlockSize, ChunkNum, ChunkRanges, ChunkRangesRef, Standard, TreeNode, + BaoTree, BlockSize, ChunkNum, ChunkRanges, ChunkRangesRef, HashMode, TreeNode, }; fn keyed_test_key(context: &[u8]) -> [u8; 32] { @@ -1300,7 +1300,7 @@ fn selection_reference_comparison_proptest( /// Reference implementation of encode_ranges_validated that uses the simple recursive impl. /// -/// Uses the [Standard] hashing strategy for unkeyed BLAKE3 mode. +/// Uses unkeyed BLAKE3 mode. fn encode_selected_reference( data: &[u8], block_size: BlockSize, @@ -1317,7 +1317,7 @@ fn encode_selected_reference( max_skip_level, true, &mut res, - Standard, + HashMode::Standard, ); (hash, res) } From ac996e625a17bd3de939880d6afc016085914730 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Aug 2026 11:08:05 +0200 Subject: [PATCH 07/12] Extend validator API for keyed outboards --- src/io/fsm.rs | 38 +++++++++++++++++++++-- src/io/sync.rs | 38 +++++++++++++++++++++-- src/tests2.rs | 81 ++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 151 insertions(+), 6 deletions(-) diff --git a/src/io/fsm.rs b/src/io/fsm.rs index 05c3d1f..f24155c 100644 --- a/src/io/fsm.rs +++ b/src/io/fsm.rs @@ -1176,11 +1176,38 @@ mod validate { outboard: O, ranges: &'a ChunkRangesRef, ) -> impl Stream>> + 'a + where + O: Outboard + 'a, + { + valid_outboard_ranges_impl(outboard, ranges, HashMode::Standard) + } + + /// Given just a keyed outboard, compute all valid ranges. + /// + /// This is not cheap since it recomputes the hashes for all chunks. + pub fn keyed_valid_outboard_ranges<'a, O>( + outboard: O, + ranges: &'a ChunkRangesRef, + key: &[u8; 32], + ) -> impl Stream>> + 'a + where + O: Outboard + 'a, + { + valid_outboard_ranges_impl(outboard, ranges, HashMode::Keyed(*key)) + } + + fn valid_outboard_ranges_impl<'a, O>( + outboard: O, + ranges: &'a ChunkRangesRef, + mode: HashMode, + ) -> impl Stream>> + 'a where O: Outboard + 'a, { Gen::new(move |co| async move { - if let Err(cause) = RecursiveOutboardValidator::validate(outboard, ranges, &co).await { + if let Err(cause) = + RecursiveOutboardValidator::validate(outboard, ranges, &co, mode).await + { co.yield_(Err(cause)).await; } }) @@ -1191,6 +1218,7 @@ mod validate { shifted_filled_size: TreeNode, outboard: O, co: &'a Co>>, + mode: HashMode, } impl RecursiveOutboardValidator<'_, O> { @@ -1198,6 +1226,7 @@ mod validate { outboard: O, ranges: &ChunkRangesRef, co: &Co>>, + mode: HashMode, ) -> io::Result<()> { let tree = outboard.tree(); if tree.blocks() == 1 { @@ -1213,6 +1242,7 @@ mod validate { shifted_filled_size, outboard, co, + mode, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -1246,7 +1276,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = HashMode::Standard.parent_cv(&l_hash, &r_hash, is_root); + let actual = self.mode.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -1272,4 +1302,6 @@ mod validate { } } #[cfg(feature = "validate")] -pub use validate::{keyed_valid_ranges, valid_outboard_ranges, valid_ranges}; +pub use validate::{ + keyed_valid_outboard_ranges, keyed_valid_ranges, valid_outboard_ranges, valid_ranges, +}; diff --git a/src/io/sync.rs b/src/io/sync.rs index 6897332..d44effb 100644 --- a/src/io/sync.rs +++ b/src/io/sync.rs @@ -1054,11 +1054,38 @@ mod validate { outboard: O, ranges: &'a ChunkRangesRef, ) -> impl IntoIterator>> + 'a + where + O: Outboard + 'a, + { + valid_outboard_ranges_impl(outboard, ranges, HashMode::Standard) + } + + /// Given just a keyed outboard, compute all valid ranges. + /// + /// This is not cheap since it recomputes the hashes for all chunks. + pub fn keyed_valid_outboard_ranges<'a, O>( + outboard: O, + ranges: &'a ChunkRangesRef, + key: &[u8; 32], + ) -> impl IntoIterator>> + 'a + where + O: Outboard + 'a, + { + valid_outboard_ranges_impl(outboard, ranges, HashMode::Keyed(*key)) + } + + fn valid_outboard_ranges_impl<'a, O>( + outboard: O, + ranges: &'a ChunkRangesRef, + mode: HashMode, + ) -> impl IntoIterator>> + 'a where O: Outboard + 'a, { Gen::new(move |co| async move { - if let Err(cause) = RecursiveOutboardValidator::validate(outboard, ranges, &co).await { + if let Err(cause) = + RecursiveOutboardValidator::validate(outboard, ranges, &co, mode).await + { co.yield_(Err(cause)).await; } }) @@ -1069,6 +1096,7 @@ mod validate { shifted_filled_size: TreeNode, outboard: O, co: &'a Co>>, + mode: HashMode, } impl RecursiveOutboardValidator<'_, O> { @@ -1076,6 +1104,7 @@ mod validate { outboard: O, ranges: &ChunkRangesRef, co: &Co>>, + mode: HashMode, ) -> io::Result<()> { let tree = outboard.tree(); if tree.blocks() == 1 { @@ -1091,6 +1120,7 @@ mod validate { shifted_filled_size, outboard, co, + mode, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -1124,7 +1154,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = HashMode::Standard.parent_cv(&l_hash, &r_hash, is_root); + let actual = self.mode.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -1150,4 +1180,6 @@ mod validate { } } #[cfg(feature = "validate")] -pub use validate::{keyed_valid_ranges, valid_outboard_ranges, valid_ranges}; +pub use validate::{ + keyed_valid_outboard_ranges, keyed_valid_ranges, valid_outboard_ranges, valid_ranges, +}; diff --git a/src/tests2.rs b/src/tests2.rs index 701b07e..3309847 100644 --- a/src/tests2.rs +++ b/src/tests2.rs @@ -515,6 +515,87 @@ mod validate { } } + fn keyed_valid_outboard_ranges_sync( + outboard: impl crate::io::sync::Outboard, + key: &[u8; 32], + ) -> ChunkRanges { + let ranges = ChunkRanges::all(); + let iter = crate::io::sync::keyed_valid_outboard_ranges(outboard, &ranges, key); + let mut res = ChunkRanges::empty(); + for item in iter { + res |= ChunkRanges::from(item.unwrap()); + } + res + } + + fn keyed_valid_outboard_ranges_fsm( + outboard: &mut PostOrderMemOutboard, + key: &[u8; 32], + ) -> ChunkRanges { + run_blocking(async move { + let ranges = ChunkRanges::all(); + let mut stream = crate::io::fsm::keyed_valid_outboard_ranges(outboard, &ranges, key); + let mut res = ChunkRanges::empty(); + while let Some(item) = stream.next().await { + let item = item?; + res |= ChunkRanges::from(item); + } + std::io::Result::Ok(res) + }) + .unwrap() + } + + fn validate_keyed_outboard_pos_impl(tree: BaoTree) { + let size = tree.size.try_into().unwrap(); + let block_size = tree.block_size; + let data = make_test_data(size); + let key = blake3::derive_key("bao-tree.test", b"valid-outboard-ranges"); + let mut outboard = PostOrderMemOutboard::create_keyed(data, block_size, &key); + let expected = ChunkRanges::from(..outboard.tree().chunks()); + let actual = keyed_valid_outboard_ranges_sync(&mut outboard, &key); + assert_eq!(expected, actual); + let actual = keyed_valid_outboard_ranges_fsm(&mut outboard, &key); + assert_eq!(expected, actual) + } + + #[proptest] + fn validate_keyed_outboard_pos_proptest(#[strategy(tree())] tree: BaoTree) { + validate_keyed_outboard_pos_impl(tree); + } + + #[test] + fn validate_keyed_outboard_pos_cases() { + let cases = [(0x10001, 0)]; + for (size, block_level) in cases { + let tree = BaoTree::new(size, BlockSize(block_level)); + validate_keyed_outboard_pos_impl(tree); + } + } + + /// Wrong key must not report a multi-block keyed outboard as fully valid. + fn validate_keyed_outboard_wrong_key_impl(tree: BaoTree) { + let size = tree.size.try_into().unwrap(); + let block_size = tree.block_size; + let data = make_test_data(size); + let key = blake3::derive_key("bao-tree.test", b"valid-outboard-ranges"); + let wrong_key = blake3::derive_key("bao-tree.test", b"wrong-key"); + let mut outboard = PostOrderMemOutboard::create_keyed(data, block_size, &key); + let expected = ChunkRanges::from(..outboard.tree().chunks()); + let actual = keyed_valid_outboard_ranges_sync(&mut outboard, &wrong_key); + assert_ne!(expected, actual); + let actual = keyed_valid_outboard_ranges_fsm(&mut outboard, &wrong_key); + assert_ne!(expected, actual) + } + + #[test] + fn validate_keyed_outboard_wrong_key_cases() { + let cases = [(0x10001, 0), (0x2001, 0), (5000, 1)]; + for (size, block_level) in cases { + let tree = BaoTree::new(size, BlockSize(block_level)); + validate_keyed_outboard_wrong_key_impl(tree); + } + } + fn validate_pos_impl(tree: BaoTree) { let size = tree.size.try_into().unwrap(); let block_size = tree.block_size; From 59f2cf1d5f1da16ec0d7256f39e8eb408a684915 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Aug 2026 11:27:46 +0200 Subject: [PATCH 08/12] Remove keyed fns from CreateOutboard keyed is a rather niche use case and I don't want to burden CreateOutboard implementers with this. We can add them later once/if we have lots of keyed outboard users. --- src/io/fsm.rs | 89 ------------- src/io/mixed.rs | 2 +- src/io/sync.rs | 81 ----------- src/lib.rs | 88 ++++-------- src/rec.rs | 348 +++++++----------------------------------------- src/tests.rs | 189 ++------------------------ src/tests2.rs | 53 +------- 7 files changed, 93 insertions(+), 757 deletions(-) diff --git a/src/io/fsm.rs b/src/io/fsm.rs index f24155c..3947656 100644 --- a/src/io/fsm.rs +++ b/src/io/fsm.rs @@ -128,37 +128,6 @@ pub trait CreateOutboard { /// /// It will only include data up the the current tree size. fn init_from(&mut self, data: impl AsyncStreamReader) -> impl Future>; - - /// Create a keyed outboard from a seekable data source. - #[allow(async_fn_in_trait)] - async fn create_keyed( - mut data: impl AsyncSliceReader, - block_size: BlockSize, - key: &[u8; 32], - ) -> io::Result - where - Self: Default + Sized, - { - let size = data.size().await?; - Self::create_sized_keyed(Cursor::new(data), size, block_size, key).await - } - - /// Create a keyed outboard from a data source with a known size. - fn create_sized_keyed( - data: impl AsyncStreamReader, - size: u64, - block_size: BlockSize, - key: &[u8; 32], - ) -> impl Future> - where - Self: Default + Sized; - - /// Init a keyed outboard from a data source. - fn init_from_keyed( - &mut self, - data: impl AsyncStreamReader, - key: &[u8; 32], - ) -> impl Future>; } impl Outboard for &mut O { @@ -280,35 +249,6 @@ impl CreateOutboard for PreOrderOutboard { this.sync().await?; Ok(()) } - - async fn create_sized_keyed( - data: impl AsyncStreamReader, - size: u64, - block_size: BlockSize, - key: &[u8; 32], - ) -> io::Result - where - Self: Default + Sized, - { - let mut res = Self { - tree: BaoTree::new(size, block_size), - ..Self::default() - }; - res.init_from_keyed(data, key).await?; - Ok(res) - } - - async fn init_from_keyed( - &mut self, - data: impl AsyncStreamReader, - key: &[u8; 32], - ) -> io::Result<()> { - let mut this = self; - let root = keyed_outboard(data, this.tree, &mut this, key).await?; - this.root = root; - this.sync().await?; - Ok(()) - } } impl CreateOutboard for PostOrderOutboard { @@ -335,35 +275,6 @@ impl CreateOutboard for PostOrderOutboard { this.sync().await?; Ok(()) } - - async fn create_sized_keyed( - data: impl AsyncStreamReader, - size: u64, - block_size: BlockSize, - key: &[u8; 32], - ) -> io::Result - where - Self: Default + Sized, - { - let mut res = Self { - tree: BaoTree::new(size, block_size), - ..Self::default() - }; - res.init_from_keyed(data, key).await?; - Ok(res) - } - - async fn init_from_keyed( - &mut self, - data: impl AsyncStreamReader, - key: &[u8; 32], - ) -> io::Result<()> { - let mut this = self; - let root = keyed_outboard(data, this.tree, &mut this, key).await?; - this.root = root; - this.sync().await?; - Ok(()) - } } impl Outboard for PostOrderOutboard { diff --git a/src/io/mixed.rs b/src/io/mixed.rs index 363237a..ed09e48 100644 --- a/src/io/mixed.rs +++ b/src/io/mixed.rs @@ -227,7 +227,7 @@ where /// This will compute the root hash, so it will have to traverse the entire tree. /// The `ranges` parameter just controls which parts of the data are written. /// -/// Except for writing to a buffer, this is the same as [crate::hash_subtree]. +/// Except for writing to a buffer, this is the same as computing the subtree hash. /// The `min_level` parameter controls the minimum level that will be emitted as a leaf. /// Set this to 0 to disable chunk groups entirely. /// The `emit_data` parameter controls whether the data is written to the buffer. diff --git a/src/io/sync.rs b/src/io/sync.rs index d44effb..59c01d9 100644 --- a/src/io/sync.rs +++ b/src/io/sync.rs @@ -96,33 +96,6 @@ pub trait CreateOutboard { /// /// It will only include data up the the current tree size. fn init_from(&mut self, data: impl Read) -> io::Result<()>; - - /// Create a keyed outboard from a data source. - fn create_keyed( - mut data: impl Read + Seek, - block_size: BlockSize, - key: &[u8; 32], - ) -> io::Result - where - Self: Default + Sized, - { - let size = data.seek(io::SeekFrom::End(0))?; - data.rewind()?; - Self::create_sized_keyed(data, size, block_size, key) - } - - /// Create a keyed outboard from a data source with a known size. - fn create_sized_keyed( - data: impl Read, - size: u64, - block_size: BlockSize, - key: &[u8; 32], - ) -> io::Result - where - Self: Default + Sized; - - /// Init a keyed outboard from a data source. - fn init_from_keyed(&mut self, data: impl Read, key: &[u8; 32]) -> io::Result<()>; } impl OutboardMut for &mut O { @@ -219,33 +192,6 @@ impl CreateOutboard for PreOrderOutboard { this.sync()?; Ok(()) } - - fn create_sized_keyed( - data: impl Read, - size: u64, - block_size: BlockSize, - key: &[u8; 32], - ) -> io::Result - where - Self: Default + Sized, - { - let tree = BaoTree::new(size, block_size); - let mut res = Self { - tree, - ..Default::default() - }; - res.init_from_keyed(data, key)?; - res.sync()?; - Ok(res) - } - - fn init_from_keyed(&mut self, data: impl Read, key: &[u8; 32]) -> io::Result<()> { - let mut this = self; - let root = keyed_outboard(data, this.tree, &mut this, key)?; - this.root = root; - this.sync()?; - Ok(()) - } } impl CreateOutboard for PostOrderOutboard { @@ -270,33 +216,6 @@ impl CreateOutboard for PostOrderOutboard { this.sync()?; Ok(()) } - - fn create_sized_keyed( - data: impl Read, - size: u64, - block_size: BlockSize, - key: &[u8; 32], - ) -> io::Result - where - Self: Default + Sized, - { - let tree = BaoTree::new(size, block_size); - let mut res = Self { - tree, - ..Default::default() - }; - res.init_from_keyed(data, key)?; - res.sync()?; - Ok(res) - } - - fn init_from_keyed(&mut self, data: impl Read, key: &[u8; 32]) -> io::Result<()> { - let mut this = self; - let root = keyed_outboard(data, this.tree, &mut this, key)?; - this.root = root; - this.sync()?; - Ok(()) - } } impl OutboardMut for PostOrderOutboard { diff --git a/src/lib.rs b/src/lib.rs index 57ccfe7..f1299f6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -262,9 +262,22 @@ impl HashMode { data: &[u8], is_root: bool, ) -> blake3::Hash { - match self { - HashMode::Standard => hash_subtree(start_chunk, data, is_root), - HashMode::Keyed(key) => keyed_hash_subtree(start_chunk, data, is_root, key), + use blake3::hazmat::{ChainingValue, HasherExt}; + if is_root { + debug_assert!(start_chunk == 0); + match self { + HashMode::Standard => blake3::hash(data), + HashMode::Keyed(key) => blake3::keyed_hash(key, data), + } + } else { + let mut hasher = match self { + HashMode::Standard => blake3::Hasher::new(), + HashMode::Keyed(key) => blake3::Hasher::new_keyed(key), + }; + hasher.set_input_offset(start_chunk * 1024); + hasher.update(data); + let non_root_hash: ChainingValue = hasher.finalize_non_root(); + blake3::Hash::from(non_root_hash) } } @@ -274,27 +287,21 @@ impl HashMode { right_child: &blake3::Hash, is_root: bool, ) -> blake3::Hash { - match self { - HashMode::Standard => parent_cv(left_child, right_child, is_root), - HashMode::Keyed(key) => keyed_parent_cv(left_child, right_child, is_root, key), + use blake3::hazmat::{merge_subtrees_non_root, merge_subtrees_root, ChainingValue, Mode}; + let left_child: ChainingValue = *left_child.as_bytes(); + let right_child: ChainingValue = *right_child.as_bytes(); + let mode = match self { + HashMode::Standard => Mode::Hash, + HashMode::Keyed(key) => Mode::KeyedHash(key), + }; + if is_root { + merge_subtrees_root(&left_child, &right_child, mode) + } else { + blake3::Hash::from(merge_subtrees_non_root(&left_child, &right_child, mode)) } } } -pub(crate) fn hash_subtree(start_chunk: u64, data: &[u8], is_root: bool) -> blake3::Hash { - use blake3::hazmat::{ChainingValue, HasherExt}; - if is_root { - debug_assert!(start_chunk == 0); - blake3::hash(data) - } else { - let mut hasher = blake3::Hasher::new(); - hasher.set_input_offset(start_chunk * 1024); - hasher.update(data); - let non_root_hash: ChainingValue = hasher.finalize_non_root(); - blake3::Hash::from(non_root_hash) - } -} - /// Compute the hash of a subtree using BLAKE3 keyed mode. /// /// See [keyed_parent_cv] for merging child hashes in keyed mode. @@ -304,36 +311,7 @@ pub fn keyed_hash_subtree( is_root: bool, key: &[u8; 32], ) -> blake3::Hash { - use blake3::hazmat::{ChainingValue, HasherExt}; - if is_root { - debug_assert!(start_chunk == 0); - blake3::keyed_hash(key, data) - } else { - let mut hasher = blake3::Hasher::new_keyed(key); - hasher.set_input_offset(start_chunk * 1024); - hasher.update(data); - let non_root_hash: ChainingValue = hasher.finalize_non_root(); - blake3::Hash::from(non_root_hash) - } -} - -pub(crate) fn parent_cv( - left_child: &blake3::Hash, - right_child: &blake3::Hash, - is_root: bool, -) -> blake3::Hash { - use blake3::hazmat::{merge_subtrees_non_root, merge_subtrees_root, ChainingValue, Mode}; - let left_child: ChainingValue = *left_child.as_bytes(); - let right_child: ChainingValue = *right_child.as_bytes(); - if is_root { - merge_subtrees_root(&left_child, &right_child, Mode::Hash) - } else { - blake3::Hash::from(merge_subtrees_non_root( - &left_child, - &right_child, - Mode::Hash, - )) - } + HashMode::Keyed(*key).hash_subtree(start_chunk, data, is_root) } /// Merge two child subtree hashes using BLAKE3 keyed mode. @@ -343,15 +321,7 @@ pub fn keyed_parent_cv( is_root: bool, key: &[u8; 32], ) -> blake3::Hash { - use blake3::hazmat::{merge_subtrees_non_root, merge_subtrees_root, ChainingValue, Mode}; - let left_child: ChainingValue = *left_child.as_bytes(); - let right_child: ChainingValue = *right_child.as_bytes(); - let mode = Mode::KeyedHash(key); - if is_root { - merge_subtrees_root(&left_child, &right_child, mode) - } else { - blake3::Hash::from(merge_subtrees_non_root(&left_child, &right_child, mode)) - } + HashMode::Keyed(*key).parent_cv(left_child, right_child, is_root) } /// Defines a Bao tree. diff --git a/src/rec.rs b/src/rec.rs index a5fa8e4..80d7e16 100644 --- a/src/rec.rs +++ b/src/rec.rs @@ -445,28 +445,8 @@ mod test_support { use std::io::Cursor; - use crate::io::outboard::{ - PostOrderMemOutboard, PostOrderOutboard, PreOrderMemOutboard, PreOrderOutboard, - }; - use crate::io::sync::{self, CreateOutboard, Outboard}; - - pub(crate) fn assert_post_order_outboard_matches_mem( - outboard: &PostOrderOutboard>, - data: &[u8], - block_size: BlockSize, - key: &[u8; 32], - ) { - let reference = PostOrderMemOutboard::create_keyed(data, block_size, key); - assert_eq!(outboard.root, reference.root); - let tree = outboard.tree; - let mut copied = PostOrderMemOutboard { - root: outboard.root, - tree, - data: vec![0; tree.outboard_hash_pairs() as usize * 64], - }; - sync::copy(outboard, &mut copied).unwrap(); - assert_eq!(copied.data, reference.data); - } + use crate::io::outboard::{PostOrderMemOutboard, PreOrderMemOutboard, PreOrderOutboard}; + use crate::io::sync::{self}; pub(crate) fn assert_pre_order_outboard_matches_mem( outboard: &PreOrderOutboard>, @@ -486,157 +466,6 @@ mod test_support { assert_eq!(copied.data, reference.data); } - fn assert_truncated_create_sized_keyed_post( - truncated: &PostOrderOutboard>, - data: &[u8], - truncated_size: u64, - block_size: BlockSize, - key: &[u8; 32], - ) { - assert_eq!(truncated.tree.size, truncated_size); - assert_eq!( - truncated.root(), - blake3::keyed_hash(key, &data[..truncated_size as usize]) - ); - assert_post_order_outboard_matches_mem( - truncated, - &data[..truncated_size as usize], - block_size, - key, - ); - } - - fn assert_truncated_create_sized_keyed_pre( - truncated: &PreOrderOutboard>, - data: &[u8], - truncated_size: u64, - block_size: BlockSize, - key: &[u8; 32], - ) { - assert_eq!(truncated.tree.size, truncated_size); - assert_eq!( - truncated.root(), - blake3::keyed_hash(key, &data[..truncated_size as usize]) - ); - assert_pre_order_outboard_matches_mem( - truncated, - &data[..truncated_size as usize], - block_size, - key, - ); - } - - pub(crate) fn keyed_create_sized_keyed_checks( - data: &[u8], - block_size: BlockSize, - key: &[u8; 32], - ) { - let size = data.len() as u64; - - let post: PostOrderOutboard> = - PostOrderOutboard::create_sized_keyed(Cursor::new(data), size, block_size, key) - .unwrap(); - assert_post_order_outboard_matches_mem(&post, data, block_size, key); - - let pre: PreOrderOutboard> = - PreOrderOutboard::create_sized_keyed(Cursor::new(data), size, block_size, key).unwrap(); - assert_pre_order_outboard_matches_mem(&pre, data, block_size, key); - - let truncated_size = 1024u64.min(size); - if truncated_size < size { - let truncated_post: PostOrderOutboard> = PostOrderOutboard::create_sized_keyed( - Cursor::new(data), - truncated_size, - BlockSize(0), - key, - ) - .unwrap(); - assert_truncated_create_sized_keyed_post( - &truncated_post, - data, - truncated_size, - BlockSize(0), - key, - ); - - let truncated_pre: PreOrderOutboard> = PreOrderOutboard::create_sized_keyed( - Cursor::new(data), - truncated_size, - BlockSize(0), - key, - ) - .unwrap(); - assert_truncated_create_sized_keyed_pre( - &truncated_pre, - data, - truncated_size, - BlockSize(0), - key, - ); - } - } - - pub(crate) fn keyed_init_from_keyed_checks(data: &[u8], block_size: BlockSize, key: &[u8; 32]) { - let tree = BaoTree::new(data.len() as u64, block_size); - let expected = blake3::keyed_hash(key, data); - - let mut post = PostOrderOutboard { - root: blake3::Hash::from([0; 32]), - tree, - data: vec![0; tree.outboard_size().try_into().unwrap()], - }; - post.init_from_keyed(Cursor::new(data), key).unwrap(); - assert_eq!(post.root(), expected); - assert_post_order_outboard_matches_mem(&post, data, block_size, key); - - let mut pre = PreOrderOutboard { - root: blake3::Hash::from([0; 32]), - tree, - data: vec![0; tree.outboard_size().try_into().unwrap()], - }; - pre.init_from_keyed(Cursor::new(data), key).unwrap(); - assert_eq!(pre.root(), expected); - assert_pre_order_outboard_matches_mem(&pre, data, block_size, key); - - let truncated_size = 1024u64.min(data.len() as u64); - if truncated_size < data.len() as u64 { - let truncated_tree = BaoTree::new(truncated_size, BlockSize(0)); - let truncated_expected = blake3::keyed_hash(key, &data[..truncated_size as usize]); - - let mut truncated_post = PostOrderOutboard { - root: blake3::Hash::from([0; 32]), - tree: truncated_tree, - data: vec![0; truncated_tree.outboard_size().try_into().unwrap()], - }; - truncated_post - .init_from_keyed(Cursor::new(data), key) - .unwrap(); - assert_eq!(truncated_post.root(), truncated_expected); - assert_post_order_outboard_matches_mem( - &truncated_post, - &data[..truncated_size as usize], - BlockSize(0), - key, - ); - - let mut truncated_pre = PreOrderOutboard { - root: blake3::Hash::from([0; 32]), - tree: truncated_tree, - data: vec![0; truncated_tree.outboard_size().try_into().unwrap()], - }; - truncated_pre - .init_from_keyed(Cursor::new(data), key) - .unwrap(); - assert_eq!(truncated_pre.root(), truncated_expected); - assert_pre_order_outboard_matches_mem( - &truncated_pre, - &data[..truncated_size as usize], - BlockSize(0), - key, - ); - } - } - pub(crate) fn keyed_outboard_functions_checks( data: &[u8], block_size: BlockSize, @@ -671,76 +500,37 @@ mod test_support { }; let pre_from_post = post_mem.flip(); assert_eq!(pre_from_post.data, pre.data); - } - - #[cfg(feature = "tokio_fsm")] - pub(crate) async fn keyed_create_sized_keyed_checks_fsm( - data: &[u8], - block_size: BlockSize, - key: &[u8; 32], - ) { - use bytes::Bytes; - - let size = data.len() as u64; - let post: PostOrderOutboard> = - > as crate::io::fsm::CreateOutboard>::create_sized_keyed( - Cursor::new(Bytes::from(data.to_vec())), - size, - block_size, - key, - ) - .await - .unwrap(); - assert_post_order_outboard_matches_mem(&post, data, block_size, key); - - let pre: PreOrderOutboard> = - > as crate::io::fsm::CreateOutboard>::create_sized_keyed( - Cursor::new(Bytes::from(data.to_vec())), - size, - block_size, - key, - ) - .await - .unwrap(); - assert_pre_order_outboard_matches_mem(&pre, data, block_size, key); - - let truncated_size = 1024u64.min(size); - if truncated_size < size { - let truncated_post: PostOrderOutboard> = - > as crate::io::fsm::CreateOutboard>::create_sized_keyed( - Cursor::new(Bytes::from(data.to_vec())), - truncated_size, - BlockSize(0), - key, - ) - .await - .unwrap(); - assert_truncated_create_sized_keyed_post( - &truncated_post, - data, - truncated_size, - BlockSize(0), - key, + let truncated_size = 1024u64.min(data.len() as u64); + if truncated_size < data.len() as u64 { + let truncated_tree = BaoTree::new(truncated_size, BlockSize(0)); + let mut truncated_pre = PreOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree: truncated_tree, + data: vec![0; truncated_tree.outboard_size().try_into().unwrap()], + }; + let root = + sync::keyed_outboard(Cursor::new(data), truncated_tree, &mut truncated_pre, key) + .unwrap(); + truncated_pre.root = root; + assert_eq!( + root, + blake3::keyed_hash(key, &data[..truncated_size as usize]) ); - - let truncated_pre: PreOrderOutboard> = - > as crate::io::fsm::CreateOutboard>::create_sized_keyed( - Cursor::new(Bytes::from(data.to_vec())), - truncated_size, - BlockSize(0), - key, - ) - .await - .unwrap(); - assert_truncated_create_sized_keyed_pre( + assert_pre_order_outboard_matches_mem( &truncated_pre, - data, - truncated_size, + &data[..truncated_size as usize], BlockSize(0), key, ); } + + let oversize_tree = BaoTree::new(data.len() as u64 + 100, BlockSize(0)); + let mut sink = Vec::new(); + assert!( + sync::keyed_outboard_post_order(Cursor::new(data), oversize_tree, &mut sink, key) + .is_err() + ); } #[cfg(feature = "tokio_fsm")] @@ -789,87 +579,28 @@ mod test_support { }; let pre_from_post = post_mem.flip(); assert_eq!(pre_from_post.data, pre.data); - } - - #[cfg(feature = "tokio_fsm")] - pub(crate) async fn keyed_init_from_keyed_checks_fsm( - data: &[u8], - block_size: BlockSize, - key: &[u8; 32], - ) { - use bytes::Bytes; - - let tree = BaoTree::new(data.len() as u64, block_size); - let expected = blake3::keyed_hash(key, data); - - let mut post = PostOrderOutboard { - root: blake3::Hash::from([0; 32]), - tree, - data: vec![0; tree.outboard_size().try_into().unwrap()], - }; - crate::io::fsm::CreateOutboard::init_from_keyed( - &mut post, - Cursor::new(Bytes::from(data.to_vec())), - key, - ) - .await - .unwrap(); - assert_eq!(post.root(), expected); - assert_post_order_outboard_matches_mem(&post, data, block_size, key); - - let mut pre = PreOrderOutboard { - root: blake3::Hash::from([0; 32]), - tree, - data: vec![0; tree.outboard_size().try_into().unwrap()], - }; - crate::io::fsm::CreateOutboard::init_from_keyed( - &mut pre, - Cursor::new(Bytes::from(data.to_vec())), - key, - ) - .await - .unwrap(); - assert_eq!(pre.root(), expected); - assert_pre_order_outboard_matches_mem(&pre, data, block_size, key); let truncated_size = 1024u64.min(data.len() as u64); if truncated_size < data.len() as u64 { let truncated_tree = BaoTree::new(truncated_size, BlockSize(0)); - let truncated_expected = blake3::keyed_hash(key, &data[..truncated_size as usize]); - - let mut truncated_post = PostOrderOutboard { - root: blake3::Hash::from([0; 32]), - tree: truncated_tree, - data: vec![0; truncated_tree.outboard_size().try_into().unwrap()], - }; - crate::io::fsm::CreateOutboard::init_from_keyed( - &mut truncated_post, - Cursor::new(Bytes::from(data.to_vec())), - key, - ) - .await - .unwrap(); - assert_eq!(truncated_post.root(), truncated_expected); - assert_post_order_outboard_matches_mem( - &truncated_post, - &data[..truncated_size as usize], - BlockSize(0), - key, - ); - let mut truncated_pre = PreOrderOutboard { root: blake3::Hash::from([0; 32]), tree: truncated_tree, data: vec![0; truncated_tree.outboard_size().try_into().unwrap()], }; - crate::io::fsm::CreateOutboard::init_from_keyed( - &mut truncated_pre, + let root = keyed_outboard( Cursor::new(Bytes::from(data.to_vec())), + truncated_tree, + &mut truncated_pre, key, ) .await .unwrap(); - assert_eq!(truncated_pre.root(), truncated_expected); + truncated_pre.root = root; + assert_eq!( + root, + blake3::keyed_hash(key, &data[..truncated_size as usize]) + ); assert_pre_order_outboard_matches_mem( &truncated_pre, &data[..truncated_size as usize], @@ -877,6 +608,17 @@ mod test_support { key, ); } + + let oversize_tree = BaoTree::new(data.len() as u64 + 100, BlockSize(0)); + let mut sink = Vec::new(); + assert!(keyed_outboard_post_order( + Cursor::new(Bytes::from(data.to_vec())), + oversize_tree, + &mut sink, + key + ) + .await + .is_err()); } /// Check that l and r of a 2-tuple are equal diff --git a/src/tests.rs b/src/tests.rs index 333b34d..c0942fd 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -19,28 +19,22 @@ use super::{ BaoTree, BlockSize, TreeNode, }; use crate::{ - assert_tuple_eq, blake3, hash_subtree, + assert_tuple_eq, blake3, io::{ - full_chunk_groups, - outboard::{PostOrderOutboard, PreOrderMemOutboard, PreOrderOutboard}, - sync::Outboard, - BaoContentItem, DecodeError, EncodeError, Leaf, + full_chunk_groups, outboard::PreOrderMemOutboard, sync::Outboard, BaoContentItem, + DecodeError, EncodeError, Leaf, }, iter::{PostOrderChunkIter, PreOrderPartialIterRef, ResponseIterRef}, - keyed_hash_subtree, keyed_parent_cv, parent_cv, prop_assert_tuple_eq, + keyed_hash_subtree, keyed_parent_cv, prop_assert_tuple_eq, rec::{ - encode_ranges_reference, encode_selected_rec, keyed_create_sized_keyed_checks, - keyed_init_from_keyed_checks, keyed_outboard_functions_checks, make_test_data, range_union, - truncate_ranges, ReferencePreOrderPartialChunkIterRef, + encode_ranges_reference, encode_selected_rec, keyed_outboard_functions_checks, + make_test_data, range_union, truncate_ranges, ReferencePreOrderPartialChunkIterRef, }, split, ChunkRanges, ChunkRangesRef, HashMode, ResponseIter, }; #[cfg(feature = "tokio_fsm")] -use crate::rec::{ - keyed_create_sized_keyed_checks_fsm, keyed_init_from_keyed_checks_fsm, - keyed_outboard_functions_checks_fsm, -}; +use crate::rec::keyed_outboard_functions_checks_fsm; /// Reference encoder using BLAKE3 keyed mode. fn keyed_encode_selected_reference( @@ -1260,11 +1254,11 @@ fn keyed_hash_subtree_differs_from_standard() { let data = make_test_data(2048); let key = blake3::derive_key("bao-tree.test", b"low-level-subtree"); - let standard = hash_subtree(0, &data, true); + let standard = HashMode::Standard.hash_subtree(0, &data, true); let keyed = keyed_hash_subtree(0, &data, true, &key); assert_ne!(standard, keyed); assert_eq!(keyed, blake3::keyed_hash(&key, &data)); - let non_root_standard = hash_subtree(1, &data[..1024], false); + let non_root_standard = HashMode::Standard.hash_subtree(1, &data[..1024], false); let non_root_keyed = keyed_hash_subtree(1, &data[..1024], false, &key); assert_ne!(non_root_standard, non_root_keyed); let mut hasher = blake3::Hasher::new_keyed(&key); @@ -1281,10 +1275,10 @@ fn keyed_parent_cv_differs_from_standard() { let left = blake3::hash(b"left"); let right = blake3::hash(b"right"); let key = blake3::derive_key("bao-tree.test", b"low-level-parent"); - let standard = parent_cv(&left, &right, true); + let standard = HashMode::Standard.parent_cv(&left, &right, true); let keyed = keyed_parent_cv(&left, &right, true, &key); assert_ne!(standard, keyed); - let standard_non_root = parent_cv(&left, &right, false); + let standard_non_root = HashMode::Standard.parent_cv(&left, &right, false); let keyed_non_root = keyed_parent_cv(&left, &right, false, &key); assert_ne!(standard_non_root, keyed_non_root); let left_cv: ChainingValue = *left.as_bytes(); @@ -1307,98 +1301,6 @@ fn keyed_pre_order_outboard_root_matches_blake3() { } } -#[test] -fn keyed_create_outboard_trait_sync() { - use crate::io::sync::CreateOutboard; - - let data = make_test_data(5000); - let key = blake3::derive_key("bao-tree.test", b"create-outboard"); - let block_size = BlockSize(2); - let post: PostOrderOutboard> = - PostOrderOutboard::create_keyed(Cursor::new(&data), block_size, &key).unwrap(); - assert_eq!(post.root(), blake3::keyed_hash(&key, &data)); - let pre: PreOrderOutboard> = - PreOrderOutboard::create_keyed(Cursor::new(&data), block_size, &key).unwrap(); - assert_eq!(pre.root(), blake3::keyed_hash(&key, &data)); -} - -#[cfg(feature = "tokio_fsm")] -#[tokio::test] -async fn keyed_create_outboard_trait_fsm() { - use crate::io::fsm::CreateOutboard; - - let data = make_test_data(5000); - let key = blake3::derive_key("bao-tree.test", b"create-outboard-fsm"); - let block_size = BlockSize(2); - let post: PostOrderOutboard> = - PostOrderOutboard::create_keyed(Bytes::from(data.clone()), block_size, &key) - .await - .unwrap(); - assert_eq!(post.root(), blake3::keyed_hash(&key, &data)); - let pre: PreOrderOutboard> = - PreOrderOutboard::create_keyed(Bytes::from(data.clone()), block_size, &key) - .await - .unwrap(); - assert_eq!(pre.root(), blake3::keyed_hash(&key, &data)); -} - -#[test] -fn keyed_create_sized_keyed_sync() { - let data = make_test_data(5000); - let key = blake3::derive_key("bao-tree.test", b"create-sized-keyed"); - keyed_create_sized_keyed_checks(&data, BlockSize(2), &key); -} - -#[test] -fn keyed_create_sized_keyed_empty_sync() { - let data: Vec = vec![]; - let key = blake3::derive_key("bao-tree.test", b"create-sized-keyed-empty"); - keyed_create_sized_keyed_checks(&data, BlockSize(0), &key); -} - -#[test] -fn keyed_create_sized_keyed_oversize_sync() { - use crate::io::sync::CreateOutboard; - - let data = make_test_data(100); - let key = blake3::derive_key("bao-tree.test", b"create-sized-keyed-oversize"); - let oversize = data.len() as u64 + 100; - assert!(PostOrderOutboard::>::create_sized_keyed( - Cursor::new(&data), - oversize, - BlockSize(0), - &key - ) - .is_err()); - let tree = BaoTree::new(oversize, BlockSize(0)); - let mut post = PostOrderOutboard { - root: blake3::Hash::from([0; 32]), - tree, - data: vec![0; tree.outboard_size().try_into().unwrap()], - }; - assert!(post.init_from_keyed(Cursor::new(&data), &key).is_err()); - assert!(PreOrderOutboard::>::create_sized_keyed( - Cursor::new(&data), - oversize, - BlockSize(0), - &key - ) - .is_err()); - let mut pre = PreOrderOutboard { - root: blake3::Hash::from([0; 32]), - tree, - data: vec![0; tree.outboard_size().try_into().unwrap()], - }; - assert!(pre.init_from_keyed(Cursor::new(&data), &key).is_err()); -} - -#[test] -fn keyed_init_from_keyed_sync() { - let data = make_test_data(5000); - let key = blake3::derive_key("bao-tree.test", b"init-from-keyed"); - keyed_init_from_keyed_checks(&data, BlockSize(2), &key); -} - #[test] fn keyed_outboard_functions_sync() { let data = make_test_data(5000); @@ -1406,75 +1308,6 @@ fn keyed_outboard_functions_sync() { keyed_outboard_functions_checks(&data, BlockSize(2), &key); } -#[cfg(feature = "tokio_fsm")] -#[tokio::test] -async fn keyed_create_sized_keyed_fsm() { - let data = make_test_data(5000); - let key = blake3::derive_key("bao-tree.test", b"create-sized-keyed-fsm"); - keyed_create_sized_keyed_checks_fsm(&data, BlockSize(2), &key).await; -} - -#[cfg(feature = "tokio_fsm")] -#[tokio::test] -async fn keyed_create_sized_keyed_empty_fsm() { - let data: Vec = vec![]; - let key = blake3::derive_key("bao-tree.test", b"create-sized-keyed-empty-fsm"); - keyed_create_sized_keyed_checks_fsm(&data, BlockSize(0), &key).await; -} - -#[cfg(feature = "tokio_fsm")] -#[tokio::test] -async fn keyed_create_sized_keyed_oversize_fsm() { - use crate::io::fsm::CreateOutboard; - - let data = make_test_data(100); - let key = blake3::derive_key("bao-tree.test", b"create-sized-keyed-oversize-fsm"); - let oversize = data.len() as u64 + 100; - assert!(PostOrderOutboard::>::create_sized_keyed( - Cursor::new(Bytes::from(data.clone())), - oversize, - BlockSize(0), - &key - ) - .await - .is_err()); - let tree = BaoTree::new(oversize, BlockSize(0)); - let mut post = PostOrderOutboard { - root: blake3::Hash::from([0; 32]), - tree, - data: vec![0; tree.outboard_size().try_into().unwrap()], - }; - assert!(post - .init_from_keyed(Cursor::new(Bytes::from(data.clone())), &key) - .await - .is_err()); - assert!(PreOrderOutboard::>::create_sized_keyed( - Cursor::new(Bytes::from(data.clone())), - oversize, - BlockSize(0), - &key - ) - .await - .is_err()); - let mut pre = PreOrderOutboard { - root: blake3::Hash::from([0; 32]), - tree, - data: vec![0; tree.outboard_size().try_into().unwrap()], - }; - assert!(pre - .init_from_keyed(Cursor::new(Bytes::from(data)), &key) - .await - .is_err()); -} - -#[cfg(feature = "tokio_fsm")] -#[tokio::test] -async fn keyed_init_from_keyed_fsm() { - let data = make_test_data(5000); - let key = blake3::derive_key("bao-tree.test", b"init-from-keyed-fsm"); - keyed_init_from_keyed_checks_fsm(&data, BlockSize(2), &key).await; -} - #[cfg(feature = "tokio_fsm")] #[tokio::test] async fn keyed_outboard_functions_fsm() { diff --git a/src/tests2.rs b/src/tests2.rs index 3309847..fd2d7d9 100644 --- a/src/tests2.rs +++ b/src/tests2.rs @@ -17,7 +17,7 @@ use smallvec::SmallVec; use test_strategy::proptest; use crate::{ - assert_tuple_eq, blake3, hash_subtree, + assert_tuple_eq, blake3, io::{ fsm::ResponseDecoderNext, outboard::{PostOrderMemOutboard, PreOrderMemOutboard}, @@ -25,10 +25,9 @@ use crate::{ BaoContentItem, Leaf, Parent, }, iter::{BaoChunk, PreOrderPartialChunkIterRef, ResponseIterRef}, - keyed_hash_subtree, keyed_parent_cv, parent_cv, prop_assert_tuple_eq, + keyed_hash_subtree, keyed_parent_cv, prop_assert_tuple_eq, rec::{ - encode_selected_rec, get_leaf_ranges, keyed_create_sized_keyed_checks, - keyed_init_from_keyed_checks, keyed_outboard_functions_checks, make_test_data, + encode_selected_rec, get_leaf_ranges, keyed_outboard_functions_checks, make_test_data, partial_chunk_iter_reference, range_union, response_iter_reference, select_nodes_rec, truncate_ranges, ReferencePreOrderPartialChunkIterRef, }, @@ -180,8 +179,8 @@ fn outboard_test_sync(data: &[u8], outboard: impl crate::io::sync::Outboard) { let start_chunk = node.chunk_range().start; let byte_range = tree.byte_range(node); let data = &data[byte_range.start.try_into().unwrap()..byte_range.end.try_into().unwrap()]; - let expected = hash_subtree(start_chunk.0, data, is_root); - let actual = parent_cv(&l_hash, &r_hash, is_root); + let expected = HashMode::Standard.hash_subtree(start_chunk.0, data, is_root); + let actual = HashMode::Standard.parent_cv(&l_hash, &r_hash, is_root); assert_eq!(actual, expected); } } @@ -224,8 +223,8 @@ async fn outboard_test_fsm(data: &[u8], mut outboard: impl crate::io::fsm::Outbo let start_chunk = node.chunk_range().start; let byte_range = tree.byte_range(node); let data = &data[byte_range.start.try_into().unwrap()..byte_range.end.try_into().unwrap()]; - let expected = hash_subtree(start_chunk.0, data, is_root); - let actual = parent_cv(&l_hash, &r_hash, is_root); + let expected = HashMode::Standard.hash_subtree(start_chunk.0, data, is_root); + let actual = HashMode::Standard.parent_cv(&l_hash, &r_hash, is_root); assert_eq!(actual, expected); } } @@ -339,20 +338,6 @@ fn keyed_pre_order_outboard_fsm_proptest(#[strategy(tree())] tree: BaoTree) { keyed_pre_order_outboard_fsm_impl(tree); } -#[proptest] -fn keyed_create_sized_keyed_proptest(#[strategy(tree())] tree: BaoTree) { - let data = make_test_data(tree.size.try_into().unwrap()); - let key = keyed_test_key(&tree.size.to_le_bytes()); - keyed_create_sized_keyed_checks(&data, tree.block_size, &key); -} - -#[proptest] -fn keyed_init_from_keyed_proptest(#[strategy(tree())] tree: BaoTree) { - let data = make_test_data(tree.size.try_into().unwrap()); - let key = keyed_test_key(&tree.size.to_le_bytes()); - keyed_init_from_keyed_checks(&data, tree.block_size, &key); -} - #[proptest] fn keyed_outboard_functions_proptest(#[strategy(tree())] tree: BaoTree) { let data = make_test_data(tree.size.try_into().unwrap()); @@ -360,30 +345,6 @@ fn keyed_outboard_functions_proptest(#[strategy(tree())] tree: BaoTree) { keyed_outboard_functions_checks(&data, tree.block_size, &key); } -#[cfg(feature = "tokio_fsm")] -#[proptest] -fn keyed_create_sized_keyed_fsm_proptest(#[strategy(tree())] tree: BaoTree) { - let data = make_test_data(tree.size.try_into().unwrap()); - let key = keyed_test_key(&tree.size.to_le_bytes()); - run_blocking(crate::rec::keyed_create_sized_keyed_checks_fsm( - &data, - tree.block_size, - &key, - )); -} - -#[cfg(feature = "tokio_fsm")] -#[proptest] -fn keyed_init_from_keyed_fsm_proptest(#[strategy(tree())] tree: BaoTree) { - let data = make_test_data(tree.size.try_into().unwrap()); - let key = keyed_test_key(&tree.size.to_le_bytes()); - run_blocking(crate::rec::keyed_init_from_keyed_checks_fsm( - &data, - tree.block_size, - &key, - )); -} - fn mem_outboard_flip_impl(tree: BaoTree) { let data = make_test_data(tree.size.try_into().unwrap()); let post = PostOrderMemOutboard::create(&data, tree.block_size); From b13747b17ad5337bc085bfef768949d752eee5a6 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Aug 2026 11:29:46 +0200 Subject: [PATCH 09/12] Remove unneeded lifetime --- src/io/fsm.rs | 2 +- src/io/sync.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/io/fsm.rs b/src/io/fsm.rs index 3947656..620342d 100644 --- a/src/io/fsm.rs +++ b/src/io/fsm.rs @@ -934,7 +934,7 @@ mod validate { outboard: O, data: D, ranges: &'a ChunkRangesRef, - key: &'a [u8; 32], + key: &[u8; 32], ) -> impl Stream>> + 'a where O: Outboard + 'a, diff --git a/src/io/sync.rs b/src/io/sync.rs index 59c01d9..f5d77dd 100644 --- a/src/io/sync.rs +++ b/src/io/sync.rs @@ -819,7 +819,7 @@ mod validate { outboard: O, data: D, ranges: &'a ChunkRangesRef, - key: &'a [u8; 32], + key: &[u8; 32], ) -> impl IntoIterator>> + 'a where O: Outboard + 'a, From 6648a939743722fa44604072530e4aa774709a8f Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Aug 2026 11:34:23 +0200 Subject: [PATCH 10/12] DRY tests --- src/tests.rs | 155 ++++++++++++-------------------------------------- src/tests2.rs | 72 ++++++----------------- 2 files changed, 53 insertions(+), 174 deletions(-) diff --git a/src/tests.rs b/src/tests.rs index c0942fd..5c3482e 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -89,19 +89,18 @@ fn keyed_encode_decode_roundtrip_sync_impl(data: &[u8], block_size: BlockSize, k fn keyed_encode_decode_roundtrip_fsm_impl(data: Vec, block_size: BlockSize, key: &[u8; 32]) { use crate::io::fsm::{keyed_decode_ranges, keyed_encode_ranges_validated}; + let rt = tokio::runtime::Runtime::new().unwrap(); let mut outboard = PostOrderMemOutboard::create_keyed(&data, block_size, key); let ranges = ChunkRanges::all(); let mut encoded = Vec::new(); - tokio::runtime::Runtime::new() - .unwrap() - .block_on(keyed_encode_ranges_validated( - Bytes::from(data.clone()), - &mut outboard, - &ranges, - &mut encoded, - key, - )) - .unwrap(); + rt.block_on(keyed_encode_ranges_validated( + Bytes::from(data.clone()), + &mut outboard, + &ranges, + &mut encoded, + key, + )) + .unwrap(); let tree = outboard.tree(); let mut decoded = bytes::BytesMut::new(); let mut ob_res = PostOrderMemOutboard { @@ -109,16 +108,14 @@ fn keyed_encode_decode_roundtrip_fsm_impl(data: Vec, block_size: BlockSize, tree, data: vec![0; tree.outboard_size().try_into().unwrap()], }; - tokio::runtime::Runtime::new() - .unwrap() - .block_on(keyed_decode_ranges( - Cursor::new(encoded.as_slice()), - ranges, - &mut decoded, - &mut ob_res, - key, - )) - .unwrap(); + rt.block_on(keyed_decode_ranges( + Cursor::new(encoded.as_slice()), + ranges, + &mut decoded, + &mut ob_res, + key, + )) + .unwrap(); assert_eq!(decoded.to_vec(), data); assert_eq!(ob_res.root(), outboard.root()); } @@ -1177,75 +1174,14 @@ fn keyed_domain_separation() { #[test] fn keyed_encode_decode_roundtrip_sync() { - use crate::io::sync::{keyed_decode_ranges, keyed_encode_ranges_validated}; - - let data = make_test_data(50_000); let key = blake3::derive_key("bao-tree.test", b"roundtrip"); - let block_size = BlockSize(2); - let outboard = PostOrderMemOutboard::create_keyed(&data, block_size, &key); - let ranges = ChunkRanges::all(); - let mut encoded = Vec::new(); - keyed_encode_ranges_validated(&data, &outboard, &ranges, &mut encoded, &key).unwrap(); - let size = outboard.tree.size; - let tree = BaoTree::new(size, block_size); - let mut decoded = Vec::new(); - let mut ob_res = PostOrderMemOutboard { - root: outboard.root(), - tree, - data: vec![0; tree.outboard_size().try_into().unwrap()], - }; - keyed_decode_ranges( - Cursor::new(encoded), - &ranges, - &mut decoded, - &mut ob_res, - &key, - ) - .unwrap(); - assert_eq!(decoded, data); - assert_eq!(ob_res.root(), outboard.root()); + keyed_encode_decode_roundtrip_sync_impl(&make_test_data(50_000), BlockSize(2), &key); } #[test] fn keyed_encode_decode_roundtrip_fsm() { - use crate::io::fsm::{keyed_decode_ranges, keyed_encode_ranges_validated}; - - let data = make_test_data(50_000); let key = blake3::derive_key("bao-tree.test", b"roundtrip"); - let block_size = BlockSize(2); - let mut outboard = PostOrderMemOutboard::create_keyed(&data, block_size, &key); - let ranges = ChunkRanges::all(); - let mut encoded = Vec::new(); - tokio::runtime::Runtime::new() - .unwrap() - .block_on(keyed_encode_ranges_validated( - Bytes::from(data.clone()), - &mut outboard, - &ranges, - &mut encoded, - &key, - )) - .unwrap(); - let size = outboard.tree.size; - let tree = BaoTree::new(size, block_size); - let mut decoded = bytes::BytesMut::new(); - let mut ob_res = PostOrderMemOutboard { - root: outboard.root(), - tree, - data: vec![0; tree.outboard_size().try_into().unwrap()], - }; - tokio::runtime::Runtime::new() - .unwrap() - .block_on(keyed_decode_ranges( - Cursor::new(encoded.as_slice()), - ranges, - &mut decoded, - &mut ob_res, - &key, - )) - .unwrap(); - assert_eq!(decoded.to_vec(), data); - assert_eq!(ob_res.root(), outboard.root()); + keyed_encode_decode_roundtrip_fsm_impl(make_test_data(50_000), BlockSize(2), &key); } #[test] @@ -1576,59 +1512,42 @@ fn keyed_encode_decode_edge_sizes_fsm() { } } -fn keyed_bao_tree_slice_roundtrip_case_table(key: &[u8; 32]) { - use make_test_data as td; +const KEYED_SLICE_ROUNDTRIP_CASES: [(usize, std::ops::Range); 8] = [ + (0, 0..1), + (1, 0..1), + (1023, 0..1), + (1024, 0..1), + (1025, 0..1), + (1025, 0..2), + (1025, 1..2), + (24 * 1024 + 1, 0..25), +]; - let cases = [ - (0, 0..1), - (1, 0..1), - (1023, 0..1), - (1024, 0..1), - (1025, 0..1), - (1025, 0..2), - (1025, 1..2), - (24 * 1024 + 1, 0..25), - ]; +#[test] +fn keyed_bao_tree_slice_roundtrip_cases() { + let key = blake3::derive_key("bao-tree.test", b"slice"); for chunk_group_log in 0..4 { let block_size = BlockSize(chunk_group_log); - for (count, range) in cases.clone() { + for (count, range) in KEYED_SLICE_ROUNDTRIP_CASES { keyed_bao_tree_slice_roundtrip_test( - td(count), + make_test_data(count), ChunkNum(range.start)..ChunkNum(range.end), block_size, - key, + &key, ); } } } -#[test] -fn keyed_bao_tree_slice_roundtrip_cases() { - let key = blake3::derive_key("bao-tree.test", b"slice"); - keyed_bao_tree_slice_roundtrip_case_table(&key); -} - #[cfg(feature = "tokio_fsm")] #[tokio::test] async fn keyed_bao_tree_slice_roundtrip_fsm_cases() { - use make_test_data as td; - let key = blake3::derive_key("bao-tree.test", b"slice-fsm"); - let cases = [ - (0, 0..1), - (1, 0..1), - (1023, 0..1), - (1024, 0..1), - (1025, 0..1), - (1025, 0..2), - (1025, 1..2), - (24 * 1024 + 1, 0..25), - ]; for chunk_group_log in 0..4 { let block_size = BlockSize(chunk_group_log); - for (count, range) in cases.clone() { + for (count, range) in KEYED_SLICE_ROUNDTRIP_CASES { keyed_bao_tree_slice_roundtrip_fsm_test( - td(count), + make_test_data(count), ChunkNum(range.start)..ChunkNum(range.end), block_size, &key, diff --git a/src/tests2.rs b/src/tests2.rs index fd2d7d9..394d04e 100644 --- a/src/tests2.rs +++ b/src/tests2.rs @@ -25,7 +25,7 @@ use crate::{ BaoContentItem, Leaf, Parent, }, iter::{BaoChunk, PreOrderPartialChunkIterRef, ResponseIterRef}, - keyed_hash_subtree, keyed_parent_cv, prop_assert_tuple_eq, + prop_assert_tuple_eq, rec::{ encode_selected_rec, get_leaf_ranges, keyed_outboard_functions_checks, make_test_data, partial_chunk_iter_reference, range_union, response_iter_reference, select_nodes_rec, @@ -145,8 +145,8 @@ fn post_traversal_chunks_iter_proptest(#[strategy(tree())] tree: BaoTree) { post_traversal_chunks_iter_impl(tree); } -/// Brute force test for a keyed outboard that computes expected hashes for each pair -fn keyed_outboard_test_sync(data: &[u8], outboard: impl crate::io::sync::Outboard, key: &[u8; 32]) { +/// Brute force test for an outboard that computes the expected hash for each pair +fn outboard_test_sync(data: &[u8], outboard: impl crate::io::sync::Outboard, mode: HashMode) { let tree = outboard.tree(); let nodes = tree .pre_order_nodes_iter() @@ -159,37 +159,17 @@ fn keyed_outboard_test_sync(data: &[u8], outboard: impl crate::io::sync::Outboar let start_chunk = node.chunk_range().start; let byte_range = tree.byte_range(node); let data = &data[byte_range.start.try_into().unwrap()..byte_range.end.try_into().unwrap()]; - let expected = keyed_hash_subtree(start_chunk.0, data, is_root, key); - let actual = keyed_parent_cv(&l_hash, &r_hash, is_root, key); + let expected = mode.hash_subtree(start_chunk.0, data, is_root); + let actual = mode.parent_cv(&l_hash, &r_hash, is_root); assert_eq!(actual, expected); } } -/// Brute force test for an outboard that just computes the expected hash for each pair -fn outboard_test_sync(data: &[u8], outboard: impl crate::io::sync::Outboard) { - let tree = outboard.tree(); - let nodes = tree - .pre_order_nodes_iter() - .enumerate() - .map(|(i, node)| (node, i == 0)) - .filter(|(node, _)| tree.is_relevant_for_outboard(*node)) - .collect::>(); - for (node, is_root) in nodes { - let (l_hash, r_hash) = outboard.load(node).unwrap().unwrap(); - let start_chunk = node.chunk_range().start; - let byte_range = tree.byte_range(node); - let data = &data[byte_range.start.try_into().unwrap()..byte_range.end.try_into().unwrap()]; - let expected = HashMode::Standard.hash_subtree(start_chunk.0, data, is_root); - let actual = HashMode::Standard.parent_cv(&l_hash, &r_hash, is_root); - assert_eq!(actual, expected); - } -} - -/// Brute force test for a keyed outboard that computes expected hashes for each pair -async fn keyed_outboard_test_fsm( +/// Brute force test for an outboard that computes the expected hash for each pair +async fn outboard_test_fsm( data: &[u8], mut outboard: impl crate::io::fsm::Outboard, - key: &[u8; 32], + mode: HashMode, ) { let tree = outboard.tree(); let nodes = tree @@ -203,28 +183,8 @@ async fn keyed_outboard_test_fsm( let start_chunk = node.chunk_range().start; let byte_range = tree.byte_range(node); let data = &data[byte_range.start.try_into().unwrap()..byte_range.end.try_into().unwrap()]; - let expected = keyed_hash_subtree(start_chunk.0, data, is_root, key); - let actual = keyed_parent_cv(&l_hash, &r_hash, is_root, key); - assert_eq!(actual, expected); - } -} - -/// Brute force test for an outboard that just computes the expected hash for each pair -async fn outboard_test_fsm(data: &[u8], mut outboard: impl crate::io::fsm::Outboard) { - let tree = outboard.tree(); - let nodes = tree - .pre_order_nodes_iter() - .enumerate() - .map(|(i, node)| (node, i == 0)) - .filter(|(node, _)| tree.is_relevant_for_outboard(*node)) - .collect::>(); - for (node, is_root) in nodes { - let (l_hash, r_hash) = outboard.load(node).await.unwrap().unwrap(); - let start_chunk = node.chunk_range().start; - let byte_range = tree.byte_range(node); - let data = &data[byte_range.start.try_into().unwrap()..byte_range.end.try_into().unwrap()]; - let expected = HashMode::Standard.hash_subtree(start_chunk.0, data, is_root); - let actual = HashMode::Standard.parent_cv(&l_hash, &r_hash, is_root); + let expected = mode.hash_subtree(start_chunk.0, data, is_root); + let actual = mode.parent_cv(&l_hash, &r_hash, is_root); assert_eq!(actual, expected); } } @@ -236,7 +196,7 @@ fn post_oder_outboard_sync_impl(tree: BaoTree) { outboard.data.len() as u64, outboard.tree().outboard_hash_pairs() * 64 ); - outboard_test_sync(&data, outboard); + outboard_test_sync(&data, outboard, HashMode::Standard); } #[test] @@ -262,7 +222,7 @@ fn post_oder_outboard_fsm_impl(tree: BaoTree) { ); tokio::runtime::Runtime::new() .unwrap() - .block_on(outboard_test_fsm(&data, outboard)); + .block_on(outboard_test_fsm(&data, outboard, HashMode::Standard)); } #[proptest] @@ -278,7 +238,7 @@ fn keyed_post_order_outboard_sync_impl(tree: BaoTree) { outboard.data.len() as u64, outboard.tree().outboard_hash_pairs() * 64 ); - keyed_outboard_test_sync(&data, outboard, &key); + outboard_test_sync(&data, outboard, HashMode::Keyed(key)); } #[proptest] @@ -296,7 +256,7 @@ fn keyed_post_order_outboard_fsm_impl(tree: BaoTree) { ); tokio::runtime::Runtime::new() .unwrap() - .block_on(keyed_outboard_test_fsm(&data, outboard, &key)); + .block_on(outboard_test_fsm(&data, outboard, HashMode::Keyed(key))); } #[proptest] @@ -312,7 +272,7 @@ fn keyed_pre_order_outboard_sync_impl(tree: BaoTree) { outboard.data.len(), outboard.tree().outboard_size().try_into().unwrap() ); - keyed_outboard_test_sync(&data, outboard, &key); + outboard_test_sync(&data, outboard, HashMode::Keyed(key)); } #[proptest] @@ -330,7 +290,7 @@ fn keyed_pre_order_outboard_fsm_impl(tree: BaoTree) { ); tokio::runtime::Runtime::new() .unwrap() - .block_on(keyed_outboard_test_fsm(&data, outboard, &key)); + .block_on(outboard_test_fsm(&data, outboard, HashMode::Keyed(key))); } #[proptest] From f1b9bed599f0abd1e0ad068247c5b6e9cbfab7c5 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Aug 2026 11:41:19 +0200 Subject: [PATCH 11/12] fixes --- src/io/fsm.rs | 8 -------- src/io/sync.rs | 23 +++++++++++++++-------- src/lib.rs | 2 ++ src/rec.rs | 2 +- 4 files changed, 18 insertions(+), 17 deletions(-) diff --git a/src/io/fsm.rs b/src/io/fsm.rs index 620342d..5d5f7cd 100644 --- a/src/io/fsm.rs +++ b/src/io/fsm.rs @@ -391,9 +391,6 @@ impl ResponseDecoder { HashMode::Keyed(*key), ))) } -} - -impl ResponseDecoder { pub(crate) fn with_mode( hash: blake3::Hash, ranges: ChunkRanges, @@ -574,7 +571,6 @@ where encode_ranges_validated_impl(data, outboard, ranges, encoded, HashMode::Keyed(*key)).await } -/// Generic encode body monomorphized over the compile time hashing strategy. async fn encode_ranges_validated_impl( mut data: D, mut outboard: O, @@ -700,7 +696,6 @@ where decode_ranges_impl(encoded, ranges, target, outboard, HashMode::Keyed(*key)).await } -/// Generic decode body monomorphized over the compile time hashing strategy. async fn decode_ranges_impl( encoded: R, ranges: ChunkRanges, @@ -774,7 +769,6 @@ async fn outboard_with_mode( outboard_impl(tree, data, &mut outboard, &mut buffer, mode).await } -/// Generic outboard traversal monomorphized over the compile time hashing strategy. async fn outboard_impl( tree: BaoTree, mut data: impl AsyncStreamReader, @@ -846,7 +840,6 @@ async fn outboard_post_order_with_mode( outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, mode).await } -/// Generic post order outboard traversal monomorphized over the compile time hashing strategy. async fn outboard_post_order_impl( tree: BaoTree, mut data: impl AsyncStreamReader, @@ -943,7 +936,6 @@ mod validate { valid_ranges_impl(outboard, data, ranges, HashMode::Keyed(*key)) } - /// Generic validation body monomorphized over the compile time hashing strategy. fn valid_ranges_impl<'a, O, D>( outboard: O, data: D, diff --git a/src/io/sync.rs b/src/io/sync.rs index f5d77dd..cdf378c 100644 --- a/src/io/sync.rs +++ b/src/io/sync.rs @@ -301,11 +301,23 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { key: &[u8; 32], ) -> Self { let buf = BytesMut::with_capacity(tree.block_size().bytes()); - DecodeResponseIter::with_mode(root, tree, encoded, ranges, buf, HashMode::Keyed(*key)) + Self::new_keyed_with_buffer(root, tree, encoded, ranges, buf, key) } -} -impl<'a, R: Read> DecodeResponseIter<'a, R> { + /// Create a new iterator to decode a keyed response. + /// + /// This is the same as [Self::new_keyed], but allows you to provide a buffer to use for decoding. + /// The buffer will be resized as needed, but it's capacity should be the [crate::BlockSize::bytes]. + pub fn new_keyed_with_buffer( + root: blake3::Hash, + tree: BaoTree, + encoded: R, + ranges: &'a ChunkRangesRef, + buf: BytesMut, + key: &[u8; 32], + ) -> Self { + DecodeResponseIter::with_mode(root, tree, encoded, ranges, buf, HashMode::Keyed(*key)) + } pub(crate) fn with_mode( root: blake3::Hash, tree: BaoTree, @@ -462,7 +474,6 @@ pub fn keyed_encode_ranges_validated( encode_ranges_validated_impl(data, outboard, ranges, encoded, HashMode::Keyed(*key)) } -/// Generic encode body monomorphized over the compile time hashing strategy. fn encode_ranges_validated_impl( data: D, outboard: O, @@ -583,7 +594,6 @@ where decode_ranges_impl(encoded, ranges, target, outboard, HashMode::Keyed(*key)) } -/// Generic decode body monomorphized over the compile time hashing strategy. fn decode_ranges_impl( encoded: R, ranges: &ChunkRangesRef, @@ -650,7 +660,6 @@ fn outboard_with_mode( outboard_impl(tree, data, &mut outboard, &mut buffer, mode) } -/// Generic outboard traversal monomorphized over the compile time hashing strategy. fn outboard_impl( tree: BaoTree, mut data: impl Read, @@ -723,7 +732,6 @@ fn outboard_post_order_with_mode( outboard_post_order_impl(tree, data, &mut outboard, &mut buffer, mode) } -/// Generic post order outboard traversal monomorphized over the compile time hashing strategy. fn outboard_post_order_impl( tree: BaoTree, mut data: impl Read, @@ -828,7 +836,6 @@ mod validate { valid_ranges_impl(outboard, data, ranges, HashMode::Keyed(*key)) } - /// Generic validation body monomorphized over the compile time hashing strategy. fn valid_ranges_impl<'a, O, D>( outboard: O, data: D, diff --git a/src/lib.rs b/src/lib.rs index f1299f6..5ca927b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -305,6 +305,7 @@ impl HashMode { /// Compute the hash of a subtree using BLAKE3 keyed mode. /// /// See [keyed_parent_cv] for merging child hashes in keyed mode. +#[inline] pub fn keyed_hash_subtree( start_chunk: u64, data: &[u8], @@ -315,6 +316,7 @@ pub fn keyed_hash_subtree( } /// Merge two child subtree hashes using BLAKE3 keyed mode. +#[inline] pub fn keyed_parent_cv( left_child: &blake3::Hash, right_child: &blake3::Hash, diff --git a/src/rec.rs b/src/rec.rs index 80d7e16..a279680 100644 --- a/src/rec.rs +++ b/src/rec.rs @@ -446,7 +446,7 @@ mod test_support { use std::io::Cursor; use crate::io::outboard::{PostOrderMemOutboard, PreOrderMemOutboard, PreOrderOutboard}; - use crate::io::sync::{self}; + use crate::io::sync; pub(crate) fn assert_pre_order_outboard_matches_mem( outboard: &PreOrderOutboard>, From e82e7442e4053277f792ad9630d66a9e60fa75e5 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Aug 2026 11:44:16 +0200 Subject: [PATCH 12/12] Add TreeNode::Hash --- src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lib.rs b/src/lib.rs index 5ca927b..4848ada 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -612,7 +612,7 @@ pub(crate) const fn blocks(size: u64, block_size: BlockSize) -> u64 { /// You typically don't have to use this, but it can be useful for debugging /// and error handling. Hash validation errors contain a `TreeNode` that allows /// you to find the position where validation failed. -#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub struct TreeNode(u64);