1use std::ops::Range;
4
5use jf_merkle_tree::MerkleTreeScheme;
6use p3_maybe_rayon::prelude::*;
7use serde::{Deserialize, Serialize};
8
9use super::{AvidmGf2Commit, AvidmGf2Share};
10use crate::{
11 VidError, VidResult, VidScheme,
12 avidm_gf2::{AvidmGf2Scheme, MerkleTree},
13};
14
15pub struct NsAvidmGf2Scheme;
17
18pub type NsAvidmGf2Commit = super::AvidmGf2Commit;
20pub type NsAvidmGf2Param = super::AvidmGf2Param;
22
23#[derive(Clone, Debug, Hash, Serialize, Deserialize, Eq, PartialEq)]
25pub struct NsAvidmGf2Common {
26 pub param: NsAvidmGf2Param,
28 pub ns_commits: Vec<AvidmGf2Commit>,
30 pub ns_lens: Vec<usize>,
32}
33
34impl NsAvidmGf2Common {
35 pub fn payload_byte_len(&self) -> usize {
37 self.ns_lens.iter().sum()
38 }
39}
40
41#[derive(Clone, Debug, Hash, Serialize, Deserialize, Eq, PartialEq, Default)]
43pub struct NsAvidmGf2Share(pub(crate) Vec<AvidmGf2Share>);
44
45impl NsAvidmGf2Share {
46 pub fn num_nss(&self) -> usize {
48 self.0.len()
49 }
50
51 pub fn weight(&self) -> usize {
53 self.0.first().map_or(0, |share| share.weight())
54 }
55
56 pub fn validate(&self) -> bool {
58 let weight = self.weight();
59 self.0
60 .iter()
61 .all(|share| share.validate() && share.weight() == weight)
62 }
63
64 pub fn contains_ns(&self, ns_index: usize) -> bool {
66 ns_index < self.num_nss()
67 }
68
69 pub fn inner_ns_share(&self, ns_index: usize) -> Option<AvidmGf2Share> {
71 self.0.get(ns_index).cloned()
72 }
73}
74
75impl NsAvidmGf2Scheme {
76 pub fn setup(recovery_threshold: usize, total_weights: usize) -> VidResult<NsAvidmGf2Param> {
78 NsAvidmGf2Param::new(recovery_threshold, total_weights)
79 }
80
81 pub fn commit(
85 param: &NsAvidmGf2Param,
86 payload: &[u8],
87 ns_table: impl IntoIterator<Item = Range<usize>>,
88 ) -> VidResult<(NsAvidmGf2Commit, NsAvidmGf2Common)> {
89 let ns_table = ns_table.into_iter().collect::<Vec<_>>();
90 let ns_lens = ns_table.iter().map(|r| r.len()).collect::<Vec<_>>();
91 let ns_commits = ns_table
92 .into_iter()
93 .map(|ns_range| AvidmGf2Scheme::commit(param, &payload[ns_range]))
94 .collect::<Result<Vec<_>, _>>()?;
95 let common = NsAvidmGf2Common {
96 param: param.clone(),
97 ns_commits,
98 ns_lens,
99 };
100 let commit = MerkleTree::from_elems(None, common.ns_commits.iter().map(|c| c.commit))
101 .map_err(|err| VidError::Internal(err.into()))?
102 .commitment();
103 Ok((NsAvidmGf2Commit { commit }, common))
104 }
105
106 pub fn is_consistent(commit: &NsAvidmGf2Commit, common: &NsAvidmGf2Common) -> bool {
108 let Ok(mt) =
109 MerkleTree::from_elems(None, common.ns_commits.iter().map(|commit| commit.commit))
110 else {
111 return false;
112 };
113 commit.commit == mt.commitment()
114 }
115
116 pub fn ns_disperse(
121 param: &NsAvidmGf2Param,
122 distribution: &[u32],
123 payload: &[u8],
124 ns_table: impl IntoIterator<Item = Range<usize>>,
125 ) -> VidResult<(NsAvidmGf2Commit, NsAvidmGf2Common, Vec<NsAvidmGf2Share>)> {
126 let num_storage_nodes = distribution.len();
127 let ns_ranges: Vec<Range<usize>> = ns_table.into_iter().collect();
128 let ns_lens: Vec<usize> = ns_ranges.iter().map(|r| r.len()).collect();
129
130 let per_ns: Vec<(AvidmGf2Commit, Vec<AvidmGf2Share>)> = ns_ranges
135 .par_iter()
136 .map(|ns_range| {
137 AvidmGf2Scheme::disperse(param, distribution, &payload[ns_range.clone()])
138 })
139 .collect::<VidResult<Vec<_>>>()?;
140
141 let (ns_commits, disperses): (Vec<_>, Vec<_>) = per_ns.into_iter().unzip();
142
143 let common = NsAvidmGf2Common {
144 param: param.clone(),
145 ns_commits,
146 ns_lens,
147 };
148 let commit = NsAvidmGf2Commit {
149 commit: MerkleTree::from_elems(None, common.ns_commits.iter().map(|c| c.commit))
150 .map_err(|err| VidError::Internal(err.into()))?
151 .commitment(),
152 };
153 let mut shares = vec![NsAvidmGf2Share::default(); num_storage_nodes];
154 disperses.into_iter().for_each(|ns_disperse| {
155 shares
156 .iter_mut()
157 .zip(ns_disperse)
158 .for_each(|(share, ns_share)| share.0.push(ns_share))
159 });
160 Ok((commit, common, shares))
161 }
162
163 pub fn verify_share_with_verified_common(
170 common: &NsAvidmGf2Common,
171 share: &NsAvidmGf2Share,
172 ) -> VidResult<crate::VerificationResult> {
173 if !(common.ns_commits.len() == common.ns_lens.len()
174 && common.ns_commits.len() == share.num_nss()
175 && share.validate())
176 {
177 return Err(VidError::InvalidShare);
178 }
179 match common
183 .ns_commits
184 .par_iter()
185 .zip(share.0.par_iter())
186 .map(|(commit, content)| AvidmGf2Scheme::verify_share(&common.param, commit, content))
187 .find_any(|r| !matches!(r, Ok(Ok(()))))
188 {
189 None => Ok(Ok(())),
190 Some(Ok(v)) => Ok(v),
191 Some(Err(e)) => Err(e),
192 }
193 }
194
195 pub fn verify_share(
197 commit: &NsAvidmGf2Commit,
198 common: &NsAvidmGf2Common,
199 share: &NsAvidmGf2Share,
200 ) -> VidResult<crate::VerificationResult> {
201 if !Self::is_consistent(commit, common) {
202 return Ok(Err(()));
203 }
204 Self::verify_share_with_verified_common(common, share)
205 }
206
207 pub fn recover(common: &NsAvidmGf2Common, shares: &[NsAvidmGf2Share]) -> VidResult<Vec<u8>> {
209 if shares.is_empty() {
210 return Err(VidError::InsufficientShares);
211 }
212 let per_ns: Vec<Vec<u8>> = (0..common.ns_lens.len())
216 .into_par_iter()
217 .map(|ns_index| Self::ns_recover(common, ns_index, shares))
218 .collect::<VidResult<Vec<_>>>()?;
219 Ok(per_ns.concat())
220 }
221
222 pub fn ns_recover(
226 common: &NsAvidmGf2Common,
227 ns_index: usize,
228 shares: &[NsAvidmGf2Share],
229 ) -> VidResult<Vec<u8>> {
230 if shares.is_empty() {
231 return Err(VidError::InsufficientShares);
232 }
233 if ns_index >= common.ns_lens.len()
234 || !shares.iter().all(|share| share.contains_ns(ns_index))
235 {
236 return Err(VidError::IndexOutOfBound);
237 }
238 let ns_commit = &common.ns_commits[ns_index];
239 let shares: Vec<_> = shares
240 .iter()
241 .filter_map(|share| share.inner_ns_share(ns_index))
242 .collect();
243 AvidmGf2Scheme::recover(&common.param, ns_commit, &shares)
244 }
245}
246
247#[cfg(test)]
249pub mod tests {
250 use rand::{RngCore, seq::SliceRandom};
251
252 use crate::avidm_gf2::namespaced::NsAvidmGf2Scheme;
253
254 fn disperse_with_payload(
255 payload: &[u8],
256 ) -> (
257 crate::avidm_gf2::namespaced::NsAvidmGf2Commit,
258 crate::avidm_gf2::namespaced::NsAvidmGf2Common,
259 Vec<crate::avidm_gf2::namespaced::NsAvidmGf2Share>,
260 ) {
261 let num_storage_nodes = 9;
262 let ns_table = [(0usize..15), (15..48)];
263
264 let mut rng = jf_utils::test_rng();
265 let weights: Vec<u32> = (0..num_storage_nodes)
266 .map(|_| rng.next_u32() % 5 + 1)
267 .collect();
268 let total_weights: u32 = weights.iter().sum();
269 let recovery_threshold = total_weights.div_ceil(3) as usize;
270 let params = NsAvidmGf2Scheme::setup(recovery_threshold, total_weights as usize).unwrap();
271
272 NsAvidmGf2Scheme::ns_disperse(¶ms, &weights, payload, ns_table.iter().cloned()).unwrap()
273 }
274
275 fn setup_test_data() -> (
276 crate::avidm_gf2::namespaced::NsAvidmGf2Commit,
277 crate::avidm_gf2::namespaced::NsAvidmGf2Common,
278 Vec<crate::avidm_gf2::namespaced::NsAvidmGf2Share>,
279 ) {
280 let payload: Vec<u8> = (0u8..48).collect();
281 disperse_with_payload(&payload)
282 }
283
284 #[test]
285 fn verify_share_with_verified_common_accepts_valid() {
286 let (commit, common, shares) = setup_test_data();
287 assert!(NsAvidmGf2Scheme::is_consistent(&commit, &common));
288 for share in &shares {
289 assert!(
290 NsAvidmGf2Scheme::verify_share_with_verified_common(&common, share)
291 .is_ok_and(|r| r.is_ok())
292 );
293 }
294 }
295
296 #[test]
297 fn verify_share_with_verified_common_rejects_tampered_share() {
298 let (_commit, common, shares) = setup_test_data();
299 let mut tampered = shares[0].clone();
301 tampered.0.pop();
302 assert!(NsAvidmGf2Scheme::verify_share_with_verified_common(&common, &tampered).is_err());
303
304 let (_commit2, _common2, shares2) = disperse_with_payload(&[0xAB; 48]);
306 let mut mixed = shares[0].clone();
307 mixed.0[0] = shares2[0].0[0].clone();
308 assert!(
309 NsAvidmGf2Scheme::verify_share_with_verified_common(&common, &mixed)
310 .is_ok_and(|r| r.is_err())
311 );
312 }
313
314 #[test]
315 fn composition_equivalence() {
316 let (commit, common, shares) = setup_test_data();
317 for share in &shares {
318 let full_result = NsAvidmGf2Scheme::verify_share(&commit, &common, share)
319 .unwrap()
320 .is_ok();
321 let composed_result = NsAvidmGf2Scheme::is_consistent(&commit, &common)
322 && NsAvidmGf2Scheme::verify_share_with_verified_common(&common, share)
323 .unwrap()
324 .is_ok();
325 assert_eq!(full_result, composed_result);
326 }
327 }
328
329 #[test]
330 fn is_consistent_rejects_tampered_commit() {
331 let (commit, common, _shares) = setup_test_data();
332 let (different_commit, ..) = disperse_with_payload(&[0xCD; 48]);
334 assert!(NsAvidmGf2Scheme::is_consistent(&commit, &common));
336 assert!(!NsAvidmGf2Scheme::is_consistent(&different_commit, &common));
338 }
339
340 #[test]
341 fn is_consistent_rejects_tampered_common() {
342 let (commit, common, _shares) = setup_test_data();
343 let (_, different_common, _) = disperse_with_payload(&[0xCD; 48]);
345 let mut tampered_common = common;
346 tampered_common.ns_commits = different_common.ns_commits;
347 assert!(!NsAvidmGf2Scheme::is_consistent(&commit, &tampered_common));
348 }
349
350 #[test]
351 fn round_trip() {
352 let num_storage_nodes = 9;
354 let ns_lens = [15, 33];
355 let ns_table = [(0usize..15), (15..48)];
356 let payload_byte_len = ns_lens.iter().sum();
357
358 let mut rng = jf_utils::test_rng();
359
360 let weights: Vec<u32> = (0..num_storage_nodes)
362 .map(|_| rng.next_u32() % 5 + 1)
363 .collect();
364 let total_weights: u32 = weights.iter().sum();
365 let recovery_threshold = total_weights.div_ceil(3) as usize;
366 let params = NsAvidmGf2Scheme::setup(recovery_threshold, total_weights as usize).unwrap();
367
368 println!(
369 "recovery_threshold:: {recovery_threshold} num_storage_nodes: {num_storage_nodes} \
370 payload_byte_len: {payload_byte_len}"
371 );
372 println!("weights: {weights:?}");
373
374 let payload = {
375 let mut bytes_random = vec![0u8; payload_byte_len];
376 rng.fill_bytes(&mut bytes_random);
377 bytes_random
378 };
379
380 let (commit, common, mut shares) =
381 NsAvidmGf2Scheme::ns_disperse(¶ms, &weights, &payload, ns_table.iter().cloned())
382 .unwrap();
383
384 assert_eq!(shares.len(), num_storage_nodes);
385
386 assert_eq!(
387 commit,
388 NsAvidmGf2Scheme::commit(¶ms, &payload, ns_table.iter().cloned())
389 .unwrap()
390 .0
391 );
392
393 shares.iter().for_each(|share| {
395 assert!(
396 NsAvidmGf2Scheme::verify_share(&commit, &common, share).is_ok_and(|r| r.is_ok())
397 )
398 });
399
400 shares.shuffle(&mut rng);
402 let mut cumulated_weights = 0;
403 let mut cut_index = 0;
404 while cumulated_weights <= recovery_threshold {
405 cumulated_weights += shares[cut_index].weight();
406 cut_index += 1;
407 }
408 let ns0_payload_recovered =
409 NsAvidmGf2Scheme::ns_recover(&common, 0, &shares[..cut_index]).unwrap();
410 assert_eq!(ns0_payload_recovered[..], payload[ns_table[0].clone()]);
411 let ns1_payload_recovered =
412 NsAvidmGf2Scheme::ns_recover(&common, 1, &shares[..cut_index]).unwrap();
413 assert_eq!(ns1_payload_recovered[..], payload[ns_table[1].clone()]);
414 let payload_recovered = NsAvidmGf2Scheme::recover(&common, &shares[..cut_index]).unwrap();
415 assert_eq!(payload_recovered, payload);
416 }
417}