Skip to main content

fuchsia_merkle/
util.rs

1// Copyright 2018 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use bssl_crypto::digest::Sha256;
6use std::mem::size_of;
7use storage_ptr_slice::PtrByteSlice;
8
9use crate::{BLOCK_SIZE, HASH_SIZE, Hash};
10
11pub const HASHES_PER_BLOCK: usize = BLOCK_SIZE / HASH_SIZE;
12
13type BlockIdentity = [u8; size_of::<u64>() + size_of::<u32>()];
14
15/// Generate the bytes representing a block's identity.
16#[inline]
17fn make_identity(length: usize, level: usize, offset: usize) -> BlockIdentity {
18    let [o0, o1, o2, o3, o4, o5, o6, o7] = (offset as u64 | level as u64).to_le_bytes();
19    let [l0, l1, l2, l3] = (length as u32).to_le_bytes();
20    [o0, o1, o2, o3, o4, o5, o6, o7, l0, l1, l2, l3]
21}
22
23/// Compute the merkle hash of a block of data.
24///
25/// A merkle hash is the SHA-256 hash of a block of data with a small header built from the length
26/// of the data, the level of the tree (0 for data blocks), and the offset into the level. The
27/// block will be zero filled if its len is less than [`BLOCK_SIZE`], except for when the first
28/// data block is completely empty.
29///
30/// # Panics
31///
32/// Panics if `block.len()` exceeds [`BLOCK_SIZE`] or if `offset` is not aligned to [`BLOCK_SIZE`]
33pub fn hash_block(block: &[u8], offset: usize) -> Hash {
34    assert!(block.len() <= BLOCK_SIZE);
35    assert!(offset.is_multiple_of(BLOCK_SIZE));
36
37    let mut hasher = Sha256::new();
38    hasher.update(&make_identity(block.len(), 0, offset));
39    hasher.update(block);
40    // Zero fill block up to BLOCK_SIZE. As a special case, if the first data block is completely
41    // empty, it is not zero filled.
42    if block.len() != BLOCK_SIZE && !(block.is_empty() && offset == 0) {
43        update_with_zeros(&mut hasher, BLOCK_SIZE - block.len());
44    }
45
46    Hash::from(hasher.digest())
47}
48
49/// Hashes a single aligned block (4096 or 8192 bytes) from `data`.
50///
51/// `unaligned_len` specifies the number of valid bytes in the block
52/// (`0 < unaligned_len <= buffer_len`).
53/// Note: Any data in `data` beyond `unaligned_len` (up to `buffer_len`) MUST be zeroed
54/// for the computed hash to match.
55///
56/// Reads directly through raw pointers in `data` without minting Rust slice references,
57/// avoiding aliasing/provenance UB on shared DMA memory.
58///
59/// # Panics
60///
61/// Panics if `buffer_len` is not 4096 or 8192, or if `unaligned_len == 0` or
62/// `unaligned_len > buffer_len`.
63pub fn hash_block_aligned(
64    offset: usize,
65    data: PtrByteSlice<'_>,
66    buffer_len: usize,
67    unaligned_len: usize,
68) -> Hash {
69    assert!(buffer_len == BLOCK_SIZE || buffer_len == 4096);
70    assert!(unaligned_len > 0 && unaligned_len <= buffer_len);
71    assert!(offset.is_multiple_of(BLOCK_SIZE));
72    assert!(data.len() >= buffer_len);
73
74    #[cfg(all(target_arch = "aarch64", target_feature = "sha2"))]
75    {
76        arm64::hash_block_arm64_neon(offset, data, buffer_len, unaligned_len)
77    }
78
79    #[cfg(not(all(target_arch = "aarch64", target_feature = "sha2")))]
80    {
81        #[repr(C)]
82        struct Sha256Ctx {
83            _h: [u32; 8],
84            _nl: u32,
85            _nh: u32,
86            _data: [u8; 64],
87            _num: core::ffi::c_uint,
88            _md_len: core::ffi::c_uint,
89        }
90
91        unsafe extern "C" {
92            fn SHA256_Init(sha: *mut Sha256Ctx) -> core::ffi::c_int;
93            fn SHA256_Update(
94                sha: *mut Sha256Ctx,
95                data: *const core::ffi::c_void,
96                len: usize,
97            ) -> core::ffi::c_int;
98            fn SHA256_Final(out: *mut u8, sha: *mut Sha256Ctx) -> core::ffi::c_int;
99        }
100
101        let mut ctx = std::mem::MaybeUninit::<Sha256Ctx>::uninit();
102        // SAFETY: SHA256_Init initializes the context. data is valid for reading buffer_len bytes.
103        unsafe {
104            SHA256_Init(ctx.as_mut_ptr());
105            let identity = make_identity(unaligned_len, 0, offset);
106            SHA256_Update(
107                ctx.as_mut_ptr(),
108                identity.as_ptr() as *const core::ffi::c_void,
109                identity.len(),
110            );
111            SHA256_Update(ctx.as_mut_ptr(), data.as_ptr() as *const core::ffi::c_void, buffer_len);
112            if buffer_len < BLOCK_SIZE {
113                const ZEROS: [u8; 4096] = [0; 4096];
114                SHA256_Update(
115                    ctx.as_mut_ptr(),
116                    ZEROS.as_ptr() as *const core::ffi::c_void,
117                    BLOCK_SIZE - buffer_len,
118                );
119            }
120            let mut out = [0u8; HASH_SIZE];
121            SHA256_Final(out.as_mut_ptr(), ctx.as_mut_ptr());
122            Hash::from_array(out)
123        }
124    }
125}
126
127#[cfg(all(target_arch = "aarch64", target_feature = "sha2"))]
128mod arm64 {
129    use super::{BLOCK_SIZE, make_identity};
130    use crate::Hash;
131    use core::arch::aarch64::*;
132    use storage_ptr_slice::PtrByteSlice;
133
134    #[repr(C, align(16))]
135    struct Aligned64U32([u32; 64]);
136
137    #[repr(C, align(16))]
138    struct Aligned4U32([u32; 4]);
139
140    static K256: Aligned64U32 = Aligned64U32([
141        0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4,
142        0xab1c5ed5, 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe,
143        0x9bdc06a7, 0xc19bf174, 0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f,
144        0x4a7484aa, 0x5cb0a9dc, 0x76f988da, 0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7,
145        0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967, 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc,
146        0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, 0xa2bfe8a1, 0xa81a664b,
147        0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070, 0x19a4c116,
148        0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
149        0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7,
150        0xc67178f2,
151    ]);
152
153    static SHA256_INIT_H_0_3: Aligned4U32 =
154        Aligned4U32([0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a]);
155    static SHA256_INIT_H_4_7: Aligned4U32 =
156        Aligned4U32([0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19]);
157
158    #[derive(Clone, Copy)]
159    struct KTable {
160        k0: uint32x4_t,
161        k1: uint32x4_t,
162        k2: uint32x4_t,
163        k3: uint32x4_t,
164        k4: uint32x4_t,
165        k5: uint32x4_t,
166        k6: uint32x4_t,
167        k7: uint32x4_t,
168        k8: uint32x4_t,
169        k9: uint32x4_t,
170        k10: uint32x4_t,
171        k11: uint32x4_t,
172        k12: uint32x4_t,
173        k13: uint32x4_t,
174        k14: uint32x4_t,
175        k15: uint32x4_t,
176    }
177
178    impl KTable {
179        #[inline(always)]
180        fn load() -> Self {
181            // SAFETY: Reading 64 32-bit constants from 16-byte aligned static array.
182            unsafe {
183                let p = K256.0.as_ptr();
184                Self {
185                    k0: vld1q_u32(p),
186                    k1: vld1q_u32(p.add(4)),
187                    k2: vld1q_u32(p.add(8)),
188                    k3: vld1q_u32(p.add(12)),
189                    k4: vld1q_u32(p.add(16)),
190                    k5: vld1q_u32(p.add(20)),
191                    k6: vld1q_u32(p.add(24)),
192                    k7: vld1q_u32(p.add(28)),
193                    k8: vld1q_u32(p.add(32)),
194                    k9: vld1q_u32(p.add(36)),
195                    k10: vld1q_u32(p.add(40)),
196                    k11: vld1q_u32(p.add(44)),
197                    k12: vld1q_u32(p.add(48)),
198                    k13: vld1q_u32(p.add(52)),
199                    k14: vld1q_u32(p.add(56)),
200                    k15: vld1q_u32(p.add(60)),
201                }
202            }
203        }
204    }
205
206    struct Sha256VectorState {
207        abcd: uint32x4_t,
208        efgh: uint32x4_t,
209    }
210
211    // Single SHA-256 round with message scheduling:
212    // 1. Adds current schedule word + round constant (W_t + K_t).
213    // 2. Begins message schedule step (sha256su0) for future round W_{t+4}.
214    // 3. Preserves `abcd` in 1 cycle via bitwise OR (`orr tmp, abcd, abcd`) for `sha256h2`.
215    // 4. Computes compression rounds (sha256h and sha256h2) across abcd and efgh.
216    // 5. Completes message schedule step (sha256su1).
217    macro_rules! sched_round {
218        ($w0:ident, $w1:ident, $w2:ident, $w3:ident, $k:ident, $vtmp:ident) => {
219            concat!(
220                "add ",
221                stringify!($vtmp),
222                ".4s, ",
223                stringify!($w0),
224                ".4s, {",
225                stringify!($k),
226                ":v}.4s\n",
227                "sha256su0 ",
228                stringify!($w0),
229                ".4s, ",
230                stringify!($w1),
231                ".4s\n",
232                "orr {tmp:v}.16b, {abcd:v}.16b, {abcd:v}.16b\n",
233                "sha256h {abcd:q}, {efgh:q}, ",
234                stringify!($vtmp),
235                ".4s\n",
236                "sha256h2 {efgh:q}, {tmp:q}, ",
237                stringify!($vtmp),
238                ".4s\n",
239                "sha256su1 ",
240                stringify!($w0),
241                ".4s, ",
242                stringify!($w2),
243                ".4s, ",
244                stringify!($w3),
245                ".4s\n",
246            )
247        };
248    }
249
250    // Final SHA-256 round without message scheduling:
251    // Rounds 12-15 are the last 4 rounds of a 64-round block. Computing W_{64..67} is
252    // unnecessary as those words are never consumed, saving vector instructions.
253    macro_rules! final_round {
254        ($w:ident, $k:ident, $vtmp:ident) => {
255            concat!(
256                "add ",
257                stringify!($vtmp),
258                ".4s, ",
259                stringify!($w),
260                ".4s, {",
261                stringify!($k),
262                ":v}.4s\n",
263                "orr {tmp:v}.16b, {abcd:v}.16b, {abcd:v}.16b\n",
264                "sha256h {abcd:q}, {efgh:q}, ",
265                stringify!($vtmp),
266                ".4s\n",
267                "sha256h2 {efgh:q}, {tmp:q}, ",
268                stringify!($vtmp),
269                ".4s\n",
270            )
271        };
272    }
273
274    // 4-round quad with cyclic message schedule rotation across v4..v7:
275    macro_rules! sched_quad {
276        ($k0:ident, $k1:ident, $k2:ident, $k3:ident) => {
277            concat!(
278                sched_round!(v4, v5, v6, v7, $k0, v16),
279                sched_round!(v5, v6, v7, v4, $k1, v17),
280                sched_round!(v6, v7, v4, v5, $k2, v16),
281                sched_round!(v7, v4, v5, v6, $k3, v17),
282            )
283        };
284    }
285
286    // Final 4-round quad for rounds 12-15 without message scheduling:
287    macro_rules! final_quad {
288        ($k0:ident, $k1:ident, $k2:ident, $k3:ident) => {
289            concat!(
290                final_round!(v4, $k0, v16),
291                final_round!(v5, $k1, v17),
292                final_round!(v6, $k2, v16),
293                final_round!(v7, $k3, v17),
294            )
295        };
296    }
297
298    // Single zero-block round:
299    // Since all input words W_t are zero, message schedule produces all zeros.
300    // Therefore W_t + K_t = K_t, allowing direct compression with the constant register.
301    macro_rules! zero_round {
302        ($k:ident) => {
303            concat!(
304                "orr {tmp:v}.16b, {abcd:v}.16b, {abcd:v}.16b\n",
305                "sha256h {abcd:q}, {efgh:q}, {",
306                stringify!($k),
307                ":v}.4s\n",
308                "sha256h2 {efgh:q}, {tmp:q}, {",
309                stringify!($k),
310                ":v}.4s\n",
311            )
312        };
313    }
314
315    // 4 zero rounds grouped together:
316    macro_rules! zero_quad {
317        ($k0:ident, $k1:ident, $k2:ident, $k3:ident) => {
318            concat!(zero_round!($k0), zero_round!($k1), zero_round!($k2), zero_round!($k3),)
319        };
320    }
321
322    // Macro to execute the 64 rounds of SHA-256 compression on quadwords in v4..v7.
323    macro_rules! sha256_compress_asm {
324        ($state:ident, $ktable:ident, $load_prefix:expr, $($input_operands:tt)*) => {
325            core::arch::asm!(
326                $load_prefix,
327                "rev32 v4.16b, v4.16b\n",
328                "rev32 v5.16b, v5.16b\n",
329                "rev32 v6.16b, v6.16b\n",
330                "rev32 v7.16b, v7.16b\n",
331
332                sched_quad!(k0, k1, k2, k3),
333                sched_quad!(k4, k5, k6, k7),
334                sched_quad!(k8, k9, k10, k11),
335                final_quad!(k12, k13, k14, k15),
336
337                $($input_operands)*
338                abcd = inout(vreg) $state.abcd,
339                efgh = inout(vreg) $state.efgh,
340                tmp = out(vreg) _,
341                k0 = in(vreg) $ktable.k0,
342                k1 = in(vreg) $ktable.k1,
343                k2 = in(vreg) $ktable.k2,
344                k3 = in(vreg) $ktable.k3,
345                k4 = in(vreg) $ktable.k4,
346                k5 = in(vreg) $ktable.k5,
347                k6 = in(vreg) $ktable.k6,
348                k7 = in(vreg) $ktable.k7,
349                k8 = in(vreg) $ktable.k8,
350                k9 = in(vreg) $ktable.k9,
351                k10 = in(vreg) $ktable.k10,
352                k11 = in(vreg) $ktable.k11,
353                k12 = in(vreg) $ktable.k12,
354                k13 = in(vreg) $ktable.k13,
355                k14 = in(vreg) $ktable.k14,
356                k15 = in(vreg) $ktable.k15,
357                out("v16") _, out("v17") _,
358            )
359        };
360    }
361
362    macro_rules! zero_compress_quads {
363        () => {
364            concat!(
365                zero_quad!(k0, k1, k2, k3),
366                zero_quad!(k4, k5, k6, k7),
367                zero_quad!(k8, k9, k10, k11),
368                zero_quad!(k12, k13, k14, k15),
369            )
370        };
371    }
372
373    impl Sha256VectorState {
374        #[inline(always)]
375        fn new() -> Self {
376            // SAFETY: Reading 16-byte aligned static constant arrays.
377            unsafe {
378                Self {
379                    abcd: vld1q_u32(SHA256_INIT_H_0_3.0.as_ptr()),
380                    efgh: vld1q_u32(SHA256_INIT_H_4_7.0.as_ptr()),
381                }
382            }
383        }
384
385        /// Performs SHA-256 transformation on a 64-byte block loaded from `ptr`.
386        ///
387        /// # Safety
388        ///
389        /// `ptr` must be valid for reading 64 bytes.
390        #[inline(always)]
391        unsafe fn transform_block_from_ptr(&mut self, ktable: &KTable, ptr: *const u8) {
392            let saved_abcd = self.abcd;
393            let saved_efgh = self.efgh;
394
395            // SAFETY: ARMv8-A SHA-256 block transformation with cycle-accurate instruction pairing.
396            unsafe {
397                sha256_compress_asm!(
398                    self,
399                    ktable,
400                    "ld1 {{v4.16b, v5.16b, v6.16b, v7.16b}}, [{ptr}]\n",
401                    ptr = in(reg) ptr,
402                    out("v4") _, out("v5") _, out("v6") _, out("v7") _,
403                );
404                self.abcd = vaddq_u32(self.abcd, saved_abcd);
405                self.efgh = vaddq_u32(self.efgh, saved_efgh);
406            }
407        }
408
409        /// Performs SHA-256 transformation on a block of all zeros without message
410        /// scheduling overhead.
411        #[inline(always)]
412        fn transform_zero_block(&mut self, ktable: &KTable) {
413            let saved_abcd = self.abcd;
414            let saved_efgh = self.efgh;
415
416            // SAFETY: For zero input, all W_t are zero and message schedule W_t produces zeros.
417            // Therefore W_t + K_t = K_t for all 64 rounds, requiring no scheduling operations.
418            unsafe {
419                core::arch::asm!(
420                    zero_compress_quads!(),
421
422                    abcd = inout(vreg) self.abcd,
423                    efgh = inout(vreg) self.efgh,
424                    tmp = out(vreg) _,
425                    k0 = in(vreg) ktable.k0,
426                    k1 = in(vreg) ktable.k1,
427                    k2 = in(vreg) ktable.k2,
428                    k3 = in(vreg) ktable.k3,
429                    k4 = in(vreg) ktable.k4,
430                    k5 = in(vreg) ktable.k5,
431                    k6 = in(vreg) ktable.k6,
432                    k7 = in(vreg) ktable.k7,
433                    k8 = in(vreg) ktable.k8,
434                    k9 = in(vreg) ktable.k9,
435                    k10 = in(vreg) ktable.k10,
436                    k11 = in(vreg) ktable.k11,
437                    k12 = in(vreg) ktable.k12,
438                    k13 = in(vreg) ktable.k13,
439                    k14 = in(vreg) ktable.k14,
440                    k15 = in(vreg) ktable.k15,
441                );
442                self.abcd = vaddq_u32(self.abcd, saved_abcd);
443                self.efgh = vaddq_u32(self.efgh, saved_efgh);
444            }
445        }
446
447        /// Performs SHA-256 transformation directly on vector registers.
448        #[inline(always)]
449        fn transform_words(
450            &mut self,
451            ktable: &KTable,
452            w0: uint8x16_t,
453            w1: uint8x16_t,
454            w2: uint8x16_t,
455            w3: uint8x16_t,
456        ) {
457            let saved_abcd = self.abcd;
458            let saved_efgh = self.efgh;
459
460            // SAFETY: ARMv8-A SHA-256 block transformation with cycle-accurate instruction pairing.
461            unsafe {
462                sha256_compress_asm!(
463                    self,
464                    ktable,
465                    "",
466                    inout("v4") w0 => _,
467                    inout("v5") w1 => _,
468                    inout("v6") w2 => _,
469                    inout("v7") w3 => _,
470                );
471                self.abcd = vaddq_u32(self.abcd, saved_abcd);
472                self.efgh = vaddq_u32(self.efgh, saved_efgh);
473            }
474        }
475
476        /// Converts final 256-bit vector registers to a big-endian `Hash`.
477        #[inline(always)]
478        fn into_hash(self) -> Hash {
479            // SAFETY: Vector register conversion and storing to aligned stack array.
480            unsafe {
481                let abcd_be = vrev32q_u8(vreinterpretq_u8_u32(self.abcd));
482                let efgh_be = vrev32q_u8(vreinterpretq_u8_u32(self.efgh));
483                let mut out = [0u8; 32];
484                vst1q_u8(out.as_mut_ptr(), abcd_be);
485                vst1q_u8(out.as_mut_ptr().add(16), efgh_be);
486                Hash::from_array(out)
487            }
488        }
489    }
490
491    /// ARMv8-A accelerated SHA-256 calculation for page-aligned Merkle blocks (4096 or 8192 bytes).
492    ///
493    /// Note: Any data in `data` beyond `unaligned_len` (up to `buffer_len`) MUST be zeroed
494    /// for the computed hash to match.
495    #[inline(always)]
496    pub fn hash_block_arm64_neon(
497        offset: usize,
498        data: PtrByteSlice<'_>,
499        buffer_len: usize,
500        unaligned_len: usize,
501    ) -> Hash {
502        assert!(buffer_len == BLOCK_SIZE || buffer_len == 4096);
503        assert!(unaligned_len > 0 && unaligned_len <= buffer_len);
504        assert!(offset.is_multiple_of(BLOCK_SIZE));
505        assert!(data.len() >= buffer_len);
506
507        let ktable = KTable::load();
508        let mut state = Sha256VectorState::new();
509
510        // Block 0: 12B identity header + first 52B of data loaded directly into vector registers.
511        let mut w0_bytes = [0u8; 16];
512        w0_bytes[0..12].copy_from_slice(&make_identity(unaligned_len, 0, offset));
513        let src = data.as_ptr();
514        // SAFETY: data has at least buffer_len >= 4096 bytes (asserted above), so offsets 0..52
515        // are valid for reading.
516        unsafe {
517            (w0_bytes.as_mut_ptr().add(12) as *mut u32)
518                .write_unaligned((src as *const u32).read_unaligned());
519            let w0 = vld1q_u8(w0_bytes.as_ptr());
520            let w1 = vld1q_u8(src.add(4));
521            let w2 = vld1q_u8(src.add(20));
522            let w3 = vld1q_u8(src.add(36));
523            state.transform_words(&ktable, w0, w1, w2, w3);
524        }
525
526        /// Loads 12 unaligned bytes from `src`, sets byte 12 to `pad_byte` (with bytes 13..15 as 0),
527        /// and loads the resulting 16-byte buffer into a vector register.
528        ///
529        /// # Safety
530        ///
531        /// `src` must be valid for reading 12 bytes.
532        #[inline(always)]
533        unsafe fn load_12b_padded(src: *const u8, pad_byte: u8) -> uint8x16_t {
534            let mut buf = [0u8; 16];
535            // SAFETY: `src` is valid for reading 12 bytes, and `buf` is a local 16-byte array.
536            unsafe {
537                (buf.as_mut_ptr() as *mut u64)
538                    .write_unaligned((src as *const u64).read_unaligned());
539                (buf.as_mut_ptr().add(8) as *mut u32)
540                    .write_unaligned((src.add(8) as *const u32).read_unaligned());
541                buf[12] = pad_byte;
542                vld1q_u8(buf.as_ptr())
543            }
544        }
545
546        /// Performs the SHA-256 transformation on the tail padding block (Block 128)
547        /// using 8 KiB Merkle length constant (12B header + 8192B payload).
548        #[inline(always)]
549        fn transform_tail_block(state: &mut Sha256VectorState, ktable: &KTable, w0: uint8x16_t) {
550            const TOTAL_BITS_8K: u64 = (12 + 8192) as u64 * 8;
551            let mut w3_bytes = [0u8; 16];
552            // SAFETY: `w3_bytes` is a 16-byte stack array, and `zero_vec` is a valid vector.
553            unsafe {
554                (w3_bytes.as_mut_ptr().add(8) as *mut u64).write_unaligned(TOTAL_BITS_8K.to_be());
555                let zero_vec = vdupq_n_u8(0);
556                let w3 = vld1q_u8(w3_bytes.as_ptr());
557                state.transform_words(ktable, w0, zero_vec, zero_vec, w3);
558            }
559        }
560
561        // Single branch check for 8 KiB (BLOCK_SIZE) vs 4 KiB buffer_len:
562        if buffer_len >= BLOCK_SIZE {
563            // Fast path for 8 KiB page: 127 full 64-byte vector blocks
564            for block_idx in 0..127 {
565                let block_offset = 52 + block_idx * 64;
566                // SAFETY: data has at least 8192 bytes, 52 + 126 * 64 + 64 = 8180 <= 8192.
567                unsafe { state.transform_block_from_ptr(&ktable, data.as_ptr().add(block_offset)) };
568            }
569
570            // Tail block (block 128): remaining 12B of data (offset 8180..8192) + SHA-256 padding
571            // SAFETY: data has at least 8192 bytes, so offset 8180 is valid for reading 12 bytes.
572            let w0 = unsafe { load_12b_padded(data.as_ptr().add(8180), 0x80) };
573            transform_tail_block(&mut state, &ktable, w0);
574        } else {
575            // Fast path for 4 KiB page: 63 full 64-byte vector blocks
576            for block_idx in 0..63 {
577                let block_offset = 52 + block_idx * 64;
578                // SAFETY: data has at least 4096 bytes, 52 + 62 * 64 + 64 = 4084 <= 4096.
579                unsafe { state.transform_block_from_ptr(&ktable, data.as_ptr().add(block_offset)) };
580            }
581
582            // Intermediate block (block 64): remaining 12B (offset 4084..4096) + 52B of zeros
583            // SAFETY: data has at least 4096 bytes, so offset 4084 is valid for reading 12 bytes.
584            unsafe {
585                let w0 = load_12b_padded(data.as_ptr().add(4084), 0x00);
586                let zero_vec = vdupq_n_u8(0);
587                state.transform_words(&ktable, w0, zero_vec, zero_vec, zero_vec);
588            }
589
590            // Remaining 63 blocks (4096 - 52 bytes) of zeros
591            for _ in 0..63 {
592                state.transform_zero_block(&ktable);
593            }
594
595            // Tail block (block 128): 12B zeros + 0x80 padding + 64-bit length
596            let mut w0_bytes = [0u8; 16];
597            w0_bytes[12] = 0x80;
598            // SAFETY: `w0_bytes` is a local 16-byte stack array.
599            let w0 = unsafe { vld1q_u8(w0_bytes.as_ptr()) };
600            transform_tail_block(&mut state, &ktable, w0);
601        }
602
603        state.into_hash()
604    }
605}
606
607/// Updates `hasher` with `count` zeros.
608pub(crate) fn update_with_zeros(hasher: &mut Sha256, mut count: usize) {
609    const BUF_SIZE: usize = 512;
610    const ZEROS: [u8; BUF_SIZE] = [0; BUF_SIZE];
611    while count >= BUF_SIZE {
612        count -= BUF_SIZE;
613        hasher.update(&ZEROS);
614    }
615    if count > 0 {
616        hasher.update(&ZEROS[0..count]);
617    }
618}
619
620/// Creates a new [`Sha256`] for hashing blocks of hashes. The hasher is initialized with the
621/// block's identity.
622pub fn make_hash_hasher(hashes_per_hash: usize, level: usize, offset: usize) -> Sha256 {
623    debug_assert!(level > 0);
624    let bytes_per_hash = hashes_per_hash * HASH_SIZE;
625    debug_assert!(offset.is_multiple_of(bytes_per_hash));
626    let mut hasher = Sha256::new();
627    hasher.update(&make_identity(bytes_per_hash, level, offset));
628    hasher
629}
630
631#[cfg(test)]
632mod tests {
633    use super::*;
634    use storage_ptr_slice::PtrByteSlice;
635
636    fn generate_pattern_data(len: usize) -> Vec<u8> {
637        (0..len).map(|i| ((i * 31 + 17) % 256) as u8).collect()
638    }
639
640    #[test]
641    fn test_hash_block_empty() {
642        let block = [];
643        let hash = hash_block(&block[..], 0);
644        let expected =
645            "15ec7bf0b50732b49f8228e07d24365338f9e3ab994b00af08e5a3bffe55fd8b".parse().unwrap();
646        assert_eq!(hash, expected);
647    }
648
649    #[test]
650    fn test_hash_block_single() {
651        let block = vec![0xFF; 8192];
652        let hash = hash_block(&block[..], 0);
653        let expected =
654            "68d131bc271f9c192d4f6dcd8fe61bef90004856da19d0f2f514a7f4098b0737".parse().unwrap();
655        assert_eq!(hash, expected);
656    }
657
658    #[test]
659    fn test_hash_block_aligned_matches_reference_4k() {
660        let offsets = [0, BLOCK_SIZE, 16384];
661        let unaligned_lens = [1, 12, 51, 52, 53, 100, 500, 4083, 4084, 4085, 4096];
662
663        for &offset in &offsets {
664            for &unaligned_len in &unaligned_lens {
665                let full_data = generate_pattern_data(unaligned_len);
666                let ref_hash = hash_block(&full_data, offset);
667
668                let mut buf = vec![0u8; 4096];
669                buf[..unaligned_len].copy_from_slice(&full_data);
670
671                let hash =
672                    hash_block_aligned(offset, PtrByteSlice::from(&buf[..]), 4096, unaligned_len);
673                assert_eq!(
674                    hash, ref_hash,
675                    "Mismatch: offset={offset}, unaligned_len={unaligned_len}"
676                );
677            }
678        }
679    }
680
681    #[test]
682    fn test_hash_block_aligned_matches_reference_8k() {
683        let offsets = [0, BLOCK_SIZE, 16384];
684        let unaligned_lens = [1, 12, 51, 52, 53, 4096, 8179, 8180, 8181, 8192];
685
686        for &offset in &offsets {
687            for &unaligned_len in &unaligned_lens {
688                let full_data = generate_pattern_data(unaligned_len);
689                let ref_hash = hash_block(&full_data, offset);
690
691                let mut buf = vec![0u8; BLOCK_SIZE];
692                buf[..unaligned_len].copy_from_slice(&full_data);
693
694                let hash = hash_block_aligned(
695                    offset,
696                    PtrByteSlice::from(&buf[..]),
697                    BLOCK_SIZE,
698                    unaligned_len,
699                );
700                assert_eq!(
701                    hash, ref_hash,
702                    "Mismatch: offset={offset}, unaligned_len={unaligned_len}"
703                );
704            }
705        }
706    }
707
708    #[test]
709    fn test_hash_block_aligned_unzeroed_tail_causes_hash_mismatch() {
710        let offset = 0;
711        let unaligned_len = 100;
712        let full_data = generate_pattern_data(unaligned_len);
713        let ref_hash = hash_block(&full_data, offset);
714
715        let mut buf = vec![0xAAu8; 4096];
716        buf[..unaligned_len].copy_from_slice(&full_data);
717
718        let hash = hash_block_aligned(offset, PtrByteSlice::from(&buf[..]), 4096, unaligned_len);
719
720        assert_ne!(hash, ref_hash, "Unzeroed tail must result in mismatched hash");
721    }
722
723    #[cfg(all(target_arch = "aarch64", target_feature = "sha2"))]
724    #[test]
725    fn test_arm64_neon_matches_reference() {
726        let offsets = [0, BLOCK_SIZE, 32768];
727        let buffer_lens = [4096, BLOCK_SIZE];
728
729        for &buffer_len in &buffer_lens {
730            let unaligned_lens = if buffer_len == 4096 {
731                vec![1, 12, 51, 52, 53, 2000, 4083, 4084, 4085, 4096]
732            } else {
733                vec![1, 12, 51, 52, 53, 4096, 8179, 8180, 8181, 8192]
734            };
735
736            for &offset in &offsets {
737                for &unaligned_len in unaligned_lens.iter() {
738                    let full_data = generate_pattern_data(unaligned_len);
739                    let ref_hash = hash_block(&full_data, offset);
740
741                    let mut buf = vec![0u8; buffer_len];
742                    buf[..unaligned_len].copy_from_slice(&full_data);
743
744                    let neon_hash = arm64::hash_block_arm64_neon(
745                        offset,
746                        PtrByteSlice::from(&buf[..]),
747                        buffer_len,
748                        unaligned_len,
749                    );
750                    assert_eq!(
751                        neon_hash, ref_hash,
752                        "NEON mismatch: buf={buffer_len}, off={offset}, len={unaligned_len}",
753                    );
754                }
755            }
756        }
757    }
758}