Add PeerDAS KZG lib integration (construction & KZG verification) (#6212)

* Add peerdas KZG library and use it for data column construction and cell kzg verification (#5701, #5941, #6118, #6179)

Co-authored-by: kevaundray <kevtheappdev@gmail.com>

* Update `rust_eth_kzg` crate to published version.

* Update kzg metrics buckets.

* Merge branch 'unstable' into peerdas-kzg

* Update KZG version to fix windows mem allocation.

* Refactor common logic from build sidecar and reconstruction. Remove unnecessary `needless_lifetimes`.

Co-authored-by: realbigsean <sean@sigmaprime.io>

* Copy existing trusted setup into `PeerDASTrustedSetup` for consistency and maintain `--trusted-setup` functionality.

* Merge branch 'unstable' into peerdas-kzg

* Merge branch 'peerdas-kzg' of github.com:jimmygchen/lighthouse into peerdas-kzg

* Merge branch 'unstable' into peerdas-kzg

* Merge branch 'unstable' into peerdas-kzg

* Load PeerDAS KZG only if PeerDAS is enabled.

beacon_node/beacon_chain/src/data_column_verification.rs

@@ -2,7 +2,8 @@ use crate::block_verification::{
     cheap_state_advance_to_obtain_committees, get_validator_pubkey_cache, process_block_slash_info,
     BlockSlashInfo,
 };
-use crate::{BeaconChain, BeaconChainError, BeaconChainTypes};
+use crate::kzg_utils::validate_data_columns;
+use crate::{metrics, BeaconChain, BeaconChainError, BeaconChainTypes};
 use derivative::Derivative;
 use fork_choice::ProtoBlock;
 use kzg::{Error as KzgError, Kzg};
@@ -11,6 +12,7 @@ use slasher::test_utils::E;
 use slog::debug;
 use slot_clock::SlotClock;
 use ssz_derive::{Decode, Encode};
+use std::iter;
 use std::sync::Arc;
 use types::data_column_sidecar::{ColumnIndex, DataColumnIdentifier};
 use types::{
@@ -255,9 +257,10 @@ impl<E: EthSpec> KzgVerifiedCustodyDataColumn<E> {
 /// Returns an error if the kzg verification check fails.
 pub fn verify_kzg_for_data_column<E: EthSpec>(
     data_column: Arc<DataColumnSidecar<E>>,
-    _kzg: &Kzg,
+    kzg: &Kzg,
 ) -> Result<KzgVerifiedDataColumn<E>, KzgError> {
-    // TODO(das): KZG verification to be implemented
+    let _timer = metrics::start_timer(&metrics::KZG_VERIFICATION_DATA_COLUMN_SINGLE_TIMES);
+    validate_data_columns(kzg, iter::once(&data_column))?;
     Ok(KzgVerifiedDataColumn { data: data_column })
 }
 
@@ -267,13 +270,14 @@ pub fn verify_kzg_for_data_column<E: EthSpec>(
 /// Note: This function should be preferred over calling `verify_kzg_for_data_column`
 /// in a loop since this function kzg verifies a list of data columns more efficiently.
 pub fn verify_kzg_for_data_column_list<'a, E: EthSpec, I>(
-    _data_column_iter: I,
-    _kzg: &'a Kzg,
+    data_column_iter: I,
+    kzg: &'a Kzg,
 ) -> Result<(), KzgError>
 where
     I: Iterator<Item = &'a Arc<DataColumnSidecar<E>>> + Clone,
 {
-    // TODO(das): implement KZG verification
+    let _timer = metrics::start_timer(&metrics::KZG_VERIFICATION_DATA_COLUMN_BATCH_TIMES);
+    validate_data_columns(kzg, data_column_iter)?;
     Ok(())
 }
 

beacon_node/beacon_chain/src/kzg_utils.rs

@@ -1,5 +1,15 @@
-use kzg::{Blob as KzgBlob, Error as KzgError, Kzg};
-use types::{Blob, EthSpec, Hash256, KzgCommitment, KzgProof};
+use kzg::{
+    Blob as KzgBlob, Bytes48, CellRef as KzgCellRef, CellsAndKzgProofs, Error as KzgError, Kzg,
+};
+use rayon::prelude::*;
+use ssz_types::FixedVector;
+use std::sync::Arc;
+use types::beacon_block_body::KzgCommitments;
+use types::data_column_sidecar::{Cell, DataColumn, DataColumnSidecarError};
+use types::{
+    Blob, BlobsList, ChainSpec, ColumnIndex, DataColumnSidecar, DataColumnSidecarList, EthSpec,
+    Hash256, KzgCommitment, KzgProof, KzgProofs, SignedBeaconBlock, SignedBeaconBlockHeader,
+};
 
 /// Converts a blob ssz List object to an array to be used with the kzg
 /// crypto library.
@@ -7,6 +17,15 @@ fn ssz_blob_to_crypto_blob<E: EthSpec>(blob: &Blob<E>) -> Result<KzgBlob, KzgErr
     KzgBlob::from_bytes(blob.as_ref()).map_err(Into::into)
 }
 
+/// Converts a cell ssz List object to an array to be used with the kzg
+/// crypto library.
+fn ssz_cell_to_crypto_cell<E: EthSpec>(cell: &Cell<E>) -> Result<KzgCellRef, KzgError> {
+    let cell_bytes: &[u8] = cell.as_ref();
+    Ok(cell_bytes
+        .try_into()
+        .expect("expected cell to have size {BYTES_PER_CELL}. This should be guaranteed by the `FixedVector type"))
+}
+
 /// Validate a single blob-commitment-proof triplet from a `BlobSidecar`.
 pub fn validate_blob<E: EthSpec>(
     kzg: &Kzg,
@@ -19,6 +38,50 @@ pub fn validate_blob<E: EthSpec>(
     kzg.verify_blob_kzg_proof(&kzg_blob, kzg_commitment, kzg_proof)
 }
 
+/// Validate a batch of `DataColumnSidecar`.
+pub fn validate_data_columns<'a, E: EthSpec, I>(
+    kzg: &Kzg,
+    data_column_iter: I,
+) -> Result<(), KzgError>
+where
+    I: Iterator<Item = &'a Arc<DataColumnSidecar<E>>> + Clone,
+{
+    let cells = data_column_iter
+        .clone()
+        .flat_map(|data_column| data_column.column.iter().map(ssz_cell_to_crypto_cell::<E>))
+        .collect::<Result<Vec<_>, KzgError>>()?;
+
+    let proofs = data_column_iter
+        .clone()
+        .flat_map(|data_column| {
+            data_column
+                .kzg_proofs
+                .iter()
+                .map(|&proof| Bytes48::from(proof))
+        })
+        .collect::<Vec<_>>();
+
+    let column_indices = data_column_iter
+        .clone()
+        .flat_map(|data_column| {
+            let col_index = data_column.index;
+            data_column.column.iter().map(move |_| col_index)
+        })
+        .collect::<Vec<ColumnIndex>>();
+
+    let commitments = data_column_iter
+        .clone()
+        .flat_map(|data_column| {
+            data_column
+                .kzg_commitments
+                .iter()
+                .map(|&commitment| Bytes48::from(commitment))
+        })
+        .collect::<Vec<_>>();
+
+    kzg.verify_cell_proof_batch(&cells, &proofs, column_indices, &commitments)
+}
+
 /// Validate a batch of blob-commitment-proof triplets from multiple `BlobSidecars`.
 pub fn validate_blobs<E: EthSpec>(
     kzg: &Kzg,
@@ -76,3 +139,264 @@ pub fn verify_kzg_proof<E: EthSpec>(
 ) -> Result<bool, KzgError> {
     kzg.verify_kzg_proof(kzg_commitment, &z.0.into(), &y.0.into(), kzg_proof)
 }
+
+/// Build data column sidecars from a signed beacon block and its blobs.
+pub fn blobs_to_data_column_sidecars<E: EthSpec>(
+    blobs: &BlobsList<E>,
+    block: &SignedBeaconBlock<E>,
+    kzg: &Kzg,
+    spec: &ChainSpec,
+) -> Result<DataColumnSidecarList<E>, DataColumnSidecarError> {
+    if blobs.is_empty() {
+        return Ok(vec![]);
+    }
+    let kzg_commitments = block
+        .message()
+        .body()
+        .blob_kzg_commitments()
+        .map_err(|_err| DataColumnSidecarError::PreDeneb)?;
+    let kzg_commitments_inclusion_proof = block.message().body().kzg_commitments_merkle_proof()?;
+    let signed_block_header = block.signed_block_header();
+
+    // NOTE: assumes blob sidecars are ordered by index
+    let blob_cells_and_proofs_vec = blobs
+        .into_par_iter()
+        .map(|blob| {
+            let blob = blob
+                .as_ref()
+                .try_into()
+                .expect("blob should have a guaranteed size due to FixedVector");
+            kzg.compute_cells_and_proofs(blob)
+        })
+        .collect::<Result<Vec<_>, KzgError>>()?;
+
+    build_data_column_sidecars(
+        kzg_commitments.clone(),
+        kzg_commitments_inclusion_proof,
+        signed_block_header,
+        blob_cells_and_proofs_vec,
+        spec,
+    )
+    .map_err(DataColumnSidecarError::BuildSidecarFailed)
+}
+
+fn build_data_column_sidecars<E: EthSpec>(
+    kzg_commitments: KzgCommitments<E>,
+    kzg_commitments_inclusion_proof: FixedVector<Hash256, E::KzgCommitmentsInclusionProofDepth>,
+    signed_block_header: SignedBeaconBlockHeader,
+    blob_cells_and_proofs_vec: Vec<CellsAndKzgProofs>,
+    spec: &ChainSpec,
+) -> Result<DataColumnSidecarList<E>, String> {
+    let number_of_columns = spec.number_of_columns;
+    let mut columns = vec![Vec::with_capacity(E::max_blobs_per_block()); number_of_columns];
+    let mut column_kzg_proofs =
+        vec![Vec::with_capacity(E::max_blobs_per_block()); number_of_columns];
+
+    for (blob_cells, blob_cell_proofs) in blob_cells_and_proofs_vec {
+        // we iterate over each column, and we construct the column from "top to bottom",
+        // pushing on the cell and the corresponding proof at each column index. we do this for
+        // each blob (i.e. the outer loop).
+        for col in 0..number_of_columns {
+            let cell = blob_cells
+                .get(col)
+                .ok_or(format!("Missing blob cell at index {col}"))?;
+            let cell: Vec<u8> = cell.to_vec();
+            let cell = Cell::<E>::from(cell);
+
+            let proof = blob_cell_proofs
+                .get(col)
+                .ok_or(format!("Missing blob cell KZG proof at index {col}"))?;
+
+            let column = columns
+                .get_mut(col)
+                .ok_or(format!("Missing data column at index {col}"))?;
+            let column_proofs = column_kzg_proofs
+                .get_mut(col)
+                .ok_or(format!("Missing data column proofs at index {col}"))?;
+
+            column.push(cell);
+            column_proofs.push(*proof);
+        }
+    }
+
+    let sidecars: Vec<Arc<DataColumnSidecar<E>>> = columns
+        .into_iter()
+        .zip(column_kzg_proofs)
+        .enumerate()
+        .map(|(index, (col, proofs))| {
+            Arc::new(DataColumnSidecar {
+                index: index as u64,
+                column: DataColumn::<E>::from(col),
+                kzg_commitments: kzg_commitments.clone(),
+                kzg_proofs: KzgProofs::<E>::from(proofs),
+                signed_block_header: signed_block_header.clone(),
+                kzg_commitments_inclusion_proof: kzg_commitments_inclusion_proof.clone(),
+            })
+        })
+        .collect();
+
+    Ok(sidecars)
+}
+
+/// Reconstruct all data columns from a subset of data column sidecars (requires at least 50%).
+pub fn reconstruct_data_columns<E: EthSpec>(
+    kzg: &Kzg,
+    data_columns: &[Arc<DataColumnSidecar<E>>],
+    spec: &ChainSpec,
+) -> Result<DataColumnSidecarList<E>, KzgError> {
+    let first_data_column = data_columns
+        .first()
+        .ok_or(KzgError::InconsistentArrayLength(
+            "data_columns should have at least one element".to_string(),
+        ))?;
+    let num_of_blobs = first_data_column.kzg_commitments.len();
+
+    let blob_cells_and_proofs_vec =
+        (0..num_of_blobs)
+            .into_par_iter()
+            .map(|row_index| {
+                let mut cells: Vec<KzgCellRef> = vec![];
+                let mut cell_ids: Vec<u64> = vec![];
+                for data_column in data_columns {
+                    let cell = data_column.column.get(row_index).ok_or(
+                        KzgError::InconsistentArrayLength(format!(
+                            "Missing data column at index {row_index}"
+                        )),
+                    )?;
+
+                    cells.push(ssz_cell_to_crypto_cell::<E>(cell)?);
+                    cell_ids.push(data_column.index);
+                }
+                kzg.recover_cells_and_compute_kzg_proofs(&cell_ids, &cells)
+            })
+            .collect::<Result<Vec<_>, KzgError>>()?;
+
+    // Clone sidecar elements from existing data column, no need to re-compute
+    build_data_column_sidecars(
+        first_data_column.kzg_commitments.clone(),
+        first_data_column.kzg_commitments_inclusion_proof.clone(),
+        first_data_column.signed_block_header.clone(),
+        blob_cells_and_proofs_vec,
+        spec,
+    )
+    .map_err(KzgError::ReconstructFailed)
+}
+
+#[cfg(test)]
+mod test {
+    use crate::kzg_utils::{blobs_to_data_column_sidecars, reconstruct_data_columns};
+    use bls::Signature;
+    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_data_columns_sidecars() {
+        let spec = E::default_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();
+        assert!(column_sidecars.is_empty());
+    }
+
+    #[track_caller]
+    fn test_build_data_columns(kzg: &Kzg, spec: &ChainSpec) {
+        let num_of_blobs = 6;
+        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();
+
+        let block_kzg_commitments = signed_block
+            .message()
+            .body()
+            .blob_kzg_commitments()
+            .unwrap()
+            .clone();
+        let block_kzg_commitments_inclusion_proof = signed_block
+            .message()
+            .body()
+            .kzg_commitments_merkle_proof()
+            .unwrap();
+
+        assert_eq!(column_sidecars.len(), spec.number_of_columns);
+        for (idx, col_sidecar) in column_sidecars.iter().enumerate() {
+            assert_eq!(col_sidecar.index, idx as u64);
+
+            assert_eq!(col_sidecar.kzg_commitments.len(), num_of_blobs);
+            assert_eq!(col_sidecar.column.len(), num_of_blobs);
+            assert_eq!(col_sidecar.kzg_proofs.len(), num_of_blobs);
+
+            assert_eq!(col_sidecar.kzg_commitments, block_kzg_commitments);
+            assert_eq!(
+                col_sidecar.kzg_commitments_inclusion_proof,
+                block_kzg_commitments_inclusion_proof
+            );
+            assert!(col_sidecar.verify_inclusion_proof());
+        }
+    }
+
+    #[track_caller]
+    fn test_reconstruct_data_columns(kzg: &Kzg, spec: &ChainSpec) {
+        let num_of_blobs = 6;
+        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();
+
+        // Now reconstruct
+        let reconstructed_columns = reconstruct_data_columns(
+            kzg,
+            &column_sidecars.iter().as_slice()[0..column_sidecars.len() / 2],
+            spec,
+        )
+        .unwrap();
+
+        for i in 0..spec.number_of_columns {
+            assert_eq!(reconstructed_columns.get(i), column_sidecars.get(i), "{i}");
+        }
+    }
+
+    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,
+    ) -> (SignedBeaconBlock<E>, BlobsList<E>) {
+        let mut block = BeaconBlock::Deneb(BeaconBlockDeneb::empty(spec));
+        let mut body = block.body_mut();
+        let blob_kzg_commitments = body.blob_kzg_commitments_mut().unwrap();
+        *blob_kzg_commitments =
+            KzgCommitments::<E>::new(vec![KzgCommitment::empty_for_testing(); num_of_blobs])
+                .unwrap();
+
+        let signed_block = SignedBeaconBlock::from_block(block, Signature::empty());
+
+        let blobs = (0..num_of_blobs)
+            .map(|_| Blob::<E>::default())
+            .collect::<Vec<_>>()
+            .into();
+
+        (signed_block, blobs)
+    }
+}

beacon_node/beacon_chain/src/metrics.rs

@@ -1645,6 +1645,13 @@ pub static BLOB_SIDECAR_INCLUSION_PROOF_COMPUTATION: LazyLock<Result<Histogram>>
             "Time taken to compute blob sidecar inclusion proof",
         )
     });
+pub static DATA_COLUMN_SIDECAR_COMPUTATION: LazyLock<Result<Histogram>> = LazyLock::new(|| {
+    try_create_histogram_with_buckets(
+        "data_column_sidecar_computation_seconds",
+        "Time taken to compute data column sidecar, including cells, proofs and inclusion proof",
+        Ok(vec![0.04, 0.05, 0.1, 0.2, 0.3, 0.5, 0.7, 1.0]),
+    )
+});
 pub static DATA_COLUMN_SIDECAR_PROCESSING_REQUESTS: LazyLock<Result<IntCounter>> =
     LazyLock::new(|| {
         try_create_int_counter(
@@ -1785,6 +1792,26 @@ pub static KZG_VERIFICATION_BATCH_TIMES: LazyLock<Result<Histogram>> = LazyLock:
         "Runtime of batched kzg verification",
     )
 });
+pub static KZG_VERIFICATION_DATA_COLUMN_SINGLE_TIMES: LazyLock<Result<Histogram>> =
+    LazyLock::new(|| {
+        try_create_histogram_with_buckets(
+            "kzg_verification_data_column_single_seconds",
+            "Runtime of single data column kzg verification",
+            Ok(vec![
+                0.0005, 0.001, 0.0015, 0.002, 0.003, 0.004, 0.005, 0.007, 0.01, 0.02, 0.05,
+            ]),
+        )
+    });
+pub static KZG_VERIFICATION_DATA_COLUMN_BATCH_TIMES: LazyLock<Result<Histogram>> =
+    LazyLock::new(|| {
+        try_create_histogram_with_buckets(
+            "kzg_verification_data_column_batch_seconds",
+            "Runtime of batched data column kzg verification",
+            Ok(vec![
+                0.002, 0.004, 0.006, 0.008, 0.01, 0.012, 0.015, 0.02, 0.03, 0.05, 0.07,
+            ]),
+        )
+    });
 
 pub static BLOCK_PRODUCTION_BLOBS_VERIFICATION_TIMES: LazyLock<Result<Histogram>> = LazyLock::new(
     || {

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
         };