commonware_cryptography/blake3/
mod.rs

1//! BLAKE3 implementation of the [Hasher] trait.
2//!
3//! This implementation uses the [blake3] crate to generate BLAKE3 digests. [Hasher::hash_across]
4//! splits a message of at least 128 KiB along the BLAKE3 tree and hashes the subtrees across the
5//! given strategy.
6//!
7//! # Example
8//! ```rust
9//! use commonware_cryptography::{Hasher, blake3::Blake3};
10//! use commonware_parallel::Rayon;
11//! use std::num::NonZeroUsize;
12//!
13//! // Hash data in a single shot
14//! let digest = Blake3::hash(&[b"hello,", b"world!"]);
15//! println!("digest: {:?}", digest);
16//!
17//! // Or stream data incrementally
18//! let mut hasher = Blake3::default();
19//! hasher.update(b"hello,");
20//! hasher.update(b"world!");
21//! let (_hasher, digest) = hasher.finalize();
22//! println!("digest: {:?}", digest);
23//!
24//! // Hash a long message across a thread pool
25//! let strategy = Rayon::new(NonZeroUsize::new(4).unwrap()).unwrap();
26//! let message = vec![7u8; 1 << 20];
27//! assert_eq!(
28//!     Blake3::hash_across(&[&message], &strategy),
29//!     Blake3::hash(&[&message]),
30//! );
31//! ```
32
33use crate::Hasher;
34use blake3::{
35    CHUNK_LEN, Hash,
36    hazmat::{ChainingValue, HasherExt as _, Mode, merge_subtrees_non_root, merge_subtrees_root},
37};
38use bytes::BufMut;
39use commonware_codec::{Buf, Error as CodecError, FixedArray, FixedSize, Read, ReadExt, Write};
40use commonware_formatting::Hex;
41use commonware_math::algebra::Random;
42use commonware_parallel::Strategy;
43use commonware_utils::{Array, Span, sequence::FixedBytes};
44use core::{
45    cmp::Ordering,
46    fmt::{Debug, Display},
47    ops::Deref,
48};
49use rand_core::CryptoRng;
50use zeroize::Zeroize;
51
52/// Re-export [blake3::Hasher] as `CoreBlake3` for external use if needed.
53pub type CoreBlake3 = blake3::Hasher;
54
55const DIGEST_LENGTH: usize = blake3::OUT_LEN;
56
57/// Minimum number of tasks a split message spreads across: [`Hasher::hash_across`] bounds the
58/// subtree that each task hashes by the message's length divided by this.
59const SUBTREES: usize = 8;
60
61/// Lower clamp, in bytes, on the subtree bound: 16 chunks, the widest batch the SIMD kernels
62/// compress at once.
63const MIN_SUBTREE_LEN: usize = 16 * CHUNK_LEN;
64
65/// Upper clamp, in bytes, on the subtree bound: 64 chunks, enough to amortize the cost of a fork.
66const MAX_SUBTREE_LEN: usize = 64 * CHUNK_LEN;
67
68/// Length, in bytes, of the shortest message [`Hasher::hash_across`] splits: [`SUBTREES`]
69/// subtrees of [`MIN_SUBTREE_LEN`] bytes.
70const MIN_SPLIT_LEN: usize = SUBTREES * MIN_SUBTREE_LEN;
71
72/// Hash the `len`-byte concatenation of `parts`, which spans more than one chunk, across `strategy`
73/// as subtrees of at most `len / SUBTREES` bytes, clamped to [`MIN_SUBTREE_LEN`] and
74/// [`MAX_SUBTREE_LEN`].
75fn hash_subtrees(parts: &[&[u8]], len: usize, strategy: &impl Strategy) -> Digest {
76    let root = Node {
77        parts,
78        skip: 0,
79        offset: 0,
80        len,
81        task_len: (len / SUBTREES).clamp(MIN_SUBTREE_LEN, MAX_SUBTREE_LEN),
82    };
83    let (left, right) = root.children(strategy);
84    merge_subtrees_root(&left, &right, Mode::Hash).into()
85}
86
87/// A node of a message's BLAKE3 tree: the `len` bytes at byte `offset` of the message, which
88/// start `skip` bytes into the first of `parts`. Subtrees of at most `task_len` bytes hash as one
89/// task.
90///
91/// `parts` holds every byte of the node, and `skip` is at most the length of its first part.
92#[derive(Clone, Copy)]
93struct Node<'a> {
94    parts: &'a [&'a [u8]],
95    skip: usize,
96    offset: usize,
97    len: usize,
98    task_len: usize,
99}
100
101impl Node<'_> {
102    /// Hash the chaining values of the children of a node that spans more than one chunk, in
103    /// parallel across `strategy`.
104    fn children(self, strategy: &impl Strategy) -> (ChainingValue, ChainingValue) {
105        // The left child is the largest power-of-two number of chunks that leaves the right
106        // child nonempty, which for a node of more than one chunk is the largest power of two
107        // below its length. The right child starts in the first part that extends past it.
108        let left_len = 1 << (self.len - 1).ilog2();
109        let mut skip = self.skip + left_len;
110        let mut index = 0;
111        while skip >= self.parts[index].len() {
112            skip -= self.parts[index].len();
113            index += 1;
114        }
115        let left = Node {
116            parts: &self.parts[..=index],
117            len: left_len,
118            ..self
119        };
120        let right = Node {
121            parts: &self.parts[index..],
122            skip,
123            offset: self.offset + left_len,
124            len: self.len - left_len,
125            task_len: self.task_len,
126        };
127        strategy.join(
128            || left.chaining_value(strategy),
129            || right.chaining_value(strategy),
130        )
131    }
132
133    /// Hash the node's chaining value, splitting it until every subtree fits in `task_len` bytes.
134    fn chaining_value(self, strategy: &impl Strategy) -> ChainingValue {
135        if self.len > self.task_len {
136            let (left, right) = self.children(strategy);
137            return merge_subtrees_non_root(&left, &right, Mode::Hash);
138        }
139
140        // Feed the node's bytes, which start `skip` bytes into the first part.
141        let mut hasher = CoreBlake3::new();
142        hasher.set_input_offset(self.offset as u64);
143        let mut skip = self.skip;
144        let mut remaining = self.len;
145        for part in self.parts {
146            let bytes = &part[skip..];
147            let take = bytes.len().min(remaining);
148            hasher.update(&bytes[..take]);
149            remaining -= take;
150            if remaining == 0 {
151                break;
152            }
153            skip = 0;
154        }
155        hasher.finalize_non_root()
156    }
157}
158
159/// BLAKE3 hasher.
160#[derive(Debug, Default)]
161pub struct Blake3 {
162    hasher: CoreBlake3,
163}
164
165impl Hasher for Blake3 {
166    type Digest = Digest;
167
168    fn hash(parts: &[&[u8]]) -> Self::Digest {
169        let mut hasher = CoreBlake3::new();
170        for part in parts {
171            hasher.update(part);
172        }
173        hasher.finalize().into()
174    }
175
176    fn hash_across(parts: &[&[u8]], strategy: &impl Strategy) -> Self::Digest {
177        // A message shorter than `MIN_SPLIT_LEN` hashes on the calling thread. A total that
178        // overflows `usize` (possible only when parts alias) hashes serially, since the streaming
179        // hasher counts bytes in a `u64`.
180        match parts
181            .iter()
182            .try_fold(0usize, |len, part| len.checked_add(part.len()))
183        {
184            Some(len) if len >= MIN_SPLIT_LEN => strategy.run(
185                len,
186                || Self::hash(parts),
187                || hash_subtrees(parts, len, strategy),
188            ),
189            _ => Self::hash(parts),
190        }
191    }
192
193    fn hash_pair(
194        left: &[&[u8]],
195        right: &[&[u8]],
196        strategy: &impl Strategy,
197    ) -> (Self::Digest, Self::Digest) {
198        (
199            Self::hash_across(left, strategy),
200            Self::hash_across(right, strategy),
201        )
202    }
203
204    fn update(&mut self, message: &[u8]) -> &mut Self {
205        self.hasher.update(message);
206        self
207    }
208
209    fn finalize(mut self) -> (Self, Self::Digest) {
210        let finalized = self.hasher.finalize();
211        self.hasher.reset();
212        let array: [u8; DIGEST_LENGTH] = finalized.into();
213        (self, Self::Digest::from(array))
214    }
215}
216
217/// Digest of a BLAKE3 hashing operation.
218#[derive(Clone, Copy, Eq, PartialEq, Hash, FixedArray)]
219#[fixed_array(infallible)]
220#[repr(transparent)]
221pub struct Digest(pub [u8; DIGEST_LENGTH]);
222
223impl Ord for Digest {
224    #[inline]
225    fn cmp(&self, other: &Self) -> Ordering {
226        FixedBytes::new(self.0).cmp(&FixedBytes::new(other.0))
227    }
228}
229
230impl PartialOrd for Digest {
231    #[inline]
232    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
233        Some(self.cmp(other))
234    }
235}
236
237#[cfg(feature = "arbitrary")]
238impl<'a> arbitrary::Arbitrary<'a> for Digest {
239    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
240        // Generate random bytes and compute their Blake3 hash
241        let len = u.int_in_range(0..=256)?;
242        let data = u.bytes(len)?;
243        Ok(Blake3::hash(&[data]))
244    }
245}
246
247impl Write for Digest {
248    fn write(&self, buf: &mut impl BufMut) {
249        self.0.write(buf);
250    }
251}
252
253impl Read for Digest {
254    type Cfg = ();
255
256    fn read_cfg(buf: &mut impl Buf, _: &()) -> Result<Self, CodecError> {
257        let array = <[u8; DIGEST_LENGTH]>::read(buf)?;
258        Ok(Self(array))
259    }
260}
261
262impl FixedSize for Digest {
263    const SIZE: usize = DIGEST_LENGTH;
264}
265
266impl Span for Digest {}
267
268impl Array for Digest {}
269
270impl From<Hash> for Digest {
271    fn from(value: Hash) -> Self {
272        Self(value.into())
273    }
274}
275
276impl AsRef<[u8]> for Digest {
277    fn as_ref(&self) -> &[u8] {
278        &self.0
279    }
280}
281
282impl Deref for Digest {
283    type Target = [u8];
284    fn deref(&self) -> &[u8] {
285        &self.0
286    }
287}
288
289impl Debug for Digest {
290    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
291        write!(f, "{}", Hex(&self.0))
292    }
293}
294
295impl Display for Digest {
296    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
297        write!(f, "{}", Hex(&self.0))
298    }
299}
300
301impl crate::Digest for Digest {
302    const EMPTY: Self = Self([0u8; DIGEST_LENGTH]);
303}
304
305impl Random for Digest {
306    fn random(mut rng: impl CryptoRng) -> Self {
307        let mut array = [0u8; DIGEST_LENGTH];
308        rng.fill_bytes(&mut array);
309        Self(array)
310    }
311}
312
313impl Zeroize for Digest {
314    fn zeroize(&mut self) {
315        self.0.zeroize();
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322    use commonware_codec::{Copying, DecodeExt, Encode};
323    use commonware_parallel::{
324        Rayon, Sequential,
325        mocks::{Explorer, explore},
326        rayon::{ThreadPoolBuilder, Yield},
327    };
328    use commonware_utils::{NZUsize, TestRng};
329    use rand::Rng as _;
330    use std::sync::{
331        Arc,
332        atomic::{AtomicBool, Ordering::Relaxed},
333        mpsc,
334    };
335
336    const HELLO_DIGEST: [u8; DIGEST_LENGTH] = commonware_formatting::hex!(
337        "d74981efa70a0c880b8d8c1985d075dbcbf679b99a5f9914e5aaf96b831a9e24"
338    );
339
340    /// Return `len` random bytes, distinct per `seed`.
341    fn random(len: usize, seed: u64) -> Vec<u8> {
342        let mut bytes = vec![0; len];
343        TestRng::new(seed).fill_bytes(&mut bytes);
344        bytes
345    }
346
347    #[test]
348    fn test_blake3() {
349        let msg = b"hello world";
350
351        // Generate initial hash
352        let mut hasher = Blake3::default();
353        hasher.update(msg);
354        let (hasher, digest) = hasher.finalize();
355        assert!(Digest::decode(Copying(&digest)).is_ok());
356        assert_eq!(digest.as_ref(), HELLO_DIGEST);
357
358        // Reuse the reset hasher
359        let mut hasher = hasher;
360        hasher.update(msg);
361        let (_, digest) = hasher.finalize();
362        assert!(Digest::decode(Copying(&digest)).is_ok());
363        assert_eq!(digest.as_ref(), HELLO_DIGEST);
364
365        // Test one-shot hasher
366        let hash = Blake3::hash(&[msg]);
367        assert_eq!(hash.as_ref(), HELLO_DIGEST);
368
369        // Test multi-part one-shot hasher
370        let hash = Blake3::hash(&[b"hello", b" world"]);
371        assert_eq!(hash.as_ref(), HELLO_DIGEST);
372    }
373
374    /// Official BLAKE3 test vectors. Hashing 16 KiB or more in one update reaches the 16-way
375    /// AVX-512 chunk kernel, and 32 KiB or more also reaches the 16-way parent kernel.
376    #[test]
377    fn test_official_vectors() {
378        const VECTORS: [(usize, [u8; DIGEST_LENGTH]); 3] = [
379            (
380                16384,
381                commonware_formatting::hex!(
382                    "f875d6646de28985646f34ee13be9a576fd515f76b5b0a26bb324735041ddde4"
383                ),
384            ),
385            (
386                31744,
387                commonware_formatting::hex!(
388                    "62b6960e1a44bcc1eb1a611a8d6235b6b4b78f32e7abc4fb4c6cdcce94895c47"
389                ),
390            ),
391            (
392                102400,
393                commonware_formatting::hex!(
394                    "bc3e3d41a1146b069abffad3c0d44860cf664390afce4d9661f7902e7943e085"
395                ),
396            ),
397        ];
398        for (len, expected) in VECTORS {
399            // The official input repeats the bytes 0 through 250.
400            let input: Vec<u8> = (0..len).map(|i| (i % 251) as u8).collect();
401            let mut hasher = Blake3::default();
402            hasher.update(&input);
403            let (_, digest) = hasher.finalize();
404            assert_eq!(digest.as_ref(), expected, "len {len}");
405        }
406    }
407
408    /// Messages around multiples of [`MIN_SUBTREE_LEN`] and around [`MIN_SPLIT_LEN`], given whole,
409    /// cut inside chunks and subtrees, padded with empty parts, and paged, hash to the reference
410    /// digest on the calling thread and across workers. Splitting into subtrees on the calling
411    /// thread also matches, down to two chunks.
412    #[test]
413    fn test_hash_long_messages() {
414        let data = random((1 << 20) + 5000, 0);
415        for len in [
416            0,
417            1,
418            CHUNK_LEN,
419            CHUNK_LEN + 1,
420            3 * CHUNK_LEN + 7,
421            MIN_SUBTREE_LEN - 1,
422            MIN_SUBTREE_LEN,
423            MIN_SUBTREE_LEN + 1,
424            MIN_SUBTREE_LEN + CHUNK_LEN,
425            2 * MIN_SUBTREE_LEN - 1,
426            2 * MIN_SUBTREE_LEN,
427            2 * MIN_SUBTREE_LEN + 1,
428            2 * MIN_SUBTREE_LEN + 3 * CHUNK_LEN + 7,
429            3 * MIN_SUBTREE_LEN + 17,
430            MIN_SPLIT_LEN - 1,
431            MIN_SPLIT_LEN,
432            MIN_SPLIT_LEN + MIN_SUBTREE_LEN + 1,
433            3 * MIN_SPLIT_LEN + 17,
434            (1 << 20) + 5000,
435        ] {
436            let message = &data[..len];
437            let expected: Digest = blake3::hash(message).into();
438            let (a, b) = (len / 3, len / 3 + MIN_SUBTREE_LEN.min(len - len / 3));
439            let splits: [Vec<&[u8]>; 4] = [
440                vec![message],
441                vec![&message[..a], &message[a..b], &message[b..]],
442                vec![&[], message, &[]],
443                message.chunks(1000).collect(),
444            ];
445            for parts in &splits {
446                assert_eq!(Blake3::hash(parts), expected, "len={len}");
447                assert_eq!(
448                    explore(|strategy| Blake3::hash_across(parts, strategy)),
449                    expected,
450                    "len={len}"
451                );
452                if len > CHUNK_LEN {
453                    assert_eq!(
454                        hash_subtrees(parts, len, &Sequential),
455                        expected,
456                        "len={len}"
457                    );
458                }
459            }
460        }
461    }
462
463    /// A message of at least [`MIN_SPLIT_LEN`] bytes is hashed across the strategy's pool, and a
464    /// shorter one on the calling thread.
465    #[test]
466    fn test_hash_splits_long_messages() {
467        // One worker, planned as four, counts the jobs it runs while looping. Splitting a message
468        // from outside the pool hands it one job, whose nested forks stay on the worker.
469        let pool = Arc::new(ThreadPoolBuilder::new().num_threads(1).build().unwrap());
470        let strategy =
471            Explorer::parallel(&Rayon::with_pool(pool.clone()).with_parallelism(NZUsize!(4)));
472        let data = random(2 * MIN_SPLIT_LEN, 0);
473        for (len, jobs) in [
474            (MIN_SPLIT_LEN - 1, 0),
475            (MIN_SPLIT_LEN, 1),
476            (2 * MIN_SPLIT_LEN, 1),
477        ] {
478            let done = Arc::new(AtomicBool::new(false));
479            let (started, ready) = mpsc::channel();
480            let (sender, receiver) = mpsc::channel();
481            pool.spawn({
482                let done = done.clone();
483                let pool = pool.clone();
484                move || {
485                    started.send(()).unwrap();
486                    let mut jobs = 0;
487                    while !done.load(Relaxed) {
488                        if pool.yield_now() == Some(Yield::Executed) {
489                            jobs += 1;
490                        }
491                    }
492                    sender.send(jobs).unwrap();
493                }
494            });
495            ready.recv().unwrap();
496
497            // Hash, then release the worker and compare its job count.
498            let message = &data[..len];
499            let digest = Blake3::hash_across(&[message], &strategy);
500            done.store(true, Relaxed);
501            assert_eq!(digest, blake3::hash(message).into(), "len={len}");
502            assert_eq!(receiver.recv().unwrap(), jobs, "len={len}");
503        }
504    }
505
506    #[test]
507    fn test_blake3_len() {
508        assert_eq!(Digest::SIZE, DIGEST_LENGTH);
509    }
510
511    #[test]
512    fn test_codec() {
513        let msg = b"hello world";
514        let mut hasher = Blake3::default();
515        hasher.update(msg);
516        let (_, digest) = hasher.finalize();
517
518        let encoded = digest.encode();
519        assert_eq!(encoded.len(), DIGEST_LENGTH);
520        assert_eq!(encoded, digest.as_ref());
521
522        let decoded = Digest::decode(encoded).unwrap();
523        assert_eq!(digest, decoded);
524    }
525
526    #[test]
527    fn test_digest_ord() {
528        let a = Digest([0; DIGEST_LENGTH]);
529        let mut b = a;
530        b.0[DIGEST_LENGTH - 1] = 1;
531
532        // The first byte decides even when the remaining bytes order the other way.
533        let mut c = a;
534        c.0[0] = 0x7f;
535        c.0[1..].fill(0xff);
536        let mut d = a;
537        d.0[0] = 0x80;
538        for (a, b, expected) in [
539            (a, a, Ordering::Equal),
540            (a, b, Ordering::Less),
541            (b, a, Ordering::Greater),
542            (c, d, Ordering::Less),
543            (d, c, Ordering::Greater),
544        ] {
545            assert_eq!(a.cmp(&b), expected);
546            assert_eq!(a.partial_cmp(&b), Some(expected));
547        }
548    }
549
550    #[cfg(feature = "arbitrary")]
551    mod conformance {
552        use super::*;
553        use commonware_codec::conformance::CodecConformance;
554
555        commonware_conformance::conformance_tests! {
556            CodecConformance<Digest>,
557        }
558    }
559}
commonware_cryptography/blake3/mod.rs

commonware_cryptography/blake3/
mod.rs
