labrador/transcript/sponges/shake/
mod.rs

1use sha3::{
2    digest::{ExtendableOutput, Update, XofReader},
3    Shake256,
4};
5
6use crate::ring::{rq::Rq, zq::Zq};
7use crate::transcript::Sponge;
8
9#[derive(Default)]
10pub struct ShakeSponge {
11    hasher: Shake256,
12}
13
14impl Sponge for ShakeSponge {
15    fn absorb_zq(&mut self, input: &[Zq]) {
16        // Convert Rq ector to u8
17        let mut u8_version_input: Vec<u8> = Vec::new();
18        for coeff in input {
19            u8_version_input.extend_from_slice(&coeff.get_value().to_be_bytes());
20        }
21        self.hasher.update(&u8_version_input);
22    }
23
24    fn absorb_rq(&mut self, input: &[Rq]) {
25        // Convert Rq ector to u8
26        let mut u8_version_input: Vec<u8> = Vec::new();
27        for rq in input {
28            for coeff in rq.get_coefficients() {
29                u8_version_input.extend_from_slice(&coeff.get_value().to_be_bytes());
30            }
31        }
32        self.hasher.update(&u8_version_input);
33    }
34
35    fn squeeze_bits(&mut self, bit_length: usize) -> Vec<bool> {
36        let byte_len = bit_length.div_ceil(8);
37        let mut reader = self.hasher.clone().finalize_xof();
38        let mut output_buffer = vec![u8::default(); byte_len];
39        reader.read(&mut output_buffer);
40
41        let mut result = Vec::with_capacity(bit_length);
42        for byte in &output_buffer {
43            let mut mask = 1u8;
44            for _ in 0..8 {
45                if result.len() == bit_length {
46                    break;
47                }
48                result.push(byte & mask != 0);
49                mask <<= 1;
50            }
51        }
52        self.hasher.update(&output_buffer);
53        result
54    }
55
56    fn squeeze_zq(&mut self, output_length: usize) -> Vec<Zq> {
57        let mut reader = self.hasher.clone().finalize_xof();
58        let mut output_buffer = vec![u8::default(); output_length * 4];
59        reader.read(&mut output_buffer);
60
61        let zq_values: Vec<Zq> = output_buffer
62            .chunks_exact(4)
63            .map(|chunk| {
64                u32::from_le_bytes(chunk.try_into().expect("Could not convert 4 u8 to one u32"))
65            })
66            .map(Zq::new)
67            .collect();
68
69        self.absorb_zq(&zq_values);
70        zq_values
71    }
72
73    fn squeeze_rq(&mut self, output_length: usize) -> Vec<Rq> {
74        let mut reader = self.hasher.clone().finalize_xof();
75        let mut output_buffer = vec![u8::default(); output_length * Rq::DEGREE * 4];
76        reader.read(&mut output_buffer);
77
78        let zq_values: Vec<Zq> = output_buffer
79            .chunks_exact(4)
80            .map(|chunk| {
81                u32::from_le_bytes(chunk.try_into().expect("Could not convert 4 u8 to one u32"))
82            })
83            .map(Zq::new)
84            .collect();
85
86        let result: Vec<Rq> = zq_values
87            .chunks_exact(Rq::DEGREE)
88            .map(|chunk| {
89                let rq_input: [Zq; Rq::DEGREE] =
90                    chunk.try_into().expect("Chunk size is Rq::DEGREE");
91                Rq::new(rq_input)
92            })
93            .collect();
94
95        self.absorb_rq(&result);
96        result
97    }
98
99    fn squeeze_bytes(&mut self, byte_length: usize) -> Vec<u8> {
100        let mut reader = self.hasher.clone().finalize_xof();
101        let mut output_buffer = vec![u8::default(); byte_length];
102        reader.read(&mut output_buffer);
103        self.hasher.update(&output_buffer);
104        output_buffer
105    }
106}
107
108#[cfg(test)]
109mod test_sponge_correctness {
110    use super::*;
111    use crate::ring::zq::UniformZq;
112    use rand::{distr::uniform::UniformSampler, rng};
113
114    fn random_zq_vector(n: usize) -> Vec<Zq> {
115        let uniform = UniformZq::new_inclusive(Zq::ZERO, Zq::NEG_ONE).unwrap();
116        (0..n).map(|_| uniform.sample(&mut rng())).collect()
117    }
118
119    #[test]
120    fn test_zq_sponge_execution() {
121        let mut sponge = ShakeSponge::default();
122        let polynomial_1 = Rq::new([Zq::new(2); Rq::DEGREE]);
123        let polynomial_2 = Rq::new([Zq::new(5); Rq::DEGREE]);
124        let polynomial_3 = Rq::new([Zq::new(1); Rq::DEGREE]);
125
126        sponge.absorb_rq(&[polynomial_1, polynomial_2, polynomial_3]);
127        let result = sponge.squeeze_rq(1);
128        assert_eq!(result.len(), 1);
129    }
130
131    #[test]
132    fn test_rq_sponge_execution() {
133        let mut sponge = ShakeSponge::default();
134        let input: Vec<Zq> = random_zq_vector(64);
135        sponge.absorb_zq(&input);
136        let result = sponge.squeeze_rq(1);
137        assert_eq!(result.len(), 1);
138    }
139
140    #[test]
141    fn test_zq_squeeze_output_size() {
142        let mut sponge = ShakeSponge::default();
143        let input: Vec<Zq> = random_zq_vector(64);
144        sponge.absorb_zq(&input);
145        let result = sponge.squeeze_zq(1000);
146        assert_eq!(result.len(), 1000);
147    }
148
149    #[test]
150    fn test_rq_squeeze_output_size() {
151        let mut sponge = ShakeSponge::default();
152        let polynomial_1 = Rq::new([Zq::new(2); Rq::DEGREE]);
153        let polynomial_2 = Rq::new([Zq::new(5); Rq::DEGREE]);
154        let polynomial_3 = Rq::new([Zq::new(1); Rq::DEGREE]);
155
156        sponge.absorb_rq(&[polynomial_1, polynomial_2, polynomial_3]);
157        let result = sponge.squeeze_rq(1000);
158        assert_eq!(result.len(), 1000);
159    }
160
161    #[test]
162    fn test_absorb_zq_is_deterministic() {
163        let input: Vec<Zq> = random_zq_vector(64);
164        let mut s1 = ShakeSponge::default();
165        let mut s2 = ShakeSponge::default();
166        s1.absorb_zq(&input);
167        s2.absorb_zq(&input);
168        let out1 = s1.squeeze_zq(8);
169        let out2 = s2.squeeze_zq(8);
170        assert_eq!(out1, out2);
171    }
172
173    #[test]
174    fn test_absorb_rq_is_deterministic() {
175        let polynomial_1 = Rq::new([Zq::new(54821); Rq::DEGREE]);
176        let polynomial_2 = Rq::new([Zq::new(2131213); Rq::DEGREE]);
177        let polynomial_3 = Rq::new([Zq::new(9891741); Rq::DEGREE]);
178
179        let mut sponge1 = ShakeSponge::default();
180        sponge1.absorb_rq(&[
181            polynomial_1.clone(),
182            polynomial_2.clone(),
183            polynomial_3.clone(),
184        ]);
185
186        let mut sponge2 = ShakeSponge::default();
187        sponge2.absorb_rq(&[polynomial_1, polynomial_2, polynomial_3]);
188        assert_eq!(sponge1.squeeze_rq(1), sponge2.squeeze_rq(1))
189    }
190
191    #[test]
192    fn test_zq_successive_squeezes_is_unique() {
193        let mut s = ShakeSponge::default();
194        let input: Vec<Zq> = random_zq_vector(32);
195        s.absorb_zq(&input);
196        let o1 = s.squeeze_zq(8);
197        let o2 = s.squeeze_zq(8);
198        assert_ne!(o1, o2);
199    }
200
201    #[test]
202    fn test_rq_successive_squeezes_is_unique() {
203        let polynomial_1 = Rq::new([Zq::new(54821); Rq::DEGREE]);
204        let polynomial_2 = Rq::new([Zq::new(2131213); Rq::DEGREE]);
205
206        let mut sponge = ShakeSponge::default();
207        sponge.absorb_rq(&[polynomial_1, polynomial_2]);
208        let result1 = sponge.squeeze_rq(1);
209        let result2 = sponge.squeeze_rq(1);
210        assert_ne!(result1, result2)
211    }
212
213    #[test]
214    fn test_zq_output_can_be_large() {
215        let mut s = ShakeSponge::default();
216        let input: Vec<Zq> = random_zq_vector(16);
217        s.absorb_zq(&input);
218        let out = s.squeeze_zq(1000);
219        assert_eq!(out.len(), 1000);
220    }
221
222    #[test]
223    fn test_rq_output_can_be_large() {
224        let polynomial_1 = Rq::new([Zq::new(54821); Rq::DEGREE]);
225        let polynomial_2 = Rq::new([Zq::new(2131213); Rq::DEGREE]);
226
227        let mut sponge = ShakeSponge::default();
228        sponge.absorb_rq(&[polynomial_1, polynomial_2]);
229
230        assert_eq!(sponge.squeeze_rq(1000).len(), 1000);
231    }
232
233    #[test]
234    fn test_squeeze_message_order() {
235        let polynomial_1 = Rq::new([Zq::new(54821); Rq::DEGREE]);
236        let polynomial_2 = Rq::new([Zq::new(2131213); Rq::DEGREE]);
237
238        let mut sponge1 = ShakeSponge::default();
239        sponge1.absorb_rq(&[polynomial_1.clone(), polynomial_2.clone()]);
240
241        let mut sponge2 = ShakeSponge::default();
242        sponge2.absorb_rq(&[polynomial_2, polynomial_1]);
243        assert_ne!(sponge1.squeeze_rq(1), sponge2.squeeze_rq(1))
244    }
245
246    #[test]
247    fn test_different_inputs_diff_outputs() {
248        let input1: Vec<Zq> = random_zq_vector(64);
249        let input2: Vec<Zq> = random_zq_vector(64);
250        let mut s1 = ShakeSponge::default();
251        let mut s2 = ShakeSponge::default();
252        s1.absorb_zq(&input1);
253        s2.absorb_zq(&input2);
254        let o1 = s1.squeeze_zq(8);
255        let o2 = s2.squeeze_zq(8);
256        assert_ne!(o1, o2);
257    }
258
259    /// Edge case: zero‑length squeeze should not panic and must return empty vecs.
260    #[test]
261    fn zero_length_squeeze() {
262        let mut s = ShakeSponge::default();
263        assert!(s.squeeze_zq(0).is_empty());
264        assert!(s.squeeze_rq(0).is_empty());
265    }
266
267    #[test]
268    fn squeeze_bits_deterministic_and_length() {
269        let mut s1 = ShakeSponge::default();
270        let mut s2 = ShakeSponge::default();
271        s1.absorb_zq(&[Zq::new(41)]);
272        s2.absorb_zq(&[Zq::new(41)]);
273
274        let b1 = s1.squeeze_bits(123);
275        let b2 = s2.squeeze_bits(123);
276        assert_eq!(b1, b2); // deterministic
277        assert_eq!(b1.len(), 123); // exact length
278    }
279}
280
281#[cfg(test)]
282mod test_sponge_randomness {
283    // We should test squeeze outputs are random-looking
284}