Load PeerDAS KZG only if PeerDAS is enabled.

Author: jimmygchen

beacon_node/beacon_chain/src/kzg_utils.rs

@@ -285,36 +285,43 @@ pub fn reconstruct_data_columns<E: EthSpec>(
 #[cfg(test)]
 mod test {
     use crate::kzg_utils::{blobs_to_data_column_sidecars, reconstruct_data_columns};
-    use crate::test_utils::KZG;
     use bls::Signature;
-    use kzg::KzgCommitment;
+    use eth2_network_config::TRUSTED_SETUP_BYTES;
+    use kzg::{Kzg, KzgCommitment, TrustedSetup};
     use types::{
         beacon_block_body::KzgCommitments, BeaconBlock, BeaconBlockDeneb, Blob, BlobsList,
         ChainSpec, EmptyBlock, EthSpec, MainnetEthSpec, SignedBeaconBlock,
     };
 
+    type E = MainnetEthSpec;
+
+    // Loading and initializing PeerDAS KZG is expensive and slow, so we group the tests together
+    // only load it once.
     #[test]
-    fn test_build_sidecars_empty() {
-        type E = MainnetEthSpec;
-        let num_of_blobs = 0;
+    fn test_build_data_columns_sidecars() {
         let spec = E::default_spec();
-        let (signed_block, blob_sidecars) = create_test_block_and_blobs::<E>(num_of_blobs, &spec);
+        let kzg = get_kzg();
+        test_build_data_columns_empty(&kzg, &spec);
+        test_build_data_columns(&kzg, &spec);
+        test_reconstruct_data_columns(&kzg, &spec);
+    }
 
+    #[track_caller]
+    fn test_build_data_columns_empty(kzg: &Kzg, spec: &ChainSpec) {
+        let num_of_blobs = 0;
+        let (signed_block, blob_sidecars) = create_test_block_and_blobs::<E>(num_of_blobs, spec);
         let column_sidecars =
-            blobs_to_data_column_sidecars(&blob_sidecars, &signed_block, &KZG, &spec).unwrap();
-
+            blobs_to_data_column_sidecars(&blob_sidecars, &signed_block, kzg, spec).unwrap();
         assert!(column_sidecars.is_empty());
     }
 
-    #[test]
-    fn test_build_sidecars() {
-        type E = MainnetEthSpec;
+    #[track_caller]
+    fn test_build_data_columns(kzg: &Kzg, spec: &ChainSpec) {
         let num_of_blobs = 6;
-        let spec = E::default_spec();
-        let (signed_block, blob_sidecars) = create_test_block_and_blobs::<E>(num_of_blobs, &spec);
+        let (signed_block, blob_sidecars) = create_test_block_and_blobs::<E>(num_of_blobs, spec);
 
         let column_sidecars =
-            blobs_to_data_column_sidecars(&blob_sidecars, &signed_block, &KZG, &spec).unwrap();
+            blobs_to_data_column_sidecars(&blob_sidecars, &signed_block, kzg, spec).unwrap();
 
         let block_kzg_commitments = signed_block
             .message()
@@ -345,21 +352,18 @@ mod test {
         }
     }
 
-    #[test]
-    fn build_and_reconstruct() {
-        type E = MainnetEthSpec;
+    #[track_caller]
+    fn test_reconstruct_data_columns(kzg: &Kzg, spec: &ChainSpec) {
         let num_of_blobs = 6;
-        let spec = E::default_spec();
-        let (signed_block, blob_sidecars) = create_test_block_and_blobs::<E>(num_of_blobs, &spec);
-
+        let (signed_block, blob_sidecars) = create_test_block_and_blobs::<E>(num_of_blobs, spec);
         let column_sidecars =
-            blobs_to_data_column_sidecars(&blob_sidecars, &signed_block, &KZG, &spec).unwrap();
+            blobs_to_data_column_sidecars(&blob_sidecars, &signed_block, kzg, spec).unwrap();
 
         // Now reconstruct
         let reconstructed_columns = reconstruct_data_columns(
-            &KZG,
+            kzg,
             &column_sidecars.iter().as_slice()[0..column_sidecars.len() / 2],
-            &spec,
+            spec,
         )
         .unwrap();
 
@@ -368,6 +372,13 @@ mod test {
         }
     }
 
+    fn get_kzg() -> Kzg {
+        let trusted_setup: TrustedSetup = serde_json::from_reader(TRUSTED_SETUP_BYTES)
+            .map_err(|e| format!("Unable to read trusted setup file: {}", e))
+            .expect("should have trusted setup");
+        Kzg::new_from_trusted_setup_das_enabled(trusted_setup).expect("should create kzg")
+    }
+
     fn create_test_block_and_blobs<E: EthSpec>(
         num_of_blobs: usize,
         spec: &ChainSpec,

beacon_node/client/src/builder.rs

@@ -10,7 +10,6 @@ use beacon_chain::graffiti_calculator::start_engine_version_cache_refresh_servic
 use beacon_chain::otb_verification_service::start_otb_verification_service;
 use beacon_chain::proposer_prep_service::start_proposer_prep_service;
 use beacon_chain::schema_change::migrate_schema;
-use beacon_chain::LightClientProducerEvent;
 use beacon_chain::{
     builder::{BeaconChainBuilder, Witness},
     eth1_chain::{CachingEth1Backend, Eth1Chain},
@@ -19,6 +18,7 @@ use beacon_chain::{
     store::{HotColdDB, ItemStore, LevelDB, StoreConfig},
     BeaconChain, BeaconChainTypes, Eth1ChainBackend, MigratorConfig, ServerSentEventHandler,
 };
+use beacon_chain::{Kzg, LightClientProducerEvent};
 use beacon_processor::{BeaconProcessor, BeaconProcessorChannels};
 use beacon_processor::{BeaconProcessorConfig, BeaconProcessorQueueLengths};
 use environment::RuntimeContext;
@@ -505,7 +505,7 @@ where
                     deposit_snapshot.and_then(|snapshot| match Eth1Service::from_deposit_snapshot(
                         config.eth1,
                         context.log().clone(),
-                        spec,
+                        spec.clone(),
                         &snapshot,
                     ) {
                         Ok(service) => {
@@ -624,12 +624,15 @@ where
         };
 
         let beacon_chain_builder = if let Some(trusted_setup) = config.trusted_setup {
-            let kzg = trusted_setup
-                .try_into()
-                .map(Arc::new)
-                .map(Some)
-                .map_err(|e| format!("Failed to load trusted setup: {:?}", e))?;
-            beacon_chain_builder.kzg(kzg)
+            let kzg_err_msg = |e| format!("Failed to load trusted setup: {:?}", e);
+
+            let kzg = if spec.is_peer_das_scheduled() {
+                Kzg::new_from_trusted_setup_das_enabled(trusted_setup).map_err(kzg_err_msg)?
+            } else {
+                Kzg::new_from_trusted_setup(trusted_setup).map_err(kzg_err_msg)?
+            };
+
+            beacon_chain_builder.kzg(Some(Arc::new(kzg)))
         } else {
             beacon_chain_builder
         };

crypto/kzg/src/lib.rs

@@ -37,6 +37,8 @@ pub enum Error {
     InconsistentArrayLength(String),
     /// Error reconstructing data columns.
     ReconstructFailed(String),
+    /// Kzg was not initialized with PeerDAS enabled.
+    DASContextUninitialized,
 }
 
 impl From<c_kzg::Error> for Error {
@@ -49,12 +51,22 @@ impl From<c_kzg::Error> for Error {
 #[derive(Debug)]
 pub struct Kzg {
     trusted_setup: KzgSettings,
-    context: DASContext,
+    context: Option<DASContext>,
 }
 
 impl Kzg {
     /// Load the kzg trusted setup parameters from a vec of G1 and G2 points.
     pub fn new_from_trusted_setup(trusted_setup: TrustedSetup) -> Result<Self, Error> {
+        Ok(Self {
+            trusted_setup: KzgSettings::load_trusted_setup(
+                &trusted_setup.g1_points(),
+                &trusted_setup.g2_points(),
+            )?,
+            context: None,
+        })
+    }
+
+    pub fn new_from_trusted_setup_das_enabled(trusted_setup: TrustedSetup) -> Result<Self, Error> {
         // Initialize the trusted setup using default parameters
         //
         // Note: One can also use `from_json` to initialize it from the consensus-specs
@@ -72,10 +84,14 @@ impl Kzg {
                 &trusted_setup.g1_points(),
                 &trusted_setup.g2_points(),
             )?,
-            context,
+            context: Some(context),
         })
     }
 
+    fn context(&self) -> Result<&DASContext, Error> {
+        self.context.as_ref().ok_or(Error::DASContextUninitialized)
+    }
+
     /// Compute the kzg proof given a blob and its kzg commitment.
     pub fn compute_blob_kzg_proof(
         &self,
@@ -180,7 +196,7 @@ impl Kzg {
         blob: KzgBlobRef<'_>,
     ) -> Result<CellsAndKzgProofs, Error> {
         let (cells, proofs) = self
-            .context
+            .context()?
             .compute_cells_and_kzg_proofs(blob)
             .map_err(Error::PeerDASKZG)?;
 
@@ -206,7 +222,7 @@ impl Kzg {
             .iter()
             .map(|commitment| commitment.as_ref())
             .collect();
-        let verification_result = self.context.verify_cell_kzg_proof_batch(
+        let verification_result = self.context()?.verify_cell_kzg_proof_batch(
             commitments.to_vec(),
             columns,
             cells.to_vec(),
@@ -227,7 +243,7 @@ impl Kzg {
         cells: &[CellRef<'_>],
     ) -> Result<CellsAndKzgProofs, Error> {
         let (cells, proofs) = self
-            .context
+            .context()?
             .recover_cells_and_proofs(cell_ids.to_vec(), cells.to_vec())
             .map_err(Error::PeerDASKZG)?;
 
@@ -236,11 +252,3 @@ impl Kzg {
         Ok((cells, c_kzg_proof))
     }
 }
-
-impl TryFrom<TrustedSetup> for Kzg {
-    type Error = Error;
-
-    fn try_from(trusted_setup: TrustedSetup) -> Result<Self, Self::Error> {
-        Kzg::new_from_trusted_setup(trusted_setup)
-    }
-}