diff --git a/src/io/fsm.rs b/src/io/fsm.rs index 5c64e01..5d5f7cd 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, + BaoTree, BlockSize, ChunkRanges, ChunkRangesRef, HashMode, TreeNode, }; /// A binary merkle tree for blake3 hashes of a blob. @@ -318,27 +317,37 @@ struct ResponseDecoderInner { iter: ResponseIter, stack: SmallVec<[blake3::Hash; 10]>, encoded: R, + mode: HashMode, } impl ResponseDecoderInner { - fn new(tree: BaoTree, hash: blake3::Hash, ranges: ChunkRanges, encoded: R) -> Self { + fn with_mode( + tree: BaoTree, + hash: blake3::Hash, + ranges: ChunkRanges, + encoded: R, + mode: HashMode, + ) -> 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, + mode, }; res.stack.push(hash); res } } -/// Response decoder +/// Response decoder. +/// +/// Keyed callers should use [Self::new_keyed]. #[derive(Debug)] pub struct ResponseDecoder(Box>); -/// Next type for ResponseDecoder. +/// Next type for [ResponseDecoder]. #[derive(Debug)] pub enum ResponseDecoderNext { /// One more item, and you get back the state machine in the next state @@ -357,8 +366,40 @@ impl ResponseDecoder { /// /// The size as well as the chunk size is given in the `tree` parameter. pub fn new(hash: blake3::Hash, ranges: ChunkRanges, tree: BaoTree, encoded: R) -> Self { - Self(Box::new(ResponseDecoderInner::new( - tree, hash, ranges, encoded, + Self(Box::new(ResponseDecoderInner::with_mode( + tree, + hash, + ranges, + encoded, + HashMode::Standard, + ))) + } + + /// Create a new keyed response decoder. + pub fn new_keyed( + hash: blake3::Hash, + ranges: ChunkRanges, + tree: BaoTree, + encoded: R, + key: &[u8; 32], + ) -> Self { + Self(Box::new(ResponseDecoderInner::with_mode( + tree, + hash, + ranges, + encoded, + HashMode::Keyed(*key), + ))) + } + pub(crate) fn with_mode( + hash: blake3::Hash, + ranges: ChunkRanges, + tree: BaoTree, + encoded: R, + mode: HashMode, + ) -> Self { + Self(Box::new(ResponseDecoderInner::with_mode( + tree, hash, ranges, encoded, mode, ))) } @@ -404,7 +445,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.mode.parent_cv(&l_hash, &r_hash, is_root); // Push the children in reverse order so they are popped in the correct order // only push right if the range intersects with the right child if right { @@ -434,7 +475,7 @@ 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.mode.hash_subtree(start_chunk.0, &data, is_root); if leaf_hash != actual { return Err(DecodeError::LeafHashMismatch(start_chunk)); } @@ -501,10 +542,41 @@ where /// This will either succeed if the requested ranges are all present, or fail /// as soon as a range is missing. pub async fn encode_ranges_validated( + data: D, + outboard: O, + ranges: &ChunkRangesRef, + encoded: W, +) -> result::Result<(), EncodeError> +where + D: AsyncSliceReader, + O: Outboard, + W: AsyncStreamWriter, +{ + encode_ranges_validated_impl(data, outboard, ranges, encoded, HashMode::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, HashMode::Keyed(*key)).await +} + +async fn encode_ranges_validated_impl( mut data: D, mut outboard: O, ranges: &ChunkRangesRef, encoded: W, + mode: HashMode, ) -> result::Result<(), EncodeError> where D: AsyncSliceReader, @@ -529,7 +601,7 @@ where .. } => { let (l_hash, r_hash) = outboard.load(node).await?.unwrap(); - let actual = parent_cv(&l_hash, &r_hash, is_root); + let actual = mode.parent_cv(&l_hash, &r_hash, is_root); let expected = stack.pop().unwrap(); if actual != expected { return Err(EncodeError::ParentHashMismatch(node)); @@ -570,10 +642,11 @@ where tree.block_size.to_u32(), true, &mut out_buf, + mode, ); (actual, out_buf.clone().into()) } else { - let actual = hash_subtree(start_chunk.0, &bytes, is_root); + let actual = mode.hash_subtree(start_chunk.0, &bytes, is_root); (actual, bytes) }; if actual != expected { @@ -594,17 +667,49 @@ 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, HashMode::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, HashMode::Keyed(*key)).await +} + +async fn decode_ranges_impl( encoded: R, ranges: ChunkRanges, mut target: W, mut outboard: O, + mode: HashMode, ) -> 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_mode(outboard.root(), ranges, outboard.tree(), encoded, mode); loop { let item = match reading.next().await { ResponseDecoderNext::Done(_reader) => break, @@ -624,6 +729,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 +741,40 @@ 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_mode(data, tree, outboard, HashMode::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_mode(data, tree, outboard, HashMode::Keyed(*key)).await +} + +/// Allocates a chunk group buffer and delegates to [outboard_impl]. +async fn outboard_with_mode( data: impl AsyncStreamReader, tree: BaoTree, mut outboard: impl OutboardMut, + mode: HashMode, ) -> 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, mode).await } -/// Internal helper for [outboard_post_order]. This takes a buffer of the chunk group size. async fn outboard_impl( tree: BaoTree, mut data: impl AsyncStreamReader, mut outboard: impl OutboardMut, buffer: &mut [u8], + mode: HashMode, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -660,7 +785,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 = mode.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -670,7 +795,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 = mode.hash_subtree(start_chunk.0, &buf, is_root); stack.push(hash); } } @@ -687,21 +812,40 @@ 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_mode(data, tree, outboard, HashMode::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_mode(data, tree, outboard, HashMode::Keyed(*key)).await +} + +/// Allocates a chunk group buffer and delegates to [outboard_post_order_impl]. +async fn outboard_post_order_with_mode( data: impl AsyncStreamReader, tree: BaoTree, mut outboard: impl AsyncStreamWriter, + mode: HashMode, ) -> 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, mode).await } -/// Internal helper for [outboard_post_order]. This takes a buffer of the chunk group size. async fn outboard_post_order_impl( tree: BaoTree, mut data: impl AsyncStreamReader, mut outboard: impl AsyncStreamWriter, buffer: &mut [u8], + mode: HashMode, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -713,7 +857,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 = mode.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -723,7 +867,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 = mode.hash_subtree(start_chunk.0, &buf, is_root); stack.push(hash); } } @@ -757,8 +901,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, BaoTree, ChunkNum, ChunkRangesRef, + HashMode, TreeNode, }; /// Given a data file and an outboard, compute all valid ranges. @@ -771,12 +915,40 @@ mod validate { data: D, ranges: &'a ChunkRangesRef, ) -> impl Stream>> + 'a + where + O: Outboard + 'a, + D: AsyncSliceReader + 'a, + { + valid_ranges_impl(outboard, data, ranges, HashMode::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: &[u8; 32], + ) -> impl Stream>> + 'a + where + O: Outboard + 'a, + D: AsyncSliceReader + 'a, + { + valid_ranges_impl(outboard, data, ranges, HashMode::Keyed(*key)) + } + + fn valid_ranges_impl<'a, O, D>( + outboard: O, + data: D, + ranges: &'a ChunkRangesRef, + mode: HashMode, + ) -> 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, mode).await { co.yield_(Err(cause)).await; } @@ -789,6 +961,7 @@ mod validate { outboard: O, data: D, co: &'a Co>>, + mode: HashMode, } impl RecursiveDataValidator<'_, O, D> { @@ -797,6 +970,7 @@ mod validate { data: D, ranges: &ChunkRangesRef, co: &Co>>, + mode: HashMode, ) -> io::Result<()> { let tree = outboard.tree(); if tree.blocks() == 1 { @@ -805,7 +979,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 = mode.hash_subtree(0, &data, true); if actual == outboard.root() { co.yield_(Ok(ChunkNum(0)..tree.chunks())).await; } @@ -820,6 +994,7 @@ mod validate { outboard, data, co, + mode, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -835,7 +1010,9 @@ mod validate { let len = (range.end - range.start).try_into().unwrap(); let data = self.data.read_exact_at(range.start, len).await?; // is_root is always false because the case of a single chunk group is handled before calling this function - let actual = hash_subtree(ChunkNum::full_chunks(range.start).0, &data, is_root); + let actual = + self.mode + .hash_subtree(ChunkNum::full_chunks(range.start).0, &data, is_root); if &actual == hash { // yield the left range self.co @@ -869,7 +1046,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.mode.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -902,11 +1079,38 @@ mod validate { outboard: O, ranges: &'a ChunkRangesRef, ) -> impl Stream>> + 'a + where + O: Outboard + 'a, + { + valid_outboard_ranges_impl(outboard, ranges, HashMode::Standard) + } + + /// Given just a keyed outboard, compute all valid ranges. + /// + /// This is not cheap since it recomputes the hashes for all chunks. + pub fn keyed_valid_outboard_ranges<'a, O>( + outboard: O, + ranges: &'a ChunkRangesRef, + key: &[u8; 32], + ) -> impl Stream>> + 'a + where + O: Outboard + 'a, + { + valid_outboard_ranges_impl(outboard, ranges, HashMode::Keyed(*key)) + } + + fn valid_outboard_ranges_impl<'a, O>( + outboard: O, + ranges: &'a ChunkRangesRef, + mode: HashMode, + ) -> impl Stream>> + 'a where O: Outboard + 'a, { Gen::new(move |co| async move { - if let Err(cause) = RecursiveOutboardValidator::validate(outboard, ranges, &co).await { + if let Err(cause) = + RecursiveOutboardValidator::validate(outboard, ranges, &co, mode).await + { co.yield_(Err(cause)).await; } }) @@ -917,6 +1121,7 @@ mod validate { shifted_filled_size: TreeNode, outboard: O, co: &'a Co>>, + mode: HashMode, } impl RecursiveOutboardValidator<'_, O> { @@ -924,6 +1129,7 @@ mod validate { outboard: O, ranges: &ChunkRangesRef, co: &Co>>, + mode: HashMode, ) -> io::Result<()> { let tree = outboard.tree(); if tree.blocks() == 1 { @@ -939,6 +1145,7 @@ mod validate { shifted_filled_size, outboard, co, + mode, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -972,7 +1179,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.mode.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -998,4 +1205,6 @@ mod validate { } } #[cfg(feature = "validate")] -pub use validate::{valid_outboard_ranges, valid_ranges}; +pub use validate::{ + keyed_valid_outboard_ranges, keyed_valid_ranges, valid_outboard_ranges, valid_ranges, +}; diff --git a/src/io/mixed.rs b/src/io/mixed.rs index 0f9f623..ed09e48 100644 --- a/src/io/mixed.rs +++ b/src/io/mixed.rs @@ -8,8 +8,7 @@ use smallvec::SmallVec; use super::{sync::Outboard, EncodeError, Leaf, Parent}; use crate::{ - hash_subtree, iter::BaoChunk, parent_cv, rec::truncate_ranges, split_inner, ChunkNum, - ChunkRangesRef, TreeNode, + iter::BaoChunk, rec::truncate_ranges, split_inner, ChunkNum, ChunkRangesRef, HashMode, TreeNode, }; /// A content item for the bao streaming protocol. @@ -84,122 +83,143 @@ where O: Outboard, F: Sender, { - send.send(EncodedItem::Size(outboard.tree().size())).await?; - let res = match traverse_ranges_validated_impl(data, outboard, ranges, send).await { - Ok(Ok(())) => EncodedItem::Done, - Err(cause) => EncodedItem::Error(cause), - Ok(Err(err)) => return Err(err), - }; - send.send(res).await + traverse_ranges_validated_impl(data, outboard, ranges, send, HashMode::Standard).await } -/// Encode ranges relevant to a query from a reader and outboard to a writer +/// Traverse ranges relevant to a query from a reader and keyed outboard to a stream /// -/// This function validates the data before writing. +/// This function validates the data before writing, using BLAKE3 keyed hashing. /// /// It is possible to encode ranges from a partial file and outboard. /// This will either succeed if the requested ranges are all present, or fail /// as soon as a range is missing. +pub async fn keyed_traverse_ranges_validated( + data: D, + outboard: O, + ranges: &ChunkRangesRef, + send: &mut F, + key: &[u8; 32], +) -> std::result::Result<(), F::Error> +where + D: ReadBytesAt, + O: Outboard, + F: Sender, +{ + traverse_ranges_validated_impl(data, outboard, ranges, send, HashMode::Keyed(*key)).await +} + async fn traverse_ranges_validated_impl( data: D, outboard: O, ranges: &ChunkRangesRef, send: &mut F, -) -> result::Result, EncodeError> + mode: HashMode, +) -> std::result::Result<(), F::Error> where D: ReadBytesAt, O: Outboard, F: Sender, { - if ranges.is_empty() { - return Ok(Ok(())); - } - let mut stack: SmallVec<[_; 10]> = SmallVec::<[blake3::Hash; 10]>::new(); - stack.push(outboard.root()); - let data = data; - let tree = outboard.tree(); - // canonicalize ranges - let ranges = truncate_ranges(ranges, tree.size()); - for item in tree.ranges_pre_order_chunks_iter_ref(ranges, 0) { - match item { - BaoChunk::Parent { - is_root, - left, - right, - node, - .. - } => { - let (l_hash, r_hash) = outboard.load(node)?.unwrap(); - let actual = parent_cv(&l_hash, &r_hash, is_root); - let expected = stack.pop().unwrap(); - if actual != expected { - return Err(EncodeError::ParentHashMismatch(node)); - } - if right { - stack.push(r_hash); - } - if left { - stack.push(l_hash); - } - let item = Parent { + send.send(EncodedItem::Size(outboard.tree().size())).await?; + let res: result::Result, EncodeError> = async { + if ranges.is_empty() { + return Ok(Ok(())); + } + let mut stack: SmallVec<[_; 10]> = SmallVec::<[blake3::Hash; 10]>::new(); + stack.push(outboard.root()); + let data = data; + let tree = outboard.tree(); + // canonicalize ranges + let ranges = truncate_ranges(ranges, tree.size()); + for item in tree.ranges_pre_order_chunks_iter_ref(ranges, 0) { + match item { + BaoChunk::Parent { + is_root, + left, + right, node, - pair: (l_hash, r_hash), - }; - if let Err(e) = send.send(item.into()).await { - return Ok(Err(e)); - } - } - BaoChunk::Leaf { - start_chunk, - size, - is_root, - ranges, - .. - } => { - let expected = stack.pop().unwrap(); - let start = start_chunk.to_bytes(); - let buffer = data.read_bytes_at(start, size)?; - if !ranges.is_all() { - // we need to encode just a part of the data - // - // write into an out buffer to ensure we detect mismatches - // before writing to the output. - let mut out_buf = Vec::new(); - let actual = traverse_selected_rec( - start_chunk, - buffer, - is_root, - ranges, - tree.block_size.to_u32(), - true, - &mut out_buf, - ); + .. + } => { + let (l_hash, r_hash) = outboard.load(node)?.unwrap(); + let actual = mode.parent_cv(&l_hash, &r_hash, is_root); + let expected = stack.pop().unwrap(); if actual != expected { - return Err(EncodeError::LeafHashMismatch(start_chunk)); + return Err(EncodeError::ParentHashMismatch(node)); } - for item in out_buf.into_iter() { - if let Err(e) = send.send(item).await { - return Ok(Err(e)); - } + if right { + stack.push(r_hash); } - } else { - let actual = hash_subtree(start_chunk.0, &buffer, is_root); - #[allow(clippy::redundant_slicing)] - if actual != expected { - return Err(EncodeError::LeafHashMismatch(start_chunk)); + if left { + stack.push(l_hash); } - let item = Leaf { - data: buffer, - offset: start_chunk.to_bytes(), + let item = Parent { + node, + pair: (l_hash, r_hash), }; if let Err(e) = send.send(item.into()).await { return Ok(Err(e)); } - }; + } + BaoChunk::Leaf { + start_chunk, + size, + is_root, + ranges, + .. + } => { + let expected = stack.pop().unwrap(); + let start = start_chunk.to_bytes(); + let buffer = data.read_bytes_at(start, size)?; + if !ranges.is_all() { + // we need to encode just a part of the data + // + // write into an out buffer to ensure we detect mismatches + // before writing to the output. + let mut out_buf = Vec::new(); + let actual = traverse_selected_rec_impl( + start_chunk, + buffer, + is_root, + ranges, + tree.block_size.to_u32(), + true, + &mut out_buf, + mode, + ); + if actual != expected { + return Err(EncodeError::LeafHashMismatch(start_chunk)); + } + for item in out_buf.into_iter() { + if let Err(e) = send.send(item).await { + return Ok(Err(e)); + } + } + } else { + let actual = mode.hash_subtree(start_chunk.0, &buffer, is_root); + #[allow(clippy::redundant_slicing)] + if actual != expected { + return Err(EncodeError::LeafHashMismatch(start_chunk)); + } + let item = Leaf { + data: buffer, + offset: start_chunk.to_bytes(), + }; + if let Err(e) = send.send(item.into()).await { + return Ok(Err(e)); + } + }; + } } } + Ok(Ok(())) } - Ok(Ok(())) + .await; + let res = match res { + Ok(Ok(())) => EncodedItem::Done, + Err(cause) => EncodedItem::Error(cause), + Ok(Err(err)) => return Err(err), + }; + send.send(res).await } /// Encode ranges relevant to a query from a slice and outboard to a buffer. @@ -207,7 +227,7 @@ where /// This will compute the root hash, so it will have to traverse the entire tree. /// The `ranges` parameter just controls which parts of the data are written. /// -/// Except for writing to a buffer, this is the same as [hash_subtree]. +/// Except for writing to a buffer, this is the same as computing the subtree hash. /// The `min_level` parameter controls the minimum level that will be emitted as a leaf. /// Set this to 0 to disable chunk groups entirely. /// The `emit_data` parameter controls whether the data is written to the buffer. @@ -228,6 +248,53 @@ pub fn traverse_selected_rec( min_level: u32, emit_data: bool, res: &mut Vec, +) -> blake3::Hash { + traverse_selected_rec_impl( + start_chunk, + data, + is_root, + query, + min_level, + emit_data, + res, + HashMode::Standard, + ) +} + +/// Keyed version of [traverse_selected_rec], using BLAKE3 keyed hashing. +#[allow(clippy::too_many_arguments)] +pub fn keyed_traverse_selected_rec( + start_chunk: ChunkNum, + data: Bytes, + is_root: bool, + query: &ChunkRangesRef, + min_level: u32, + emit_data: bool, + res: &mut Vec, + key: &[u8; 32], +) -> blake3::Hash { + traverse_selected_rec_impl( + start_chunk, + data, + is_root, + query, + min_level, + emit_data, + res, + HashMode::Keyed(*key), + ) +} + +#[allow(clippy::too_many_arguments)] +fn traverse_selected_rec_impl( + start_chunk: ChunkNum, + data: Bytes, + is_root: bool, + query: &ChunkRangesRef, + min_level: u32, + emit_data: bool, + res: &mut Vec, + mode: HashMode, ) -> blake3::Hash { use blake3::CHUNK_LEN; if data.len() <= CHUNK_LEN { @@ -240,7 +307,7 @@ pub fn traverse_selected_rec( .into(), ); } - hash_subtree(start_chunk.0, &data, is_root) + mode.hash_subtree(start_chunk.0, &data, is_root) } else { let chunks = data.len() / CHUNK_LEN + (data.len() % CHUNK_LEN != 0) as usize; let chunks = chunks.next_power_of_two(); @@ -268,7 +335,7 @@ pub fn traverse_selected_rec( None }; // recurse to the left and right to compute the hashes and emit data - let left = traverse_selected_rec( + let left = traverse_selected_rec_impl( start_chunk, data.slice(..mid_bytes), false, @@ -276,8 +343,9 @@ pub fn traverse_selected_rec( min_level, emit_data, res, + mode, ); - let right = traverse_selected_rec( + let right = traverse_selected_rec_impl( mid_chunk, data.slice(mid_bytes..), false, @@ -285,6 +353,7 @@ pub fn traverse_selected_rec( min_level, emit_data, res, + mode, ); // backfill the hashes if needed if let Some(o) = hash_offset { @@ -296,7 +365,7 @@ pub fn traverse_selected_rec( } .into(); } - parent_cv(&left, &right, is_root) + mode.parent_cv(&left, &right, is_root) } } @@ -304,7 +373,10 @@ pub fn traverse_selected_rec( mod tests { use super::*; use crate::{ - io::{outboard::PreOrderMemOutboard, sync::encode_ranges_validated}, + io::{ + outboard::PreOrderMemOutboard, + sync::{encode_ranges_validated, keyed_encode_ranges_validated}, + }, BlockSize, ChunkRanges, }; @@ -343,6 +415,40 @@ mod tests { let encoded2 = flatten(res); assert_eq!(encoded, encoded2); } + + #[tokio::test] + async fn keyed_smoke() { + let data = [0u8; 100000]; + let key = blake3::derive_key("bao-tree.test", b"mixed"); + let outboard = PreOrderMemOutboard::create_keyed(data, BlockSize::from_chunk_log(4), &key); + let (mut tx, mut rx) = tokio::sync::mpsc::channel(10); + let mut encoded = Vec::new(); + keyed_encode_ranges_validated( + &data[..], + &outboard, + &ChunkRanges::all(), + &mut encoded, + &key, + ) + .unwrap(); + tokio::spawn(async move { + keyed_traverse_ranges_validated( + &data[..], + &outboard, + &ChunkRanges::all(), + &mut tx, + &key, + ) + .await + .unwrap(); + }); + let mut res = Vec::new(); + while let Some(item) = rx.recv().await { + res.push(item); + } + let encoded2 = flatten(res); + assert_eq!(encoded, encoded2); + } } /// Trait identical to `ReadAt` but returning `Bytes` instead of reading into a buffer. 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..cdf378c 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, + BaoTree, BlockSize, ChunkRangesRef, HashMode, TreeNode, }; /// A binary merkle tree for blake3 hashes of a blob. @@ -257,13 +256,16 @@ 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 [Self::new_keyed]. #[derive(Debug)] pub struct DecodeResponseIter<'a, R> { inner: ResponseIterRef<'a>, stack: SmallVec<[blake3::Hash; 10]>, encoded: R, buf: BytesMut, + mode: HashMode, } impl<'a, R: Read> DecodeResponseIter<'a, R> { @@ -286,6 +288,43 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { encoded: R, ranges: &'a ChunkRangesRef, buf: BytesMut, + ) -> Self { + DecodeResponseIter::with_mode(root, tree, encoded, ranges, buf, HashMode::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], + ) -> Self { + let buf = BytesMut::with_capacity(tree.block_size().bytes()); + Self::new_keyed_with_buffer(root, tree, encoded, ranges, buf, key) + } + + /// Create a new iterator to decode a keyed response. + /// + /// This is the same as [Self::new_keyed], but allows you to provide a buffer to use for decoding. + /// The buffer will be resized as needed, but it's capacity should be the [crate::BlockSize::bytes]. + pub fn new_keyed_with_buffer( + root: blake3::Hash, + tree: BaoTree, + encoded: R, + ranges: &'a ChunkRangesRef, + buf: BytesMut, + key: &[u8; 32], + ) -> Self { + DecodeResponseIter::with_mode(root, tree, encoded, ranges, buf, HashMode::Keyed(*key)) + } + pub(crate) fn with_mode( + root: blake3::Hash, + tree: BaoTree, + encoded: R, + ranges: &'a ChunkRangesRef, + buf: BytesMut, + mode: HashMode, ) -> Self { let ranges = truncate_ranges(ranges, tree.size()); let mut stack = SmallVec::new(); @@ -295,6 +334,7 @@ impl<'a, R: Read> DecodeResponseIter<'a, R> { inner: ResponseIterRef::new(tree, ranges), encoded, buf, + mode, } } @@ -322,7 +362,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.mode.parent_cv(&l_hash, &r_hash, is_root); if parent_hash != actual { return Err(DecodeError::ParentHashMismatch(node)); } @@ -344,7 +384,7 @@ 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.mode.hash_subtree(start_chunk.0, &self.buf, is_root); let leaf_hash = self.stack.pop().unwrap(); if leaf_hash != actual { return Err(DecodeError::LeafHashMismatch(start_chunk)); @@ -419,6 +459,27 @@ pub fn encode_ranges_validated( outboard: O, ranges: &ChunkRangesRef, encoded: W, +) -> result::Result<(), EncodeError> { + encode_ranges_validated_impl(data, outboard, ranges, encoded, HashMode::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, HashMode::Keyed(*key)) +} + +fn encode_ranges_validated_impl( + data: D, + outboard: O, + ranges: &ChunkRangesRef, + encoded: W, + mode: HashMode, ) -> result::Result<(), EncodeError> { if ranges.is_empty() { return Ok(()); @@ -442,7 +503,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 = mode.parent_cv(&l_hash, &r_hash, is_root); let expected = stack.pop().unwrap(); if actual != expected { return Err(EncodeError::ParentHashMismatch(node)); @@ -481,10 +542,11 @@ pub fn encode_ranges_validated( tree.block_size.to_u32(), true, &mut out_buf, + mode, ); (actual, &out_buf[..]) } else { - let actual = hash_subtree(start_chunk.0, buf, is_root); + let actual = mode.hash_subtree(start_chunk.0, buf, is_root); #[allow(clippy::redundant_slicing)] (actual, &buf[..]) }; @@ -503,17 +565,55 @@ 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, HashMode::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, HashMode::Keyed(*key)) +} + +fn decode_ranges_impl( encoded: R, ranges: &ChunkRangesRef, mut target: W, mut outboard: O, + mode: HashMode, ) -> 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_mode( + outboard.root(), + outboard.tree(), + encoded, + ranges, + BytesMut::with_capacity(outboard.tree().block_size().bytes()), + mode, + ); for item in iter { match item? { BaoContentItem::Parent(Parent { node, pair }) => { @@ -532,21 +632,40 @@ 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_mode(data, tree, outboard, HashMode::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_mode(data, tree, outboard, HashMode::Keyed(*key)) +} + +/// Allocates a chunk group buffer and delegates to [outboard_impl]. +fn outboard_with_mode( data: impl Read, tree: BaoTree, mut outboard: impl OutboardMut, + mode: HashMode, ) -> 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, mode) } -/// Internal helper for [outboard_post_order]. This takes a buffer of the chunk group size. fn outboard_impl( tree: BaoTree, mut data: impl Read, mut outboard: impl OutboardMut, buffer: &mut [u8], + mode: HashMode, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -557,7 +676,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 = mode.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -568,7 +687,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 = mode.hash_subtree(start_chunk.0, buf, is_root); stack.push(hash); } } @@ -585,21 +704,40 @@ 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_mode(data, tree, outboard, HashMode::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_mode(data, tree, outboard, HashMode::Keyed(*key)) +} + +/// Allocates a chunk group buffer and delegates to [outboard_post_order_impl]. +fn outboard_post_order_with_mode( data: impl Read, tree: BaoTree, mut outboard: impl Write, + mode: HashMode, ) -> 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, mode) } -/// Internal helper for [outboard_post_order]. This takes a buffer of the chunk group size. fn outboard_post_order_impl( tree: BaoTree, mut data: impl Read, mut outboard: impl Write, buffer: &mut [u8], + mode: HashMode, ) -> io::Result { // do not allocate for small trees let mut stack = SmallVec::<[blake3::Hash; 10]>::new(); @@ -611,7 +749,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 = mode.parent_cv(&left_hash, &right_hash, is_root); stack.push(parent); } BaoChunk::Leaf { @@ -622,7 +760,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 = mode.hash_subtree(start_chunk.0, buf, is_root); stack.push(hash); } } @@ -663,8 +801,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, BaoTree, ChunkNum, ChunkRangesRef, + HashMode, TreeNode, }; /// Given a data file and an outboard, compute all valid ranges. @@ -677,12 +815,40 @@ mod validate { data: D, ranges: &'a ChunkRangesRef, ) -> impl IntoIterator>> + 'a + where + O: Outboard + 'a, + D: ReadAt + 'a, + { + valid_ranges_impl(outboard, data, ranges, HashMode::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: &[u8; 32], + ) -> impl IntoIterator>> + 'a + where + O: Outboard + 'a, + D: ReadAt + 'a, + { + valid_ranges_impl(outboard, data, ranges, HashMode::Keyed(*key)) + } + + fn valid_ranges_impl<'a, O, D>( + outboard: O, + data: D, + ranges: &'a ChunkRangesRef, + mode: HashMode, + ) -> 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, mode).await { co.yield_(Err(cause)).await; } @@ -696,6 +862,7 @@ mod validate { data: D, buffer: Vec, co: &'a Co>>, + mode: HashMode, } impl RecursiveDataValidator<'_, O, D> { @@ -704,6 +871,7 @@ mod validate { data: D, ranges: &ChunkRangesRef, co: &Co>>, + mode: HashMode, ) -> io::Result<()> { let tree = outboard.tree(); let mut buffer = vec![0u8; tree.chunk_group_bytes()]; @@ -711,7 +879,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 = mode.hash_subtree(0, tmp, true); if actual == outboard.root() { co.yield_(Ok(ChunkNum(0)..tree.chunks())).await; } @@ -727,6 +895,7 @@ mod validate { data, buffer, co, + mode, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -743,7 +912,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 + .mode + .hash_subtree(ChunkNum::full_chunks(range.start).0, tmp, is_root); if &actual == hash { // yield the left range self.co @@ -777,7 +948,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.mode.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -809,11 +980,38 @@ mod validate { outboard: O, ranges: &'a ChunkRangesRef, ) -> impl IntoIterator>> + 'a + where + O: Outboard + 'a, + { + valid_outboard_ranges_impl(outboard, ranges, HashMode::Standard) + } + + /// Given just a keyed outboard, compute all valid ranges. + /// + /// This is not cheap since it recomputes the hashes for all chunks. + pub fn keyed_valid_outboard_ranges<'a, O>( + outboard: O, + ranges: &'a ChunkRangesRef, + key: &[u8; 32], + ) -> impl IntoIterator>> + 'a + where + O: Outboard + 'a, + { + valid_outboard_ranges_impl(outboard, ranges, HashMode::Keyed(*key)) + } + + fn valid_outboard_ranges_impl<'a, O>( + outboard: O, + ranges: &'a ChunkRangesRef, + mode: HashMode, + ) -> impl IntoIterator>> + 'a where O: Outboard + 'a, { Gen::new(move |co| async move { - if let Err(cause) = RecursiveOutboardValidator::validate(outboard, ranges, &co).await { + if let Err(cause) = + RecursiveOutboardValidator::validate(outboard, ranges, &co, mode).await + { co.yield_(Err(cause)).await; } }) @@ -824,6 +1022,7 @@ mod validate { shifted_filled_size: TreeNode, outboard: O, co: &'a Co>>, + mode: HashMode, } impl RecursiveOutboardValidator<'_, O> { @@ -831,6 +1030,7 @@ mod validate { outboard: O, ranges: &ChunkRangesRef, co: &Co>>, + mode: HashMode, ) -> io::Result<()> { let tree = outboard.tree(); if tree.blocks() == 1 { @@ -846,6 +1046,7 @@ mod validate { shifted_filled_size, outboard, co, + mode, }; validator .validate_rec(&root_hash, shifted_root, true, ranges) @@ -879,7 +1080,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.mode.parent_cv(&l_hash, &r_hash, is_root); if &actual != parent_hash { // hash mismatch, we can't validate return Ok(()); @@ -905,4 +1106,6 @@ mod validate { } } #[cfg(feature = "validate")] -pub use validate::{valid_outboard_ranges, valid_ranges}; +pub use validate::{ + keyed_valid_outboard_ranges, keyed_valid_ranges, valid_outboard_ranges, valid_ranges, +}; diff --git a/src/lib.rs b/src/lib.rs index b0d8ec1..4848ada 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,33 +239,91 @@ 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 { - use blake3::hazmat::{ChainingValue, HasherExt}; - if is_root { - debug_assert!(start_chunk == 0); - blake3::hash(data) - } else { - let mut hasher = blake3::Hasher::new(); - hasher.set_input_offset(start_chunk * 1024); - hasher.update(data); - let non_root_hash: ChainingValue = hasher.finalize_non_root(); - blake3::Hash::from(non_root_hash) +/// Hashing mode for shared encode and decode paths, either standard or keyed BLAKE3. +#[derive(Clone, Copy)] +pub(crate) enum HashMode { + Standard, + Keyed([u8; 32]), +} + +impl std::fmt::Debug for HashMode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + HashMode::Standard => f.write_str("Standard"), + HashMode::Keyed(_) => f.debug_struct("Keyed").finish_non_exhaustive(), + } } } -fn parent_cv(left_child: &blake3::Hash, right_child: &blake3::Hash, is_root: bool) -> blake3::Hash { - use blake3::hazmat::{merge_subtrees_non_root, merge_subtrees_root, ChainingValue, Mode}; - let left_child: ChainingValue = *left_child.as_bytes(); - let right_child: ChainingValue = *right_child.as_bytes(); - if is_root { - merge_subtrees_root(&left_child, &right_child, Mode::Hash) - } else { - blake3::Hash::from(merge_subtrees_non_root( - &left_child, - &right_child, - Mode::Hash, - )) +impl HashMode { + pub(crate) fn hash_subtree( + &self, + start_chunk: u64, + data: &[u8], + is_root: bool, + ) -> blake3::Hash { + use blake3::hazmat::{ChainingValue, HasherExt}; + if is_root { + debug_assert!(start_chunk == 0); + match self { + HashMode::Standard => blake3::hash(data), + HashMode::Keyed(key) => blake3::keyed_hash(key, data), + } + } else { + let mut hasher = match self { + HashMode::Standard => blake3::Hasher::new(), + HashMode::Keyed(key) => blake3::Hasher::new_keyed(key), + }; + hasher.set_input_offset(start_chunk * 1024); + hasher.update(data); + let non_root_hash: ChainingValue = hasher.finalize_non_root(); + blake3::Hash::from(non_root_hash) + } } + + pub(crate) fn parent_cv( + &self, + 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(); + let mode = match self { + HashMode::Standard => Mode::Hash, + HashMode::Keyed(key) => Mode::KeyedHash(key), + }; + if is_root { + merge_subtrees_root(&left_child, &right_child, mode) + } else { + blake3::Hash::from(merge_subtrees_non_root(&left_child, &right_child, mode)) + } + } +} + +/// Compute the hash of a subtree using BLAKE3 keyed mode. +/// +/// See [keyed_parent_cv] for merging child hashes in keyed mode. +#[inline] +pub fn keyed_hash_subtree( + start_chunk: u64, + data: &[u8], + is_root: bool, + key: &[u8; 32], +) -> blake3::Hash { + HashMode::Keyed(*key).hash_subtree(start_chunk, data, is_root) +} + +/// Merge two child subtree hashes using BLAKE3 keyed mode. +#[inline] +pub fn keyed_parent_cv( + left_child: &blake3::Hash, + right_child: &blake3::Hash, + is_root: bool, + key: &[u8; 32], +) -> blake3::Hash { + HashMode::Keyed(*key).parent_cv(left_child, right_child, is_root) } /// Defines a Bao tree. @@ -547,7 +612,7 @@ pub(crate) const fn blocks(size: u64, block_size: BlockSize) -> u64 { /// You typically don't have to use this, but it can be useful for debugging /// and error handling. Hash validation errors contain a `TreeNode` that allows /// you to find the position where validation failed. -#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub struct TreeNode(u64); diff --git a/src/rec.rs b/src/rec.rs index 4673991..a279680 100644 --- a/src/rec.rs +++ b/src/rec.rs @@ -2,7 +2,7 @@ //! //! Encocding is used to compute hashes, decoding is only used in tests as a //! reference implementation. -use crate::{blake3, hash_subtree, parent_cv, split_inner, ChunkNum, ChunkRangesRef}; +use crate::{blake3, split_inner, ChunkNum, ChunkRangesRef}; /// Given a set of chunk ranges, adapt them for a tree of the given size. /// @@ -96,6 +96,9 @@ 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. +/// +/// `mode` is the hashing mode for subtree and parent hashes. +#[allow(clippy::too_many_arguments)] pub(crate) fn encode_selected_rec( start_chunk: ChunkNum, data: &[u8], @@ -104,13 +107,14 @@ pub(crate) fn encode_selected_rec( min_level: u32, emit_data: bool, res: &mut Vec, + mode: crate::HashMode, ) -> blake3::Hash { use blake3::CHUNK_LEN; if data.len() <= CHUNK_LEN { if emit_data && !query.is_empty() { res.extend_from_slice(data); } - hash_subtree(start_chunk.0, data, is_root) + mode.hash_subtree(start_chunk.0, data, is_root) } else { let chunks = data.len() / CHUNK_LEN + (data.len() % CHUNK_LEN != 0) as usize; let chunks = chunks.next_power_of_two(); @@ -142,6 +146,7 @@ pub(crate) fn encode_selected_rec( min_level, emit_data, res, + mode, ); let right = encode_selected_rec( mid_chunk, @@ -151,13 +156,14 @@ pub(crate) fn encode_selected_rec( min_level, emit_data, res, + mode, ); // backfill the hashes if needed if let Some(o) = hash_offset { res[o..o + 32].copy_from_slice(left.as_bytes()); res[o + 32..o + 64].copy_from_slice(right.as_bytes()); } - parent_cv(&left, &right, is_root) + mode.parent_cv(&left, &right, is_root) } } @@ -275,6 +281,7 @@ mod test_support { 0, false, &mut res, + crate::HashMode::Standard, ); (res, hash) } @@ -290,6 +297,7 @@ mod test_support { 0, true, &mut res, + crate::HashMode::Standard, ); (res, hash) } @@ -430,10 +438,189 @@ mod test_support { block_size.to_u32(), true, &mut res, + crate::HashMode::Standard, ); (res, hash) } + use std::io::Cursor; + + use crate::io::outboard::{PostOrderMemOutboard, PreOrderMemOutboard, PreOrderOutboard}; + use crate::io::sync; + + 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); + } + + 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); + + let truncated_size = 1024u64.min(data.len() as u64); + if truncated_size < data.len() as u64 { + let truncated_tree = BaoTree::new(truncated_size, BlockSize(0)); + let mut truncated_pre = PreOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree: truncated_tree, + data: vec![0; truncated_tree.outboard_size().try_into().unwrap()], + }; + let root = + sync::keyed_outboard(Cursor::new(data), truncated_tree, &mut truncated_pre, key) + .unwrap(); + truncated_pre.root = root; + assert_eq!( + root, + blake3::keyed_hash(key, &data[..truncated_size as usize]) + ); + assert_pre_order_outboard_matches_mem( + &truncated_pre, + &data[..truncated_size as usize], + BlockSize(0), + key, + ); + } + + let oversize_tree = BaoTree::new(data.len() as u64 + 100, BlockSize(0)); + let mut sink = Vec::new(); + assert!( + sync::keyed_outboard_post_order(Cursor::new(data), oversize_tree, &mut sink, key) + .is_err() + ); + } + + #[cfg(feature = "tokio_fsm")] + 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); + + let truncated_size = 1024u64.min(data.len() as u64); + if truncated_size < data.len() as u64 { + let truncated_tree = BaoTree::new(truncated_size, BlockSize(0)); + let mut truncated_pre = PreOrderOutboard { + root: blake3::Hash::from([0; 32]), + tree: truncated_tree, + data: vec![0; truncated_tree.outboard_size().try_into().unwrap()], + }; + let root = keyed_outboard( + Cursor::new(Bytes::from(data.to_vec())), + truncated_tree, + &mut truncated_pre, + key, + ) + .await + .unwrap(); + truncated_pre.root = root; + assert_eq!( + root, + blake3::keyed_hash(key, &data[..truncated_size as usize]) + ); + assert_pre_order_outboard_matches_mem( + &truncated_pre, + &data[..truncated_size as usize], + BlockSize(0), + key, + ); + } + + let oversize_tree = BaoTree::new(data.len() as u64 + 100, BlockSize(0)); + let mut sink = Vec::new(); + assert!(keyed_outboard_post_order( + Cursor::new(Bytes::from(data.to_vec())), + oversize_tree, + &mut sink, + key + ) + .await + .is_err()); + } + /// Check that l and r of a 2-tuple are equal #[macro_export] macro_rules! assert_tuple_eq { diff --git a/src/tests.rs b/src/tests.rs index 4959f69..5c3482e 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -20,16 +20,423 @@ use super::{ }; use crate::{ assert_tuple_eq, blake3, - io::{full_chunk_groups, outboard::PreOrderMemOutboard, sync::Outboard, BaoContentItem, Leaf}, + io::{ + full_chunk_groups, outboard::PreOrderMemOutboard, sync::Outboard, BaoContentItem, + DecodeError, EncodeError, Leaf, + }, iter::{PostOrderChunkIter, PreOrderPartialIterRef, ResponseIterRef}, - prop_assert_tuple_eq, + keyed_hash_subtree, keyed_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_outboard_functions_checks, + make_test_data, range_union, truncate_ranges, ReferencePreOrderPartialChunkIterRef, }, - split, ChunkRanges, ChunkRangesRef, ResponseIter, + split, ChunkRanges, ChunkRangesRef, HashMode, ResponseIter, }; +#[cfg(feature = "tokio_fsm")] +use crate::rec::keyed_outboard_functions_checks_fsm; + +/// Reference encoder using BLAKE3 keyed mode. +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, + HashMode::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 rt = tokio::runtime::Runtime::new().unwrap(); + let mut outboard = PostOrderMemOutboard::create_keyed(&data, block_size, key); + let ranges = ChunkRanges::all(); + let mut encoded = Vec::new(); + rt.block_on(keyed_encode_ranges_validated( + Bytes::from(data.clone()), + &mut outboard, + &ranges, + &mut encoded, + key, + )) + .unwrap(); + let tree = outboard.tree(); + let mut decoded = bytes::BytesMut::new(); + let mut ob_res = PostOrderMemOutboard { + root: outboard.root(), + tree, + data: vec![0; tree.outboard_size().try_into().unwrap()], + }; + rt.block_on(keyed_decode_ranges( + Cursor::new(encoded.as_slice()), + ranges, + &mut decoded, + &mut ob_res, + key, + )) + .unwrap(); + assert_eq!(decoded.to_vec(), data); + assert_eq!(ob_res.root(), outboard.root()); +} + +/// 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 +1036,7 @@ fn encode_selected_rec_cases() { min_level, true, &mut actual_encoded, + HashMode::Standard, ); actual_encoded.len() - data.len() }; @@ -654,6 +1062,7 @@ fn encode_selected_reference( max_skip_level, true, &mut res, + HashMode::Standard, ); (hash, res) } @@ -742,6 +1151,412 @@ 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() { + let key = blake3::derive_key("bao-tree.test", b"roundtrip"); + keyed_encode_decode_roundtrip_sync_impl(&make_test_data(50_000), BlockSize(2), &key); +} + +#[test] +fn keyed_encode_decode_roundtrip_fsm() { + let key = blake3::derive_key("bao-tree.test", b"roundtrip"); + keyed_encode_decode_roundtrip_fsm_impl(make_test_data(50_000), BlockSize(2), &key); +} + +#[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 = HashMode::Standard.hash_subtree(0, &data, true); + let keyed = keyed_hash_subtree(0, &data, true, &key); + assert_ne!(standard, keyed); + assert_eq!(keyed, blake3::keyed_hash(&key, &data)); + let non_root_standard = HashMode::Standard.hash_subtree(1, &data[..1024], false); + let non_root_keyed = keyed_hash_subtree(1, &data[..1024], false, &key); + assert_ne!(non_root_standard, non_root_keyed); + let mut hasher = blake3::Hasher::new_keyed(&key); + 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 = HashMode::Standard.parent_cv(&left, &right, true); + let keyed = keyed_parent_cv(&left, &right, true, &key); + assert_ne!(standard, keyed); + let standard_non_root = HashMode::Standard.parent_cv(&left, &right, false); + let keyed_non_root = keyed_parent_cv(&left, &right, false, &key); + assert_ne!(standard_non_root, keyed_non_root); + let left_cv: ChainingValue = *left.as_bytes(); + 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_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_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); + } +} + +const KEYED_SLICE_ROUNDTRIP_CASES: [(usize, std::ops::Range); 8] = [ + (0, 0..1), + (1, 0..1), + (1023, 0..1), + (1024, 0..1), + (1025, 0..1), + (1025, 0..2), + (1025, 1..2), + (24 * 1024 + 1, 0..25), +]; + +#[test] +fn keyed_bao_tree_slice_roundtrip_cases() { + let key = blake3::derive_key("bao-tree.test", b"slice"); + for chunk_group_log in 0..4 { + let block_size = BlockSize(chunk_group_log); + for (count, range) in KEYED_SLICE_ROUNDTRIP_CASES { + keyed_bao_tree_slice_roundtrip_test( + make_test_data(count), + ChunkNum(range.start)..ChunkNum(range.end), + block_size, + &key, + ); + } + } +} + +#[cfg(feature = "tokio_fsm")] +#[tokio::test] +async fn keyed_bao_tree_slice_roundtrip_fsm_cases() { + let key = blake3::derive_key("bao-tree.test", b"slice-fsm"); + for chunk_group_log in 0..4 { + let block_size = BlockSize(chunk_group_log); + for (count, range) in KEYED_SLICE_ROUNDTRIP_CASES { + keyed_bao_tree_slice_roundtrip_fsm_test( + make_test_data(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 +1730,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..394d04e 100644 --- a/src/tests2.rs +++ b/src/tests2.rs @@ -17,7 +17,7 @@ use smallvec::SmallVec; use test_strategy::proptest; use crate::{ - assert_tuple_eq, blake3, hash_subtree, + assert_tuple_eq, blake3, io::{ fsm::ResponseDecoderNext, outboard::{PostOrderMemOutboard, PreOrderMemOutboard}, @@ -25,15 +25,19 @@ use crate::{ BaoContentItem, Leaf, Parent, }, iter::{BaoChunk, PreOrderPartialChunkIterRef, ResponseIterRef}, - parent_cv, prop_assert_tuple_eq, + 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_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, HashMode, 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,8 +145,8 @@ fn post_traversal_chunks_iter_proptest(#[strategy(tree())] tree: BaoTree) { post_traversal_chunks_iter_impl(tree); } -/// 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) { +/// Brute force test for an outboard that computes the expected hash for each pair +fn outboard_test_sync(data: &[u8], outboard: impl crate::io::sync::Outboard, mode: HashMode) { let tree = outboard.tree(); let nodes = tree .pre_order_nodes_iter() @@ -155,14 +159,18 @@ fn outboard_test_sync(data: &[u8], outboard: impl crate::io::sync::Outboard) { let start_chunk = node.chunk_range().start; let byte_range = tree.byte_range(node); let data = &data[byte_range.start.try_into().unwrap()..byte_range.end.try_into().unwrap()]; - let expected = hash_subtree(start_chunk.0, data, is_root); - let actual = parent_cv(&l_hash, &r_hash, is_root); + let expected = mode.hash_subtree(start_chunk.0, data, is_root); + let actual = mode.parent_cv(&l_hash, &r_hash, is_root); assert_eq!(actual, expected); } } -/// Brute force test for an outboard that just computes the expected hash for each pair -async fn outboard_test_fsm(data: &[u8], mut outboard: impl crate::io::fsm::Outboard) { +/// Brute force test for an outboard that computes the expected hash for each pair +async fn outboard_test_fsm( + data: &[u8], + mut outboard: impl crate::io::fsm::Outboard, + mode: HashMode, +) { let tree = outboard.tree(); let nodes = tree .pre_order_nodes_iter() @@ -175,8 +183,8 @@ async fn outboard_test_fsm(data: &[u8], mut outboard: impl crate::io::fsm::Outbo let start_chunk = node.chunk_range().start; let byte_range = tree.byte_range(node); let data = &data[byte_range.start.try_into().unwrap()..byte_range.end.try_into().unwrap()]; - let expected = hash_subtree(start_chunk.0, data, is_root); - let actual = parent_cv(&l_hash, &r_hash, is_root); + let expected = mode.hash_subtree(start_chunk.0, data, is_root); + let actual = mode.parent_cv(&l_hash, &r_hash, is_root); assert_eq!(actual, expected); } } @@ -188,7 +196,7 @@ fn post_oder_outboard_sync_impl(tree: BaoTree) { outboard.data.len() as u64, outboard.tree().outboard_hash_pairs() * 64 ); - outboard_test_sync(&data, outboard); + outboard_test_sync(&data, outboard, HashMode::Standard); } #[test] @@ -214,7 +222,7 @@ fn post_oder_outboard_fsm_impl(tree: BaoTree) { ); tokio::runtime::Runtime::new() .unwrap() - .block_on(outboard_test_fsm(&data, outboard)); + .block_on(outboard_test_fsm(&data, outboard, HashMode::Standard)); } #[proptest] @@ -222,6 +230,81 @@ 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 + ); + outboard_test_sync(&data, outboard, HashMode::Keyed(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(outboard_test_fsm(&data, outboard, HashMode::Keyed(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() + ); + outboard_test_sync(&data, outboard, HashMode::Keyed(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(outboard_test_fsm(&data, outboard, HashMode::Keyed(key))); +} + +#[proptest] +fn keyed_pre_order_outboard_fsm_proptest(#[strategy(tree())] tree: BaoTree) { + keyed_pre_order_outboard_fsm_impl(tree); +} + +#[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); +} + 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 +362,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 { @@ -320,6 +436,87 @@ mod validate { } } + fn keyed_valid_outboard_ranges_sync( + outboard: impl crate::io::sync::Outboard, + key: &[u8; 32], + ) -> ChunkRanges { + let ranges = ChunkRanges::all(); + let iter = crate::io::sync::keyed_valid_outboard_ranges(outboard, &ranges, key); + let mut res = ChunkRanges::empty(); + for item in iter { + res |= ChunkRanges::from(item.unwrap()); + } + res + } + + fn keyed_valid_outboard_ranges_fsm( + outboard: &mut PostOrderMemOutboard, + key: &[u8; 32], + ) -> ChunkRanges { + run_blocking(async move { + let ranges = ChunkRanges::all(); + let mut stream = crate::io::fsm::keyed_valid_outboard_ranges(outboard, &ranges, key); + let mut res = ChunkRanges::empty(); + while let Some(item) = stream.next().await { + let item = item?; + res |= ChunkRanges::from(item); + } + std::io::Result::Ok(res) + }) + .unwrap() + } + + fn validate_keyed_outboard_pos_impl(tree: BaoTree) { + let size = tree.size.try_into().unwrap(); + let block_size = tree.block_size; + let data = make_test_data(size); + let key = blake3::derive_key("bao-tree.test", b"valid-outboard-ranges"); + let mut outboard = PostOrderMemOutboard::create_keyed(data, block_size, &key); + let expected = ChunkRanges::from(..outboard.tree().chunks()); + let actual = keyed_valid_outboard_ranges_sync(&mut outboard, &key); + assert_eq!(expected, actual); + let actual = keyed_valid_outboard_ranges_fsm(&mut outboard, &key); + assert_eq!(expected, actual) + } + + #[proptest] + fn validate_keyed_outboard_pos_proptest(#[strategy(tree())] tree: BaoTree) { + validate_keyed_outboard_pos_impl(tree); + } + + #[test] + fn validate_keyed_outboard_pos_cases() { + let cases = [(0x10001, 0)]; + for (size, block_level) in cases { + let tree = BaoTree::new(size, BlockSize(block_level)); + validate_keyed_outboard_pos_impl(tree); + } + } + + /// Wrong key must not report a multi-block keyed outboard as fully valid. + fn validate_keyed_outboard_wrong_key_impl(tree: BaoTree) { + let size = tree.size.try_into().unwrap(); + let block_size = tree.block_size; + let data = make_test_data(size); + let key = blake3::derive_key("bao-tree.test", b"valid-outboard-ranges"); + let wrong_key = blake3::derive_key("bao-tree.test", b"wrong-key"); + let mut outboard = PostOrderMemOutboard::create_keyed(data, block_size, &key); + let expected = ChunkRanges::from(..outboard.tree().chunks()); + let actual = keyed_valid_outboard_ranges_sync(&mut outboard, &wrong_key); + assert_ne!(expected, actual); + let actual = keyed_valid_outboard_ranges_fsm(&mut outboard, &wrong_key); + assert_ne!(expected, actual) + } + + #[test] + fn validate_keyed_outboard_wrong_key_cases() { + let cases = [(0x10001, 0), (0x2001, 0), (5000, 1)]; + for (size, block_level) in cases { + let tree = BaoTree::new(size, BlockSize(block_level)); + validate_keyed_outboard_wrong_key_impl(tree); + } + } + fn validate_pos_impl(tree: BaoTree) { let size = tree.size.try_into().unwrap(); let block_size = tree.block_size; @@ -349,6 +546,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 +849,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 +1146,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 +1204,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 +1300,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 unkeyed BLAKE3 mode. fn encode_selected_reference( data: &[u8], block_size: BlockSize, @@ -783,6 +1319,7 @@ fn encode_selected_reference( max_skip_level, true, &mut res, + HashMode::Standard, ); (hash, res) }