diff --git a/src/io/fsm.rs b/src/io/fsm.rs index 5c64e01..21fc1a0 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, + blake3, io::{ error::EncodeError, outboard::{PostOrderOutboard, PreOrderOutboard}, Leaf, Parent, }, iter::{BaoChunk, ResponseIter}, - parent_cv, 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. @@ -129,6 +128,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 +280,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 +335,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 { @@ -314,37 +402,52 @@ 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, } -impl ResponseDecoderInner { - fn new(tree: BaoTree, hash: blake3::Hash, ranges: ChunkRanges, encoded: R) -> Self { +impl ResponseDecoderInner { + fn with_hash_strategy( + tree: BaoTree, + hash: blake3::Hash, + ranges: ChunkRanges, + encoded: R, + hash_strategy: H, + ) -> 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, + hash_strategy, }; res.stack.push(hash); res } } -/// Response decoder +/// Response decoder. +/// +/// Keyed callers should use [KeyedResponseDecoder] via [Self::new_keyed]. #[derive(Debug)] -pub struct ResponseDecoder(Box>); +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 { +pub enum ResponseDecoderNext { /// One more item, and you get back the state machine in the next state More( ( - ResponseDecoder, + ResponseDecoder, std::result::Result, ), ), @@ -352,18 +455,58 @@ 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::new( - tree, hash, ranges, encoded, + Self(Box::new(ResponseDecoderInner::with_hash_strategy( + tree, hash, ranges, encoded, Standard, + ))) + } + + /// Create a new keyed response decoder. + pub fn new_keyed( + hash: blake3::Hash, + ranges: ChunkRanges, + tree: BaoTree, + encoded: R, + key: &[u8; 32], + ) -> KeyedResponseDecoder { + ResponseDecoder(Box::new(ResponseDecoderInner::with_hash_strategy( + tree, + hash, + ranges, + encoded, + Keyed(*key), + ))) + } +} + +impl ResponseDecoder { + pub(crate) fn with_hash_strategy( + hash: blake3::Hash, + ranges: ChunkRanges, + tree: BaoTree, + encoded: R, + hash_strategy: H, + ) -> Self { + Self(Box::new(ResponseDecoderInner::with_hash_strategy( + tree, + hash, + ranges, + encoded, + 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)) @@ -404,7 +547,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 = 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 { @@ -434,7 +577,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(start_chunk.0, &data, is_root); + let actual = this + .hash_strategy + .hash_subtree(start_chunk.0, &data, is_root); if leaf_hash != actual { return Err(DecodeError::LeafHashMismatch(start_chunk)); } @@ -501,10 +646,42 @@ 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_impl(data, outboard, ranges, encoded, Standard).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_impl(data, outboard, ranges, encoded, 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, ranges: &ChunkRangesRef, encoded: W, + hash_strategy: H, ) -> result::Result<(), EncodeError> where D: AsyncSliceReader, @@ -529,7 +706,7 @@ where .. } => { let (l_hash, r_hash) = outboard.load(node).await?.unwrap(); - let actual = parent_cv(&l_hash, &r_hash, is_root); + 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)); @@ -570,10 +747,11 @@ where tree.block_size.to_u32(), true, &mut out_buf, + hash_strategy, ); (actual, out_buf.clone().into()) } else { - let actual = hash_subtree(start_chunk.0, &bytes, is_root); + let actual = hash_strategy.hash_subtree(start_chunk.0, &bytes, is_root); (actual, bytes) }; if actual != expected { @@ -594,17 +772,55 @@ 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_impl(encoded, ranges, target, outboard, Standard).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_impl(encoded, ranges, target, outboard, Keyed(*key)).await +} + +/// Generic decode body monomorphized over the compile time hashing strategy. +async fn decode_ranges_impl( encoded: R, ranges: ChunkRanges, mut target: W, mut outboard: O, + hash_strategy: H, ) -> 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 = ResponseDecoder::with_hash_strategy( + outboard.root(), + ranges, + outboard.tree(), + encoded, + hash_strategy, + ); loop { let item = match reading.next().await { ResponseDecoderNext::Done(_reader) => break, @@ -624,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()); @@ -635,21 +852,41 @@ 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_hash_strategy(data, tree, outboard, Standard).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_hash_strategy(data, tree, outboard, Keyed(*key)).await +} + +/// 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, + hash_strategy: H, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - let hash = outboard_impl(tree, data, &mut outboard, &mut buffer).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( +/// Generic outboard traversal monomorphized over the compile time hashing strategy. +async fn outboard_impl( tree: BaoTree, mut data: impl AsyncStreamReader, mut outboard: impl OutboardMut, buffer: &mut [u8], + hash_strategy: H, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -660,7 +897,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 = hash_strategy.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -670,7 +907,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_strategy.hash_subtree(start_chunk.0, &buf, is_root); stack.push(hash); } } @@ -687,21 +924,41 @@ 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_hash_strategy(data, tree, outboard, Standard).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_hash_strategy(data, tree, outboard, Keyed(*key)).await +} + +/// 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, + 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).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( +/// Generic post order outboard traversal monomorphized over the compile time hashing strategy. +async fn outboard_post_order_impl( tree: BaoTree, mut data: impl AsyncStreamReader, mut outboard: impl AsyncStreamWriter, buffer: &mut [u8], + hash_strategy: H, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -713,7 +970,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 = hash_strategy.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -723,7 +980,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_strategy.hash_subtree(start_chunk.0, &buf, is_root); stack.push(hash); } } @@ -757,8 +1014,8 @@ mod validate { use super::Outboard; use crate::{ - blake3, hash_subtree, io::LocalBoxFuture, parent_cv, 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. @@ -771,32 +1028,63 @@ mod validate { data: D, ranges: &'a ChunkRangesRef, ) -> impl Stream>> + 'a + where + O: Outboard + 'a, + D: AsyncSliceReader + 'a, + { + valid_ranges_impl(outboard, data, ranges, Standard) + } + + /// 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_impl(outboard, data, ranges, Keyed(*key)) + } + + /// 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, + ranges: &'a ChunkRangesRef, + hash_strategy: H, + ) -> 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, 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>>, + hash_strategy: H, } - impl RecursiveDataValidator<'_, O, D> { + impl RecursiveDataValidator<'_, O, D, H> { async fn validate( outboard: O, data: D, ranges: &ChunkRangesRef, co: &Co>>, + hash_strategy: H, ) -> io::Result<()> { let tree = outboard.tree(); if tree.blocks() == 1 { @@ -805,7 +1093,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_strategy.hash_subtree(0, &data, true); if actual == outboard.root() { co.yield_(Ok(ChunkNum(0)..tree.chunks())).await; } @@ -820,6 +1108,7 @@ mod validate { outboard, data, co, + hash_strategy, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -835,7 +1124,11 @@ 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 = self.hash_strategy.hash_subtree( + ChunkNum::full_chunks(range.start).0, + &data, + is_root, + ); if &actual == hash { // yield the left range self.co @@ -869,7 +1162,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = parent_cv(&l_hash, &r_hash, is_root); + 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(()); @@ -972,7 +1265,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = parent_cv(&l_hash, &r_hash, is_root); + let actual = Standard.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -998,4 +1291,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..2058b58 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, + blake3, io::{ error::EncodeError, outboard::{parse_hash_pair, PostOrderOutboard, PreOrderOutboard}, Leaf, Parent, }, iter::{BaoChunk, ResponseIterRef}, - parent_cv, 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. @@ -97,6 +96,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 +219,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 +270,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 { @@ -257,15 +337,23 @@ 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. +/// +/// Keyed callers should use [KeyedDecodeResponseIter] via [Self::new_keyed]. #[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, + 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. /// @@ -286,6 +374,31 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { encoded: R, ranges: &'a ChunkRangesRef, buf: BytesMut, + ) -> Self { + DecodeResponseIter::with_hash_strategy(root, tree, encoded, ranges, buf, Standard) + } + + /// 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], + ) -> 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> { + pub(crate) fn with_hash_strategy( + root: blake3::Hash, + tree: BaoTree, + encoded: R, + ranges: &'a ChunkRangesRef, + buf: BytesMut, + hash_strategy: H, ) -> Self { let ranges = truncate_ranges(ranges, tree.size()); let mut stack = SmallVec::new(); @@ -295,6 +408,7 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { inner: ResponseIterRef::new(tree, ranges), encoded, buf, + hash_strategy, } } @@ -322,7 +436,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 = self.hash_strategy.parent_cv(&l_hash, &r_hash, is_root); if parent_hash != actual { return Err(DecodeError::ParentHashMismatch(node)); } @@ -344,7 +458,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(start_chunk.0, &self.buf, is_root); + 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)); @@ -362,7 +478,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 { @@ -419,6 +535,28 @@ pub fn encode_ranges_validated( outboard: O, ranges: &ChunkRangesRef, encoded: W, +) -> result::Result<(), EncodeError> { + encode_ranges_validated_impl(data, outboard, ranges, encoded, Standard) +} + +/// 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_impl(data, outboard, ranges, encoded, Keyed(*key)) +} + +/// Generic encode body monomorphized over the compile time hashing strategy. +fn encode_ranges_validated_impl( + data: D, + outboard: O, + ranges: &ChunkRangesRef, + encoded: W, + hash_strategy: H, ) -> result::Result<(), EncodeError> { if ranges.is_empty() { return Ok(()); @@ -442,7 +580,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 = hash_strategy.parent_cv(&l_hash, &r_hash, is_root); let expected = stack.pop().unwrap(); if actual != expected { return Err(EncodeError::ParentHashMismatch(node)); @@ -481,10 +619,11 @@ pub fn encode_ranges_validated( tree.block_size.to_u32(), true, &mut out_buf, + hash_strategy, ); (actual, &out_buf[..]) } else { - let actual = hash_subtree(start_chunk.0, buf, is_root); + let actual = hash_strategy.hash_subtree(start_chunk.0, buf, is_root); #[allow(clippy::redundant_slicing)] (actual, &buf[..]) }; @@ -503,17 +642,56 @@ 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_impl(encoded, ranges, target, outboard, Standard) +} + +/// 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_impl(encoded, ranges, target, outboard, Keyed(*key)) +} + +/// Generic decode body monomorphized over the compile time hashing strategy. +fn decode_ranges_impl( encoded: R, ranges: &ChunkRangesRef, mut target: W, mut outboard: O, + hash_strategy: H, ) -> std::result::Result<(), DecodeError> where O: OutboardMut + Outboard, R: Read, W: WriteAt, { - let iter = DecodeResponseIter::new(outboard.root(), outboard.tree(), encoded, ranges); + 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 }) => { @@ -532,21 +710,41 @@ 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_hash_strategy(data, tree, outboard, Standard) +} + +/// 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_hash_strategy(data, tree, outboard, Keyed(*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, + hash_strategy: H, ) -> io::Result { let mut buffer = vec![0u8; tree.chunk_group_bytes()]; - let hash = outboard_impl(tree, data, &mut outboard, &mut buffer)?; - 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( +/// Generic outboard traversal monomorphized over the compile time hashing strategy. +fn outboard_impl( tree: BaoTree, mut data: impl Read, mut outboard: impl OutboardMut, buffer: &mut [u8], + hash_strategy: H, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -557,7 +755,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 = hash_strategy.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -568,7 +766,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_strategy.hash_subtree(start_chunk.0, buf, is_root); stack.push(hash); } } @@ -585,21 +783,41 @@ 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_hash_strategy(data, tree, outboard, Standard) +} + +/// 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_hash_strategy(data, tree, outboard, Keyed(*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, + 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)?; - 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( +/// Generic post order outboard traversal monomorphized over the compile time hashing strategy. +fn outboard_post_order_impl( tree: BaoTree, mut data: impl Read, mut outboard: impl Write, buffer: &mut [u8], + hash_strategy: H, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -611,7 +829,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 = hash_strategy.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -622,7 +840,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_strategy.hash_subtree(start_chunk.0, buf, is_root); stack.push(hash); } } @@ -663,8 +881,8 @@ mod validate { use super::Outboard; use crate::{ - blake3, hash_subtree, io::LocalBoxFuture, parent_cv, 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. @@ -677,33 +895,64 @@ mod validate { data: D, ranges: &'a ChunkRangesRef, ) -> impl IntoIterator>> + 'a + where + O: Outboard + 'a, + D: ReadAt + 'a, + { + valid_ranges_impl(outboard, data, ranges, Standard) + } + + /// 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_impl(outboard, data, ranges, Keyed(*key)) + } + + /// 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, + ranges: &'a ChunkRangesRef, + hash_strategy: H, + ) -> 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, 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>>, + hash_strategy: H, } - impl RecursiveDataValidator<'_, O, D> { + impl RecursiveDataValidator<'_, O, D, H> { async fn validate( outboard: O, data: D, ranges: &ChunkRangesRef, co: &Co>>, + hash_strategy: H, ) -> io::Result<()> { let tree = outboard.tree(); let mut buffer = vec![0u8; tree.chunk_group_bytes()]; @@ -711,7 +960,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_strategy.hash_subtree(0, tmp, true); if actual == outboard.root() { co.yield_(Ok(ChunkNum(0)..tree.chunks())).await; } @@ -727,6 +976,7 @@ mod validate { data, buffer, co, + hash_strategy, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -743,7 +993,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 = hash_subtree(ChunkNum::full_chunks(range.start).0, tmp, is_root); + let actual = + self.hash_strategy + .hash_subtree(ChunkNum::full_chunks(range.start).0, tmp, is_root); if &actual == hash { // yield the left range self.co @@ -777,7 +1029,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = parent_cv(&l_hash, &r_hash, is_root); + 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(()); @@ -879,7 +1131,7 @@ mod validate { // outboard is incomplete, we can't validate return Ok(()); }; - let actual = parent_cv(&l_hash, &r_hash, is_root); + let actual = Standard.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -905,4 +1157,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..134cdcb 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,7 +239,63 @@ 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 { +/// 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; +} + +/// BLAKE3 hash mode strategy for unkeyed APIs. +#[doc(hidden)] +#[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 strategy. Wraps a 32 byte key for domain separated hashing. +#[doc(hidden)] +#[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 { use blake3::hazmat::{ChainingValue, HasherExt}; if is_root { debug_assert!(start_chunk == 0); @@ -246,7 +309,33 @@ 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 { +/// 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 { + 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(); @@ -261,6 +350,24 @@ fn parent_cv(left_child: &blake3::Hash, right_child: &blake3::Hash, is_root: boo } } +/// 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 { + 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)) + } +} + /// Defines a Bao tree. /// /// This is just the specification of the tree, it does not contain any actual data. diff --git a/src/rec.rs b/src/rec.rs index 4673991..fc1c0a6 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,7 +96,10 @@ 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. -pub(crate) fn encode_selected_rec( +/// +/// `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, data: &[u8], is_root: bool, @@ -104,13 +107,14 @@ pub(crate) fn encode_selected_rec( min_level: u32, emit_data: bool, res: &mut Vec, + 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); } - 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(); @@ -142,6 +146,7 @@ pub(crate) fn encode_selected_rec( min_level, emit_data, res, + hash_strategy, ); let right = encode_selected_rec( mid_chunk, @@ -151,13 +156,14 @@ pub(crate) fn encode_selected_rec( min_level, emit_data, res, + 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()); } - parent_cv(&left, &right, is_root) + hash_strategy.parent_cv(&left, &right, is_root) } } @@ -275,6 +281,7 @@ mod test_support { 0, false, &mut res, + crate::Standard, ); (res, hash) } @@ -290,6 +297,7 @@ mod test_support { 0, true, &mut res, + crate::Standard, ); (res, hash) } @@ -430,10 +438,447 @@ mod test_support { block_size.to_u32(), true, &mut res, + crate::Standard, ); (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..08faab9 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -19,17 +19,433 @@ 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, + split, ChunkRanges, ChunkRangesRef, Keyed, ResponseIter, Standard, +}; + +#[cfg(feature = "tokio_fsm")] +use crate::rec::{ + keyed_create_sized_keyed_checks_fsm, keyed_init_from_keyed_checks_fsm, + keyed_outboard_functions_checks_fsm, }; +/// Reference encoder using the [Keyed] hashing strategy. +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, + Keyed(*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 +1045,7 @@ fn encode_selected_rec_cases() { min_level, true, &mut actual_encoded, + Standard, ); actual_encoded.len() - data.len() }; @@ -654,6 +1071,7 @@ fn encode_selected_reference( max_skip_level, true, &mut res, + Standard, ); (hash, res) } @@ -742,6 +1160,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 +1978,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..616d9b8 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, + BaoTree, BlockSize, ChunkNum, ChunkRanges, ChunkRangesRef, Standard, 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( @@ -766,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, @@ -783,6 +1317,7 @@ fn encode_selected_reference( max_skip_level, true, &mut res, + Standard, ); (hash, res) }