Skip to main content

vid/avidm_gf2/
namespaced.rs

1//! This file implements the namespaced AvidmGf2 scheme.
2
3use 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
15/// Dummy struct for namespaced AvidmGf2 scheme
16pub struct NsAvidmGf2Scheme;
17
18/// Namespaced commitment type
19pub type NsAvidmGf2Commit = super::AvidmGf2Commit;
20/// Namespaced parameter type
21pub type NsAvidmGf2Param = super::AvidmGf2Param;
22
23/// VID Common data that needs to be broadcasted to all storage nodes
24#[derive(Clone, Debug, Hash, Serialize, Deserialize, Eq, PartialEq)]
25pub struct NsAvidmGf2Common {
26    /// The AvidmGf2 parameters
27    pub param: NsAvidmGf2Param,
28    /// The list of all namespace commitments
29    pub ns_commits: Vec<AvidmGf2Commit>,
30    /// The size of each namespace
31    pub ns_lens: Vec<usize>,
32}
33
34impl NsAvidmGf2Common {
35    /// Return the total payload byte length
36    pub fn payload_byte_len(&self) -> usize {
37        self.ns_lens.iter().sum()
38    }
39}
40
41/// Namespaced share for each storage node, contains one [`AvidmGf2Share`] for each namespace.
42#[derive(Clone, Debug, Hash, Serialize, Deserialize, Eq, PartialEq, Default)]
43pub struct NsAvidmGf2Share(pub(crate) Vec<AvidmGf2Share>);
44
45impl NsAvidmGf2Share {
46    /// Return the number of namespaces in this share
47    pub fn num_nss(&self) -> usize {
48        self.0.len()
49    }
50
51    /// Return the weight of this share
52    pub fn weight(&self) -> usize {
53        self.0.first().map_or(0, |share| share.weight())
54    }
55
56    /// Validate the share structure
57    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    /// Check whether this share contains a given namespace
65    pub fn contains_ns(&self, ns_index: usize) -> bool {
66        ns_index < self.num_nss()
67    }
68
69    /// Return the inner share for a given namespace if there exists one.
70    pub fn inner_ns_share(&self, ns_index: usize) -> Option<AvidmGf2Share> {
71        self.0.get(ns_index).cloned()
72    }
73}
74
75impl NsAvidmGf2Scheme {
76    /// Setup an instance for AVID-M scheme
77    pub fn setup(recovery_threshold: usize, total_weights: usize) -> VidResult<NsAvidmGf2Param> {
78        NsAvidmGf2Param::new(recovery_threshold, total_weights)
79    }
80
81    /// Commit to a payload given namespace table.
82    /// WARN: it assumes that the namespace table is well formed, i.e. ranges
83    /// are non-overlapping and cover the whole payload.
84    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    /// Check whether the namespaced commitment is consistent with the common data
107    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    /// Disperse a payload according to a distribution table and a namespace
117    /// table.
118    /// WARN: it assumes that the namespace table is well formed, i.e. ranges
119    /// are non-overlapping and cover the whole payload.
120    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        // Per-namespace dispersals are independent; run them on rayon.
131        // Inner `AvidmGf2Scheme::disperse` also uses `par_iter` internally for
132        // leaf hashing and share assembly — rayon's work-stealing pool
133        // coordinates the nested parallelism.
134        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    /// Verify a namespaced share given already-verified common data.
164    ///
165    /// # Safety Contract
166    /// Caller MUST ensure `is_consistent(commit, common)` returned `true`
167    /// before calling this. Without that check, a malicious common could
168    /// make an invalid share appear valid.
169    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        // Per-namespace verifications are independent. `find_any` short-
180        // circuits on the first failing namespace and avoids allocating an
181        // intermediate `Vec<VerificationResult>`.
182        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    /// Verify a namespaced share
196    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    /// Recover the entire payload from enough share
208    pub fn recover(common: &NsAvidmGf2Common, shares: &[NsAvidmGf2Share]) -> VidResult<Vec<u8>> {
209        if shares.is_empty() {
210            return Err(VidError::InsufficientShares);
211        }
212        // Each `ns_recover` is independent: distinct decoder state, distinct
213        // output bytes. Run them on the rayon pool so multi-namespace blocks
214        // actually use more than one core.
215        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    /// Recover the payload for a given namespace.
223    /// Given namespace ID should be valid for all shares, i.e. `ns_commits` and `content` have
224    /// at least `ns_index` elements for all shares.
225    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/// Unit tests
248#[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(&params, &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        // Create a tampered share by removing one namespace entry
300        let mut tampered = shares[0].clone();
301        tampered.0.pop();
302        assert!(NsAvidmGf2Scheme::verify_share_with_verified_common(&common, &tampered).is_err());
303
304        // Create a tampered share by dispersing a different payload and swapping
305        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        // Use commit from a different dispersal
333        let (different_commit, ..) = disperse_with_payload(&[0xCD; 48]);
334        // Verify original is consistent
335        assert!(NsAvidmGf2Scheme::is_consistent(&commit, &common));
336        // Verify different commit is inconsistent with original common
337        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        // Swap in ns_commits from a different dispersal
344        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        // play with these items
353        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        // more items as a function of the above
361        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(&params, &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(&params, &payload, ns_table.iter().cloned())
389                .unwrap()
390                .0
391        );
392
393        // verify shares
394        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        // test payload recovery on a random subset of shares
401        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}