Skip to main content

ab_blake3/
single_block.rs

1//! BLAKE3 functions that process at most a single block.
2//!
3//! This module and submodules are copied with modifications from the official [`blake3`] crate, but
4//! are unlikely to be upstreamed.
5
6#[cfg(test)]
7mod tests;
8
9use crate::platform::{le_bytes_from_words_32, words_from_le_bytes_32};
10use crate::{
11    BLOCK_LEN, BlockBytes, BlockWords, CHUNK_END, CHUNK_START, CVWords, DERIVE_KEY_CONTEXT,
12    DERIVE_KEY_MATERIAL, IV, KEY_LEN, KEYED_HASH, OUT_LEN, ROOT, portable,
13};
14use blake3::IncrementCounter;
15use blake3::platform::Platform;
16
17/// Hash single block worth of values
18#[inline(always)]
19fn hash_block(input: &[u8], key: CVWords, flags: u8) -> Option<[u8; OUT_LEN]> {
20    // If the whole subtree is one block, hash it directly with a ChunkState.
21    if input.len() > BLOCK_LEN {
22        return None;
23    }
24
25    let mut cv = key;
26
27    let mut block = [0; BLOCK_LEN];
28    block[..input.len()].copy_from_slice(input);
29    Platform::detect().compress_in_place(
30        &mut cv,
31        &block,
32        input.len() as u8,
33        0,
34        flags | CHUNK_START | CHUNK_END | ROOT,
35    );
36
37    Some(*le_bytes_from_words_32(&cv))
38}
39
40/// Hash multiple single block-sized inputs
41#[inline(always)]
42fn hash_block_many_exact<const NUM_BLOCKS: usize>(
43    inputs: &[BlockBytes; NUM_BLOCKS],
44    // TODO: `&mut [MaybeUninit<[u8; OUT_LEN]>; N]` would make more sense, but doesn't match
45    //  `blake3` API
46    outputs: &mut [[u8; OUT_LEN]; NUM_BLOCKS],
47    key: CVWords,
48    flags: u8,
49) {
50    let platform = Platform::detect();
51
52    let (input_chunks, remaining_inputs) = inputs.as_chunks::<16>();
53    let (output_chunks, remaining_output_chunks) = outputs.as_chunks_mut::<16>();
54
55    for (inputs, outputs) in input_chunks.iter().zip(output_chunks) {
56        // TODO: This is a very awkward API, ideally we wouldn't have this array allocated inline
57        //  for no good reason
58        platform.hash_many(
59            &inputs.each_ref(),
60            &key,
61            0,
62            IncrementCounter::No,
63            flags | CHUNK_START | CHUNK_END | ROOT,
64            0,
65            0,
66            outputs.as_flattened_mut(),
67        );
68    }
69
70    for (input, output) in remaining_inputs.iter().zip(remaining_output_chunks) {
71        let mut cv = key;
72
73        platform.compress_in_place(
74            &mut cv,
75            input,
76            BLOCK_LEN as u8,
77            0,
78            flags | CHUNK_START | CHUNK_END | ROOT,
79        );
80
81        output.copy_from_slice(le_bytes_from_words_32(&cv));
82    }
83}
84
85/// Hashing function for at most single block worth of bytes.
86///
87/// Returns `None` if the input length exceeds one block.
88#[inline]
89#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
90pub fn single_block_hash(input: &[u8]) -> Option<[u8; OUT_LEN]> {
91    hash_block(input, *IV, 0)
92}
93
94/// Hashing function for many single-block inputs
95#[inline]
96#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
97pub fn single_block_hash_many_exact<const NUM_BLOCKS: usize>(
98    inputs: &[BlockBytes; NUM_BLOCKS],
99    // TODO: `&mut [MaybeUninit<[u8; OUT_LEN]>; N]` would make more sense, but doesn't match
100    //  `blake3` API
101    outputs: &mut [[u8; OUT_LEN]; NUM_BLOCKS],
102) {
103    hash_block_many_exact(inputs, outputs, *IV, 0);
104}
105
106/// The keyed hash function for at most single block worth of bytes.
107///
108/// Returns `None` if the input length exceeds one block.
109#[inline]
110#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
111pub fn single_block_keyed_hash(key: &[u8; KEY_LEN], input: &[u8]) -> Option<[u8; OUT_LEN]> {
112    let key_words = words_from_le_bytes_32(key);
113    hash_block(input, key_words, KEYED_HASH)
114}
115
116/// Keyed hash function for many single-block inputs
117#[inline]
118#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
119pub fn single_block_keyed_hash_many_exact<const NUM_BLOCKS: usize>(
120    key: &[u8; KEY_LEN],
121    inputs: &[BlockBytes; NUM_BLOCKS],
122    // TODO: `&mut [MaybeUninit<[u8; OUT_LEN]>; N]` would make more sense, but doesn't match
123    //  `blake3` API
124    outputs: &mut [[u8; OUT_LEN]; NUM_BLOCKS],
125) {
126    let key_words = words_from_le_bytes_32(key);
127    hash_block_many_exact(inputs, outputs, key_words, KEYED_HASH);
128}
129
130// The key derivation function for at most a single block worth of bytes.
131//
132// Returns `None` if either context or key material length exceed one block.
133#[inline]
134#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
135pub fn single_block_derive_key(context: &str, key_material: &[u8]) -> Option<[u8; OUT_LEN]> {
136    let context_key = hash_block(context.as_bytes(), *IV, DERIVE_KEY_CONTEXT)?;
137    let context_key_words = words_from_le_bytes_32(&context_key);
138    hash_block(key_material, context_key_words, DERIVE_KEY_MATERIAL)
139}
140
141/// Hashing function for at most single block worth of words using portable implementation.
142///
143/// This API operates on words and is GPU-friendly.
144///
145/// `num_bytes` specifies how many actual bytes are occupied by useful value in `input`. Bytes
146/// outside that must be set to `0`.
147///
148/// NOTE: If unused bytes are not set to `0` or an invalid number of bytes is specified, it'll
149/// simply result in an invalid hash.
150///
151/// [`words_from_le_bytes_32()`], [`words_from_le_bytes_64()`] and [`le_bytes_from_words_32()`] can
152/// be used to convert bytes to words and back if necessary.
153///
154/// [`words_from_le_bytes_64()`]: crate::words_from_le_bytes_64
155#[inline]
156#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
157pub fn single_block_hash_portable_words(input: &BlockWords, num_bytes: u32) -> CVWords {
158    let mut cv = *IV;
159
160    portable::compress_in_place_u32(
161        &mut cv,
162        input,
163        num_bytes,
164        0,
165        u32::from(CHUNK_START | CHUNK_END | ROOT),
166    );
167
168    cv
169}