Skip to content

Commit

Permalink
Load PeerDAS KZG only if PeerDAS is enabled.
Browse files Browse the repository at this point in the history
  • Loading branch information
jimmygchen committed Aug 12, 2024
1 parent 8156092 commit 856fa89
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 44 deletions.
57 changes: 34 additions & 23 deletions beacon_node/beacon_chain/src/kzg_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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();

Expand All @@ -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,
Expand Down
19 changes: 11 additions & 8 deletions beacon_node/client/src/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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;
Expand Down Expand Up @@ -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) => {
Expand Down Expand Up @@ -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
};
Expand Down
34 changes: 21 additions & 13 deletions crypto/kzg/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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)?;

Expand All @@ -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(),
Expand All @@ -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)?;

Expand All @@ -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)
}
}

0 comments on commit 856fa89

Please sign in to comment.