From bb38e21021a9eea82aa1881a47c6e57136083da5 Mon Sep 17 00:00:00 2001 From: Noah Citron Date: Mon, 26 Aug 2024 17:07:12 -0400 Subject: [PATCH] refactor: unify bls types (#357) * refactor: unify bls types * ci * remove clone --- consensus-core/src/consensus_core.rs | 33 +++++++-------- consensus-core/src/lib.rs | 2 +- consensus-core/src/types/bls.rs | 40 +++++++++++++++++++ .../src/types/{primitives.rs => bytes.rs} | 0 consensus-core/src/types/mod.rs | 34 ++++++++-------- consensus-core/src/utils.rs | 14 +------ consensus/src/consensus.rs | 17 ++++---- 7 files changed, 86 insertions(+), 54 deletions(-) create mode 100644 consensus-core/src/types/bls.rs rename consensus-core/src/types/{primitives.rs => bytes.rs} (100%) diff --git a/consensus-core/src/consensus_core.rs b/consensus-core/src/consensus_core.rs index f51880a7..7acbe727 100644 --- a/consensus-core/src/consensus_core.rs +++ b/consensus-core/src/consensus_core.rs @@ -2,7 +2,6 @@ use std::cmp; use alloy::primitives::B256; use eyre::Result; -use milagro_bls::PublicKey; use ssz_types::{BitVector, FixedVector}; use tracing::{info, warn}; use tree_hash::TreeHash; @@ -11,13 +10,13 @@ use zduny_wasm_timer::{SystemTime, UNIX_EPOCH}; use common::config::types::Forks; use crate::errors::ConsensusError; +use crate::types::bls::{PublicKey, Signature}; use crate::types::{ - FinalityUpdate, GenericUpdate, Header, LightClientStore, OptimisticUpdate, SignatureBytes, - SyncCommittee, Update, + FinalityUpdate, GenericUpdate, Header, LightClientStore, OptimisticUpdate, SyncCommittee, + Update, }; use crate::utils::{ - calc_sync_period, compute_domain, compute_fork_data_root, compute_signing_root, - is_aggregate_valid, is_proof_valid, + calc_sync_period, compute_domain, compute_fork_data_root, compute_signing_root, is_proof_valid, }; pub fn get_participating_keys( @@ -28,8 +27,7 @@ pub fn get_participating_keys( bitfield.iter().enumerate().for_each(|(i, bit)| { if bit { - let pk = &committee.pubkeys[i]; - let pk = PublicKey::from_bytes_unchecked(&pk.inner).unwrap(); + let pk = committee.pubkeys[i].clone(); pks.push(pk); } }); @@ -43,7 +41,7 @@ pub fn get_bits(bitfield: &BitVector) -> u64 { pub fn is_finality_proof_valid( attested_header: &Header, - finality_header: &mut Header, + finality_header: &Header, finality_branch: &[B256], ) -> bool { is_proof_valid(attested_header, finality_header, finality_branch, 6, 41) @@ -51,7 +49,7 @@ pub fn is_finality_proof_valid( pub fn is_next_committee_proof_valid( attested_header: &Header, - next_committee: &mut SyncCommittee, + next_committee: &SyncCommittee, next_committee_branch: &[B256], ) -> bool { is_proof_valid( @@ -65,7 +63,7 @@ pub fn is_next_committee_proof_valid( pub fn is_current_committee_proof_valid( attested_header: &Header, - current_committee: &mut SyncCommittee, + current_committee: &SyncCommittee, current_committee_branch: &[B256], ) -> bool { is_proof_valid( @@ -216,8 +214,8 @@ pub fn verify_generic_update( if update.finalized_header.is_some() && update.finality_branch.is_some() { let is_valid = is_finality_proof_valid( &update.attested_header, - &mut update.finalized_header.clone().unwrap(), - &update.finality_branch.clone().unwrap(), + update.finalized_header.as_ref().unwrap(), + update.finality_branch.as_ref().unwrap(), ); if !is_valid { @@ -228,8 +226,8 @@ pub fn verify_generic_update( if update.next_sync_committee.is_some() && update.next_sync_committee_branch.is_some() { let is_valid = is_next_committee_proof_valid( &update.attested_header, - &mut update.next_sync_committee.clone().unwrap(), - &update.next_sync_committee_branch.clone().unwrap(), + update.next_sync_committee.as_ref().unwrap(), + update.next_sync_committee_branch.as_ref().unwrap(), ); if !is_valid { @@ -316,13 +314,12 @@ pub fn expected_current_slot(now: SystemTime, genesis_time: u64) -> u64 { pub fn verify_sync_committee_signture( pks: &[PublicKey], attested_header: &Header, - signature: &SignatureBytes, + signature: &Signature, fork_data_root: B256, ) -> bool { - let pks: Vec<&PublicKey> = pks.iter().collect(); - let header_root = attested_header.clone().tree_hash_root(); + let header_root = attested_header.tree_hash_root(); let signing_root = compute_committee_sign_root(header_root, fork_data_root); - is_aggregate_valid(signature, signing_root.as_ref(), &pks) + signature.verify(signing_root.as_slice(), pks) } pub fn compute_committee_sign_root(header: B256, fork_data_root: B256) -> B256 { diff --git a/consensus-core/src/lib.rs b/consensus-core/src/lib.rs index c2a321d4..cd984650 100644 --- a/consensus-core/src/lib.rs +++ b/consensus-core/src/lib.rs @@ -2,5 +2,5 @@ pub mod errors; pub mod types; pub mod utils; -pub mod consensus_core; +mod consensus_core; pub use crate::consensus_core::*; diff --git a/consensus-core/src/types/bls.rs b/consensus-core/src/types/bls.rs new file mode 100644 index 00000000..3de2a3b7 --- /dev/null +++ b/consensus-core/src/types/bls.rs @@ -0,0 +1,40 @@ +use milagro_bls::{AggregateSignature, PublicKey as MilagroPK}; +use serde::{Deserialize, Serialize}; +use ssz_derive::{Decode, Encode}; +use tree_hash_derive::TreeHash; + +use super::bytes::ByteVector; + +#[derive(Debug, Clone, Default, Serialize, Deserialize, Encode, Decode, TreeHash)] +#[ssz(struct_behaviour = "transparent")] +#[serde(transparent)] +pub struct PublicKey { + inner: ByteVector, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, Encode, Decode, TreeHash)] +#[ssz(struct_behaviour = "transparent")] +#[serde(transparent)] +pub struct Signature { + inner: ByteVector, +} + +impl Signature { + pub fn verify(&self, msg: &[u8], pks: &[PublicKey]) -> bool { + if let Ok(agg) = AggregateSignature::from_bytes(&self.inner.inner) { + let pks_res = pks + .iter() + .map(|pk| MilagroPK::from_bytes(&pk.inner.inner)) + .collect::, _>>(); + + if let Ok(pks) = pks_res { + let pks = pks.iter().collect::>(); + agg.fast_aggregate_verify(msg, &pks) + } else { + false + } + } else { + false + } + } +} diff --git a/consensus-core/src/types/primitives.rs b/consensus-core/src/types/bytes.rs similarity index 100% rename from consensus-core/src/types/primitives.rs rename to consensus-core/src/types/bytes.rs diff --git a/consensus-core/src/types/mod.rs b/consensus-core/src/types/mod.rs index 9082e367..b3116bc7 100644 --- a/consensus-core/src/types/mod.rs +++ b/consensus-core/src/types/mod.rs @@ -6,16 +6,18 @@ use ssz_types::{serde_utils::quoted_u64_var_list, BitList, BitVector, FixedVecto use superstruct::superstruct; use tree_hash_derive::TreeHash; -use self::primitives::{ByteList, ByteVector}; +use self::{ + bls::{PublicKey, Signature}, + bytes::{ByteList, ByteVector}, +}; -pub mod primitives; +pub mod bls; +pub mod bytes; mod serde_utils; mod utils; pub type LogsBloom = ByteVector; -pub type BLSPubKey = ByteVector; pub type KZGCommitment = ByteVector; -pub type SignatureBytes = ByteVector; pub type Transaction = ByteList; #[derive(Debug, Default, Clone, Deserialize)] @@ -51,7 +53,7 @@ pub struct BeaconBlock { #[ssz(enum_behaviour = "transparent")] #[tree_hash(enum_behaviour = "transparent")] pub struct BeaconBlockBody { - randao_reveal: SignatureBytes, + randao_reveal: Signature, eth1_data: Eth1Data, graffiti: B256, proposer_slashings: VariableList, @@ -76,14 +78,14 @@ impl Default for BeaconBlockBody { #[derive(Default, Clone, Debug, Encode, TreeHash, Deserialize)] pub struct SignedBlsToExecutionChange { message: BlsToExecutionChange, - signature: SignatureBytes, + signature: Signature, } #[derive(Default, Clone, Debug, Encode, TreeHash, Deserialize)] pub struct BlsToExecutionChange { #[serde(with = "serde_utils::u64")] validator_index: u64, - from_bls_pubkey: BLSPubKey, + from_bls_pubkey: PublicKey, to_execution_address: Address, } @@ -154,7 +156,7 @@ pub struct ProposerSlashing { #[derive(Deserialize, Debug, Default, Encode, TreeHash, Clone)] struct SignedBeaconBlockHeader { message: BeaconBlockHeader, - signature: SignatureBytes, + signature: Signature, } #[derive(Deserialize, Debug, Default, Encode, TreeHash, Clone)] @@ -179,14 +181,14 @@ struct IndexedAttestation { #[serde(with = "quoted_u64_var_list")] attesting_indices: VariableList, data: AttestationData, - signature: SignatureBytes, + signature: Signature, } #[derive(Deserialize, Debug, Encode, TreeHash, Clone)] pub struct Attestation { aggregation_bits: BitList, data: AttestationData, - signature: SignatureBytes, + signature: Signature, } #[derive(Deserialize, Debug, Default, Encode, TreeHash, Clone)] @@ -210,7 +212,7 @@ struct Checkpoint { #[derive(Deserialize, Debug, Default, Encode, TreeHash, Clone)] pub struct SignedVoluntaryExit { message: VoluntaryExit, - signature: SignatureBytes, + signature: Signature, } #[derive(Deserialize, Debug, Default, Encode, TreeHash, Clone)] @@ -229,11 +231,11 @@ pub struct Deposit { #[derive(Deserialize, Default, Debug, Encode, TreeHash, Clone)] struct DepositData { - pubkey: BLSPubKey, + pubkey: PublicKey, withdrawal_credentials: B256, #[serde(with = "serde_utils::u64")] amount: u64, - signature: SignatureBytes, + signature: Signature, } #[derive(Deserialize, Debug, Default, Encode, TreeHash, Clone)] @@ -300,14 +302,14 @@ pub struct Header { #[derive(Debug, Clone, Default, Encode, TreeHash, Serialize, Deserialize)] pub struct SyncCommittee { - pub pubkeys: FixedVector, - pub aggregate_pubkey: BLSPubKey, + pub pubkeys: FixedVector, + pub aggregate_pubkey: PublicKey, } #[derive(Serialize, Deserialize, Debug, Clone, Default, Encode, TreeHash)] pub struct SyncAggregate { pub sync_committee_bits: BitVector, - pub sync_committee_signature: SignatureBytes, + pub sync_committee_signature: Signature, } pub struct GenericUpdate { diff --git a/consensus-core/src/utils.rs b/consensus-core/src/utils.rs index c214f0b3..5bf10001 100644 --- a/consensus-core/src/utils.rs +++ b/consensus-core/src/utils.rs @@ -1,11 +1,10 @@ use alloy::primitives::B256; -use milagro_bls::{AggregateSignature, PublicKey}; use sha2::{Digest, Sha256}; use ssz_types::FixedVector; use tree_hash::TreeHash; use tree_hash_derive::TreeHash; -use crate::types::{Header, SignatureBytes}; +use crate::types::Header; pub fn calc_sync_period(slot: u64) -> u64 { // 32 slots per epoch @@ -14,18 +13,9 @@ pub fn calc_sync_period(slot: u64) -> u64 { epoch / 256 } -pub fn is_aggregate_valid(sig_bytes: &SignatureBytes, msg: &[u8], pks: &[&PublicKey]) -> bool { - let sig_res = AggregateSignature::from_bytes(&sig_bytes.inner); - - match sig_res { - Ok(sig) => sig.fast_aggregate_verify(msg, pks), - Err(_) => false, - } -} - pub fn is_proof_valid( attested_header: &Header, - leaf_object: &mut L, + leaf_object: &L, branch: &[B256], depth: usize, index: usize, diff --git a/consensus/src/consensus.rs b/consensus/src/consensus.rs index 2d9cc94c..efdcaef6 100644 --- a/consensus/src/consensus.rs +++ b/consensus/src/consensus.rs @@ -356,7 +356,7 @@ impl Inner { } pub async fn bootstrap(&mut self, checkpoint: B256) -> Result<()> { - let mut bootstrap = self + let bootstrap = self .rpc .get_bootstrap(checkpoint) .await @@ -374,7 +374,7 @@ impl Inner { let committee_valid = is_current_committee_proof_valid( &bootstrap.header, - &mut bootstrap.current_sync_committee, + &bootstrap.current_sync_committee, &bootstrap.current_sync_committee_branch, ); @@ -542,7 +542,10 @@ mod tests { }; use alloy::primitives::b256; use consensus_core::errors::ConsensusError; - use consensus_core::types::{BLSPubKey, Header, SignatureBytes}; + use consensus_core::types::{ + bls::{PublicKey, Signature}, + Header, + }; use config::{networks, Config}; use tokio::sync::{mpsc::channel, watch}; @@ -606,7 +609,7 @@ mod tests { .unwrap(); let mut update = updates[0].clone(); - update.next_sync_committee.pubkeys[0] = BLSPubKey::default(); + update.next_sync_committee.pubkeys[0] = PublicKey::default(); let err = client.verify_update(&update).err().unwrap(); assert_eq!( @@ -646,7 +649,7 @@ mod tests { .unwrap(); let mut update = updates[0].clone(); - update.sync_aggregate.sync_committee_signature = SignatureBytes::default(); + update.sync_aggregate.sync_committee_signature = Signature::default(); let err = client.verify_update(&update).err().unwrap(); assert_eq!( @@ -683,7 +686,7 @@ mod tests { let client = get_client(false, true).await; let mut update = client.rpc.get_finality_update().await.unwrap(); - update.sync_aggregate.sync_committee_signature = SignatureBytes::default(); + update.sync_aggregate.sync_committee_signature = Signature::default(); let err = client.verify_finality_update(&update).err().unwrap(); assert_eq!( @@ -705,7 +708,7 @@ mod tests { let client = get_client(false, true).await; let mut update = client.rpc.get_optimistic_update().await.unwrap(); - update.sync_aggregate.sync_committee_signature = SignatureBytes::default(); + update.sync_aggregate.sync_committee_signature = Signature::default(); let err = client.verify_optimistic_update(&update).err().unwrap(); assert_eq!(