risc0_zkp/core/hash/sha/
cpu.rs

1// Copyright 2024 RISC Zero, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Simple wrappers for a CPU-based SHA-256 implementation.
16
17use alloc::{boxed::Box, vec::Vec};
18use core::slice;
19
20use sha2::{
21    digest::generic_array::{typenum::U64, GenericArray},
22    Digest as _,
23};
24
25use super::{Block, Sha256, SHA256_INIT};
26use crate::core::digest::{Digest, DIGEST_WORDS};
27
28/// A CPU-based [Sha256] implementation.
29#[derive(Default, Clone)]
30pub struct Impl {}
31
32fn set_word(buf: &mut [u8], idx: usize, word: u32) {
33    buf[(4 * idx)..(4 * idx + 4)].copy_from_slice(&word.to_ne_bytes());
34}
35
36impl Sha256 for Impl {
37    type DigestPtr = Box<Digest>;
38
39    fn hash_bytes(bytes: &[u8]) -> Self::DigestPtr {
40        let digest = sha2::Sha256::digest(bytes);
41        let words: Vec<u32> = digest
42            .as_slice()
43            .chunks(4)
44            .map(|chunk| u32::from_ne_bytes(chunk.try_into().unwrap()))
45            .collect();
46        Box::new(Digest::from(
47            <[u32; DIGEST_WORDS]>::try_from(words).unwrap(),
48        ))
49    }
50
51    fn hash_words(words: &[u32]) -> Self::DigestPtr {
52        Self::hash_bytes(bytemuck::cast_slice(words))
53    }
54
55    #[inline]
56    fn hash_raw_data_slice<T: bytemuck::NoUninit>(data: &[T]) -> Self::DigestPtr {
57        let u8s: &[u8] = bytemuck::cast_slice(data);
58        let mut state: [u32; DIGEST_WORDS] = SHA256_INIT.into();
59        for word in state.iter_mut() {
60            *word = word.to_be();
61        }
62        let mut blocks = u8s.chunks_exact(64);
63        for block in blocks.by_ref() {
64            sha2::compress256(&mut state, slice::from_ref(GenericArray::from_slice(block)));
65        }
66        let remainder = blocks.remainder();
67        if !remainder.is_empty() {
68            let mut last_block: GenericArray<u8, U64> = GenericArray::default();
69            bytemuck::cast_slice_mut(last_block.as_mut_slice())[..remainder.len()]
70                .clone_from_slice(remainder);
71            sha2::compress256(&mut state, slice::from_ref(&last_block));
72        }
73        for word in state.iter_mut() {
74            *word = word.to_be();
75        }
76        Box::new(Digest::from(state))
77    }
78
79    // Digest two digest into one
80    #[inline]
81    fn compress(
82        orig_state: &Digest,
83        block_half1: &Digest,
84        block_half2: &Digest,
85    ) -> Self::DigestPtr {
86        // Convert the state from big-endian to native byte order.
87        let mut state: [u32; DIGEST_WORDS] = *orig_state.as_ref();
88        for word in state.iter_mut() {
89            *word = word.to_be();
90        }
91
92        // Half-blocks may not be contiguous so they must be copied here.
93        let mut block: GenericArray<u8, U64> = GenericArray::default();
94        for i in 0..8 {
95            set_word(block.as_mut_slice(), i, block_half1.as_words()[i]);
96            set_word(block.as_mut_slice(), 8 + i, block_half2.as_words()[i]);
97        }
98        sha2::compress256(&mut state, slice::from_ref(&block));
99
100        // Convert the state from big-endian to native byte order.
101        for word in state.iter_mut() {
102            *word = word.to_be();
103        }
104        Box::new(Digest::from(state))
105    }
106
107    #[inline]
108    fn compress_slice(orig_state: &Digest, blocks: &[Block]) -> Self::DigestPtr {
109        // Convert the state from big-endian to native byte order.
110        let mut state: [u32; DIGEST_WORDS] = *orig_state.as_ref();
111        for word in state.iter_mut() {
112            *word = word.to_be();
113        }
114
115        // Reinterpret the RISC Zero blocks as GenericArray<u8, U64>.
116        // SAFETY: We know that the two types have the same memory layout, so this
117        // conversion is known to be safe.
118        match unsafe { blocks.align_to::<GenericArray<u8, U64>>() } {
119            (&[], aligned_blocks, &[]) => sha2::compress256(&mut state, aligned_blocks),
120            _ => unreachable!("alignment will always be satisfied for block conversion"),
121        };
122
123        // Convert the native byte order result to big-endian.
124        for word in state.iter_mut() {
125            *word = word.to_be();
126        }
127        Box::new(Digest::from(state))
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::Impl;
134
135    #[test]
136    fn test_impl() {
137        crate::core::hash::sha::testutil::test_sha_impl::<Impl>();
138    }
139}
risc0_zkp/core/hash/sha/cpu.rs

risc0_zkp/core/hash/sha/
cpu.rs
