Implement Coding Abstraction (#1657)

This adds an abstraction over various coding schemes. There's a concrete implementation for both the Reed-Solomon case, and for the trivial case where no coding is performed at all.

Tests and benchmarks generic over the scheme are provided.

Author: cronokirby

Cargo.lock

@@ -62,6 +62,7 @@ bytes = { version = "1.7.1", default-features = false }
 sha2 = { version = "0.10.8", default-features = false }
 blake3 = { version = "1.8.2", default-features = false }
 rand = { version = "0.8.5", default-features = false }
+rand_chacha = { version = "0.3", default-features = false }
 rand_core = "0.6.4"
 rand_distr = "0.4.3"
 futures = "0.3.31"

Cargo.toml

@@ -17,15 +17,22 @@ commonware-storage = { workspace = true, features = ["std"] }
 bytes = { workspace = true }
 thiserror = { workspace = true }
 reed-solomon-simd = "3.0.1"
+rand_core = { workspace = true }
 
 [lib]
 bench = false
 
 [dev-dependencies]
 criterion = { workspace = true }
 rand = { workspace = true }
+rand_chacha = { workspace = true }
 
 [[bench]]
-name = "reed_solomon"
+name = "coding_scheme_times"
 harness = false
-path = "src/reed_solomon/benches/bench.rs"
+path = "src/benches/bench.rs"
+
+[[bench]]
+name = "coding_scheme_sizes"
+harness = false
+path = "src/benches/bench_size.rs"

coding/Cargo.toml

@@ -12,13 +12,24 @@ cargo-fuzz = true
 commonware-cryptography = { workspace = true }
 libfuzzer-sys = { workspace = true }
 arbitrary = { workspace = true, features = ["derive"] }
+rand_chacha = { workspace = true }
 
 [dependencies.commonware-coding]
 path = ".."
 
+[lib]
+bench = false
+
 [[bin]]
 name = "reed_solomon"
 path = "fuzz_targets/reed_solomon.rs"
 test = false
 doc = false
 bench = false
+
+[[bin]]
+name = "no_coding"
+path = "fuzz_targets/no_coding.rs"
+test = false
+doc = false
+bench = false

coding/fuzz/Cargo.toml

@@ -0,0 +1,10 @@
+#![no_main]
+
+use commonware_coding::NoCoding;
+use commonware_coding_fuzz::{fuzz, FuzzInput};
+use commonware_cryptography::Sha256;
+use libfuzzer_sys::fuzz_target;
+
+fuzz_target!(|input: FuzzInput| {
+    fuzz::<NoCoding<Sha256>>(input);
+});

coding/fuzz/fuzz_targets/no_coding.rs

@@ -1,92 +1,10 @@
 #![no_main]
 
-use arbitrary::{Arbitrary, Unstructured};
-use commonware_coding::reed_solomon::{decode, encode, Chunk};
+use commonware_coding::ReedSolomon;
+use commonware_coding_fuzz::{fuzz, FuzzInput};
 use commonware_cryptography::Sha256;
 use libfuzzer_sys::fuzz_target;
 
-#[derive(Debug)]
-struct FuzzInput {
-    total: u16,
-    min: u16,
-    data: Vec<u8>,
-    shuffle_bytes: Vec<u8>,
-}
-
-impl<'a> Arbitrary<'a> for FuzzInput {
-    fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
-        let min = u.int_in_range(1..=u16::MAX - 1)?; // min > 0
-        let total = u.int_in_range(min + 1..=u16::MAX)?; // total > min
-        let data_len = u.int_in_range(0..=u32::MAX)?; // data.len() <= u32:Max
-        let data = u.bytes(data_len as usize)?.to_vec();
-        let shuffle_bytes = u.bytes(8)?.to_vec();
-
-        Ok(FuzzInput {
-            total,
-            min,
-            data,
-            shuffle_bytes,
-        })
-    }
-}
-
-#[derive(Clone)]
-pub struct ShuffledChunks {
-    pub chunks: Vec<Chunk<Sha256>>,
-}
-
-impl ShuffledChunks {
-    pub fn from_chunks<I>(chunks: I, fuzz_bytes: &[u8]) -> arbitrary::Result<Self>
-    where
-        I: IntoIterator<Item = Chunk<Sha256>>,
-    {
-        let mut chunks: Vec<_> = chunks.into_iter().collect();
-        let mut u = Unstructured::new(fuzz_bytes);
-
-        for i in (1..chunks.len()).rev() {
-            let j = u.int_in_range(0..=i)?;
-            chunks.swap(i, j);
-        }
-
-        Ok(ShuffledChunks { chunks })
-    }
-}
-
-fn fuzz(input: FuzzInput) {
-    let total = input.total;
-    let min = input.min;
-    let data = input.data;
-    let shuffle_bytes = input.shuffle_bytes;
-
-    // if encode returns Digest then we should be able to decode it later.
-    // we return in Error case, because the underlying library can panic on arbitrary input.
-    let (root, chunks) = match encode::<Sha256>(total, min, data.to_vec()) {
-        Ok(result) => result,
-        Err(_) => return,
-    };
-
-    assert_eq!(chunks.len(), total as usize);
-
-    for (i, chunk) in chunks.iter().enumerate() {
-        assert!(chunk.verify(i as u16, &root), "failed to verify chunk");
-    }
-
-    let decoded = match decode::<Sha256>(total, min, &root, chunks.clone()) {
-        Ok(data) => data,
-        Err(e) => panic!("decode with all chunks failed: {e:?}"),
-    };
-    assert_eq!(decoded, data, "decode with all chunks failed");
-
-    let subset =
-        ShuffledChunks::from_chunks(chunks, &shuffle_bytes).expect("failed to shuffle chunks");
-
-    let decoded_subset = match decode::<Sha256>(total, min, &root, subset.chunks) {
-        Ok(data) => data,
-        Err(e) => panic!("decode with min chunks failed: {e:?}"),
-    };
-    assert_eq!(decoded_subset, data);
-}
-
 fuzz_target!(|input: FuzzInput| {
-    fuzz(input);
+    fuzz::<ReedSolomon<Sha256>>(input);
 });

coding/fuzz/fuzz_targets/reed_solomon.rs

@@ -0,0 +1,86 @@
+use arbitrary::{Arbitrary, Unstructured};
+use commonware_coding::{Config, Scheme};
+
+#[derive(Debug)]
+pub struct FuzzInput {
+    min: u16,
+    recovery: u16,
+    to_use: u16,
+    data: Vec<u8>,
+    shuffle_bytes: Vec<u8>,
+}
+
+impl<'a> Arbitrary<'a> for FuzzInput {
+    fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
+        // We need to generate parameters which satisfy the conditions of valid RS coding,
+        // which are such that if min <= 2^16 - 2^n, then recovery <= 2^n.
+        let n: u64 = u.int_in_range(0..=15)?;
+        let mut min = u.int_in_range(1..=u16::try_from((1 << 16) - (1 << n)).unwrap())?;
+        let mut recovery = u.int_in_range(1..=u16::try_from(1 << n).unwrap())?;
+        // Correction to make sure that we can fit min + recovery in a u16.
+        if min.checked_add(recovery).is_none() {
+            if recovery > 1 {
+                recovery -= 1;
+            } else {
+                min -= 1;
+            }
+        }
+        let to_use = u.int_in_range(min..=min + recovery)?;
+        let data_len = u.int_in_range(0..=u32::MAX)?; // data.len() <= u32:Max
+        let data = u.bytes(data_len as usize)?.to_vec();
+        let shuffle_bytes = u.bytes(8)?.to_vec();
+
+        Ok(FuzzInput {
+            recovery,
+            min,
+            to_use,
+            data,
+            shuffle_bytes,
+        })
+    }
+}
+
+fn shuffle<T>(shuffle_bytes: &[u8], data: &mut [T]) {
+    let mut u = Unstructured::new(shuffle_bytes);
+
+    for i in (1..data.len()).rev() {
+        let j = u.int_in_range(0..=i).unwrap();
+        data.swap(i, j);
+    }
+}
+
+pub fn fuzz<S: Scheme>(input: FuzzInput) {
+    let FuzzInput {
+        recovery,
+        min,
+        to_use,
+        data,
+        shuffle_bytes,
+    } = input;
+
+    let config = Config {
+        minimum_shards: min,
+        extra_shards: recovery,
+    };
+    let (commitment, mut shards) = S::encode(&config, data.as_slice()).unwrap();
+    assert_eq!(shards.len(), (recovery + min) as usize);
+    // Each participant checks their shard
+    let mut reshards = shards
+        .iter()
+        .map(|(shard, proof)| S::check(&commitment, proof, shard).unwrap())
+        .collect::<Vec<_>>();
+    // The last shard is "ours"
+    let (my_shard, _) = shards.pop().unwrap();
+    // We shuffle the remaining reshards
+    reshards.truncate(reshards.len() - 1);
+    shuffle(&shuffle_bytes, &mut reshards);
+    // We decode using the specified number of reshards, and ours.
+    let decoded = S::decode(
+        &config,
+        &commitment,
+        my_shard,
+        &reshards[..(to_use - 1) as usize],
+    )
+    .unwrap();
+    assert_eq!(&decoded, &data);
+}

coding/fuzz/src/lib.rs

@@ -0,0 +1,84 @@
+use commonware_coding::{Config, Scheme};
+use criterion::{criterion_main, BatchSize, Criterion};
+use rand::{seq::SliceRandom, RngCore, SeedableRng as _};
+use rand_chacha::ChaCha8Rng;
+
+mod no_coding;
+mod reed_solomon;
+
+pub(crate) fn benchmark_encode_generic<S: Scheme>(name: &str, c: &mut Criterion) {
+    let mut rng = ChaCha8Rng::seed_from_u64(0);
+    let cases = [8, 12, 16, 19, 20, 24].map(|i| 2usize.pow(i));
+    for data_length in cases.into_iter() {
+        for chunks in [10, 25, 50, 100, 250] {
+            let min = chunks / 3;
+            let config = Config {
+                minimum_shards: min as u16,
+                extra_shards: (chunks - min) as u16,
+            };
+            c.bench_function(
+                &format!("{}/msg_len={} chunks={}", name, data_length, chunks),
+                |b| {
+                    b.iter_batched(
+                        || {
+                            // Generate random data
+                            let mut data = vec![0u8; data_length];
+                            rng.fill_bytes(&mut data);
+                            data
+                        },
+                        |data| S::encode(&config, data.as_slice()),
+                        BatchSize::SmallInput,
+                    );
+                },
+            );
+        }
+    }
+}
+
+pub(crate) fn benchmark_decode_generic<S: Scheme>(name: &str, c: &mut Criterion) {
+    let mut rng = ChaCha8Rng::seed_from_u64(0);
+    let cases = [8, 12, 16, 19, 20, 24].map(|i| 2usize.pow(i));
+    for data_length in cases.into_iter() {
+        for chunks in [10, 25, 50, 100, 250] {
+            let min = chunks / 3;
+            let config = Config {
+                minimum_shards: min as u16,
+                extra_shards: (chunks - min) as u16,
+            };
+            c.bench_function(
+                &format!("{}/msg_len={} chunks={}", name, data_length, chunks),
+                |b| {
+                    b.iter_batched(
+                        || {
+                            // Generate random data
+                            let mut data = vec![0u8; data_length];
+                            rng.fill_bytes(&mut data);
+
+                            // Encode data
+                            let (commitment, mut shards) =
+                                S::encode(&config, data.as_slice()).unwrap();
+
+                            shards.shuffle(&mut rng);
+                            let my_shard_and_proof = shards.pop().unwrap();
+                            let reshards = shards
+                                .iter()
+                                .take(min)
+                                .map(|(shard, proof)| S::check(&commitment, proof, shard).unwrap())
+                                .collect::<Vec<_>>();
+
+                            (commitment, my_shard_and_proof, reshards)
+                        },
+                        // We include the cost of checking your shard as part of decoding
+                        |(commitment, (my_shard, my_proof), reshards)| {
+                            S::check(&commitment, &my_proof, &my_shard).unwrap();
+                            S::decode(&config, &commitment, my_shard, &reshards).unwrap();
+                        },
+                        BatchSize::SmallInput,
+                    );
+                },
+            );
+        }
+    }
+}
+
+criterion_main!(reed_solomon::benches, no_coding::benches);