risc0_zkp/core/hash/sha/
rng.rs

1// Copyright 2025 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//! A SHA-256 based CRNG used in Fiat-Shamir.
16
17use alloc::boxed::Box;
18
19use rand_core::{impls, RngCore};
20use risc0_core::field::{Elem, Field};
21
22use super::{Digest, Impl, Sha256, DIGEST_WORDS};
23use crate::core::hash::Rng;
24
25/// A random number generator driven by a [Sha256].
26#[derive(Clone, Debug)]
27pub struct ShaRng {
28    // Pool 0 receives new entropy and is where values are drawn from.
29    pool0: Box<Digest>,
30    // Pool 1 provides secret state in the step function. It is never observable.
31    pool1: Box<Digest>,
32    pool_used: usize,
33}
34
35impl Default for ShaRng {
36    fn default() -> Self {
37        Self::new()
38    }
39}
40
41impl ShaRng {
42    /// Create a new [ShaRng] from a given [Sha256].
43    pub fn new() -> Self {
44        Self {
45            pool0: (*Impl::hash_bytes(b"Hello")).into(),
46            pool1: (*Impl::hash_bytes(b"World")).into(),
47            pool_used: 0,
48        }
49    }
50
51    /// Mix the pool with a specified [Digest].
52    pub fn inner_mix(&mut self, val: &Digest) {
53        for i in 0..DIGEST_WORDS {
54            self.pool0.as_mut_words()[i] = self.pool0.as_words()[i] ^ val.as_words()[i];
55        }
56        self.step();
57    }
58
59    fn step(&mut self) {
60        self.pool0 = (*Impl::hash_pair(&self.pool0, &self.pool1)).into();
61        self.pool1 = (*Impl::hash_pair(&self.pool0, &self.pool1)).into();
62        self.pool_used = 0;
63    }
64}
65
66impl RngCore for ShaRng {
67    fn next_u32(&mut self) -> u32 {
68        if self.pool_used == DIGEST_WORDS {
69            self.step();
70        }
71        let out = self.pool0.as_words()[self.pool_used];
72        // Mark this word as used.
73        self.pool_used += 1;
74        out
75    }
76
77    fn next_u64(&mut self) -> u64 {
78        ((self.next_u32() as u64) << 32) | (self.next_u32() as u64)
79    }
80
81    fn fill_bytes(&mut self, dest: &mut [u8]) {
82        impls::fill_bytes_via_next(self, dest);
83    }
84}
85
86impl<F: Field> Rng<F> for ShaRng {
87    fn mix(&mut self, val: &Digest) {
88        self.inner_mix(val);
89    }
90
91    fn random_bits(&mut self, bits: usize) -> u32 {
92        ((1 << bits) - 1) & self.next_u32()
93    }
94
95    fn random_elem(&mut self) -> F::Elem {
96        F::Elem::random(self)
97    }
98
99    fn random_ext_elem(&mut self) -> F::ExtElem {
100        F::ExtElem::random(self)
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use rand_core::RngCore;
107
108    use super::ShaRng;
109    use crate::core::hash::sha::{cpu::Impl, Sha256};
110
111    // Runs conformance test on a SHA implementation to make sure it
112    // properly behaves for generating pseudo-random numbers.
113    #[test]
114    fn test_sha_rng_impl() {
115        let mut x = ShaRng::new();
116        for _ in 0..10 {
117            x.next_u32();
118        }
119        assert_eq!(x.next_u32(), 785921476);
120        x.inner_mix(&Impl::hash_bytes(b"foo"));
121        assert_eq!(x.next_u32(), 4167871101);
122    }
123}