From ed0106b6fd96832f425bd5608adcc5b69de1eff9 Mon Sep 17 00:00:00 2001
From: Anders Konring <anders.konring@gmail.com>
Date: Fri, 19 Apr 2024 09:30:28 +0200
Subject: [PATCH] feat: add method to verify consistent precomputed data

---
 vid/src/advz.rs               | 14 ++++++------
 vid/src/advz/precomputable.rs | 40 +++++++++++++++++++++++++++++++++--
 vid/src/precomputable.rs      |  9 ++++++++
 3 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/vid/src/advz.rs b/vid/src/advz.rs
index f7b184fa1..684e3e04c 100644
--- a/vid/src/advz.rs
+++ b/vid/src/advz.rs
@@ -1022,7 +1022,7 @@ mod tests {
 
     #[test]
     fn sad_path_verify_share_corrupt_share() {
-        let (mut advz, bytes_random) = advz_init();
+        let (mut advz, bytes_random, _) = advz_init();
         let disperse = advz.disperse(bytes_random).unwrap();
         let (shares, common, commit) = (disperse.shares, disperse.common, disperse.commit);
 
@@ -1088,7 +1088,7 @@ mod tests {
 
     #[test]
     fn sad_path_verify_share_corrupt_commit() {
-        let (mut advz, bytes_random) = advz_init();
+        let (mut advz, bytes_random, _) = advz_init();
         let disperse = advz.disperse(bytes_random).unwrap();
         let (shares, common, commit) = (disperse.shares, disperse.common, disperse.commit);
 
@@ -1134,7 +1134,7 @@ mod tests {
 
     #[test]
     fn sad_path_verify_share_corrupt_share_and_commit() {
-        let (mut advz, bytes_random) = advz_init();
+        let (mut advz, bytes_random, _) = advz_init();
         let disperse = advz.disperse(bytes_random).unwrap();
         let (mut shares, mut common, commit) = (disperse.shares, disperse.common, disperse.commit);
 
@@ -1159,8 +1159,8 @@ mod tests {
 
     #[test]
     fn sad_path_recover_payload_corrupt_shares() {
-        let (mut advz, bytes_random) = advz_init();
-        let disperse = advz.disperse(&bytes_random).unwrap();
+        let (mut advz, bytes_random, _) = advz_init();
+        let disperse = advz.disperse(bytes_random.clone()).unwrap();
         let (shares, common) = (disperse.shares, disperse.common);
 
         {
@@ -1219,13 +1219,13 @@ mod tests {
     /// Returns the following tuple:
     /// 1. An initialized [`Advz`] instance.
     /// 2. A `Vec<u8>` filled with random bytes.
-    pub(super) fn advz_init() -> (Advz<Bls12_381, Sha256>, Vec<u8>) {
+    pub(super) fn advz_init() -> (Advz<Bls12_381, Sha256>, Vec<u8>, u32) {
         let (recovery_threshold, num_storage_nodes) = (4, 6);
         let mut rng = jf_utils::test_rng();
         let srs = init_srs(recovery_threshold as usize, &mut rng);
         let advz = Advz::new(num_storage_nodes, recovery_threshold, srs).unwrap();
         let bytes_random = init_random_payload(4000, &mut rng);
-        (advz, bytes_random)
+        (advz, bytes_random, num_storage_nodes)
     }
 
     /// Convenience wrapper to assert [`VidError::Argument`] return value.
diff --git a/vid/src/advz/precomputable.rs b/vid/src/advz/precomputable.rs
index f53b9c283..b080faaaf 100644
--- a/vid/src/advz/precomputable.rs
+++ b/vid/src/advz/precomputable.rs
@@ -6,6 +6,7 @@
 
 //! Implementations of [`Precomputable`] for `Advz`.
 
+use crate::VidError;
 use crate::{
     advz::{
         polynomial_eval, AdvzInternal, Common, HasherDigest, KzgCommit, KzgEvalsMerkleTree,
@@ -14,14 +15,16 @@ use crate::{
     precomputable::Precomputable,
     vid, VidDisperse, VidResult,
 };
+use alloc::string::ToString;
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 use ark_std::{end_timer, start_timer, vec::Vec};
 use jf_merkle_tree::{MerkleCommitment, MerkleTreeScheme};
 use jf_pcs::{prelude::Commitment, PolynomialCommitmentScheme, UnivariatePCS};
-
 use jf_utils::canonical;
 use serde::{Deserialize, Serialize};
 
+use super::Advz;
+
 impl<E, H, T> Precomputable for AdvzInternal<E, H, T>
 where
     E: Pairing,
@@ -139,6 +142,26 @@ where
             commit,
         })
     }
+
+    fn is_consistent_precompute(
+        commit: &Self::Commit,
+        precompute_data: &Self::PrecomputeData,
+        payload_byte_len: u32,
+        num_storage_nodes: u32,
+    ) -> VidResult<()> {
+        if *commit
+            != Advz::<E, H>::derive_commit(
+                &precompute_data.poly_commits,
+                payload_byte_len,
+                num_storage_nodes,
+            )?
+        {
+            return Err(VidError::Argument(
+                "precompute data inconsistent with commit".to_string(),
+            ));
+        }
+        Ok(())
+    }
 }
 
 #[derive(
@@ -217,7 +240,7 @@ mod tests {
 
     #[test]
     fn commit_disperse_recover_with_precomputed_data() {
-        let (advz, bytes_random) = advz_init();
+        let (advz, bytes_random, _) = advz_init();
         let (commit, data) = advz.commit_only_precompute(&bytes_random).unwrap();
         let disperse = advz.disperse_precompute(&bytes_random, &data).unwrap();
         let (shares, common) = (disperse.shares, disperse.common);
@@ -231,4 +254,17 @@ mod tests {
             .expect("recover_payload should succeed");
         assert_eq!(bytes_recovered, bytes_random);
     }
+
+    #[test]
+    fn commit_and_verify_consistent_precomputed_data() {
+        let (advz, bytes_random, num_storage_nodes) = advz_init();
+        let (commit, data) = advz.commit_only_precompute(&bytes_random).unwrap();
+        assert!(Advz::is_consistent_precompute(
+            &commit,
+            &data,
+            bytes_random.len() as u32,
+            num_storage_nodes
+        )
+        .is_ok())
+    }
 }
diff --git a/vid/src/precomputable.rs b/vid/src/precomputable.rs
index 584704770..564fa0ec0 100644
--- a/vid/src/precomputable.rs
+++ b/vid/src/precomputable.rs
@@ -36,4 +36,13 @@ pub trait Precomputable: VidScheme {
     ) -> VidResult<VidDisperse<Self>>
     where
         B: AsRef<[u8]>;
+
+    /// Check that a [`Precomputable::PrecomputeData`] is consistent with a
+    /// [`VidScheme::Commit`].
+    fn is_consistent_precompute(
+        commit: &Self::Commit,
+        precompute_data: &Self::PrecomputeData,
+        payload_byte_len: u32,
+        num_storage_nodes: u32,
+    ) -> VidResult<()>;
 }