From bd682be4ce5c0d4ac111853e4e0fbea2d07c5d82 Mon Sep 17 00:00:00 2001 From: dapplion <35266934+dapplion@users.noreply.github.com> Date: Mon, 26 Feb 2024 15:45:13 +0800 Subject: [PATCH] Handle key re-additions --- .../beacon_chain/src/block_verification.rs | 3 +- .../src/sync_committee_verification.rs | 8 ++-- beacon_node/beacon_chain/src/test_utils.rs | 13 +++++- beacon_node/http_api/src/lib.rs | 18 ++++---- beacon_node/operation_pool/src/lib.rs | 12 +++-- consensus/types/src/beacon_state.rs | 45 ++++++++++++++----- .../types/src/beacon_state/pubkey_cache.rs | 22 +++++---- lcli/src/transition_blocks.rs | 2 +- 8 files changed, 84 insertions(+), 39 deletions(-) diff --git a/beacon_node/beacon_chain/src/block_verification.rs b/beacon_node/beacon_chain/src/block_verification.rs index 9d2679f5d58..56b212db403 100644 --- a/beacon_node/beacon_chain/src/block_verification.rs +++ b/beacon_node/beacon_chain/src/block_verification.rs @@ -2083,7 +2083,8 @@ fn get_signature_verifier<'a, T: BeaconChainTypes>( let decompressor = move |pk_bytes| { // Map compressed pubkey to validator index. - let validator_index = state.get_validator_index_readonly(pk_bytes)?; + // TODO(lion): How to ensure the cache is updated here? + let validator_index = state.get_validator_index_readonly_unchecked(pk_bytes)?; // Map validator index to pubkey (respecting guard on unknown validators). get_pubkey(validator_index) }; diff --git a/beacon_node/beacon_chain/src/sync_committee_verification.rs b/beacon_node/beacon_chain/src/sync_committee_verification.rs index 4eec3848cb2..9b3c3603f75 100644 --- a/beacon_node/beacon_chain/src/sync_committee_verification.rs +++ b/beacon_node/beacon_chain/src/sync_committee_verification.rs @@ -383,11 +383,11 @@ impl VerifiedSyncContribution { participant_pubkeys .iter() .map(|pubkey| { - head_state - .get_validator_index_readonly(pubkey) - .ok_or(Error::UnknownValidatorPubkey(*pubkey)) + Ok(head_state + .get_validator_index_readonly(pubkey)? + .ok_or(Error::UnknownValidatorPubkey(*pubkey))?) }) - .collect::, _>>()? + .collect::, Error>>()? }; // Ensure that all signatures are valid. diff --git a/beacon_node/beacon_chain/src/test_utils.rs b/beacon_node/beacon_chain/src/test_utils.rs index fe8d722bb9f..c3bfd838726 100644 --- a/beacon_node/beacon_chain/src/test_utils.rs +++ b/beacon_node/beacon_chain/src/test_utils.rs @@ -736,7 +736,16 @@ where } pub fn get_current_state(&self) -> BeaconState { - self.chain.head_beacon_state_cloned() + self.chain + .head_snapshot() + .beacon_state + .clone_with(CloneConfig { + committee_caches: true, + pubkey_cache: true, + exit_cache: false, + tree_hash_cache: false, + progressive_balances_cache: false, + }) } pub fn get_timestamp_at_slot(&self) -> u64 { @@ -1193,6 +1202,7 @@ where .map(|(subcommittee_position, pubkey)| { let validator_index = state .get_validator_index_readonly(pubkey) + .expect("pubkey cache not updated") .expect("should find validator index"); let sync_message = SyncCommitteeMessage::new::( @@ -1403,6 +1413,7 @@ where .find_map(|pubkey| { let validator_index = state .get_validator_index_readonly(pubkey) + .expect("pubkey cache not updated") .expect("should find validator index"); let selection_proof = SyncSelectionProof::new::( diff --git a/beacon_node/http_api/src/lib.rs b/beacon_node/http_api/src/lib.rs index 38dd9f290a4..a6dcc8f3b83 100644 --- a/beacon_node/http_api/src/lib.rs +++ b/beacon_node/http_api/src/lib.rs @@ -769,9 +769,9 @@ pub fn serve( &chain, |state, execution_optimistic, finalized| { let index_opt = match &validator_id { - ValidatorId::PublicKey(pubkey) => { - state.get_validator_index_readonly(pubkey) - } + ValidatorId::PublicKey(pubkey) => state + .get_validator_index_readonly(pubkey) + .map_err(|e| warp_utils::reject::beacon_state_error(e))?, ValidatorId::Index(index) => Some(*index as usize), }; @@ -1040,14 +1040,14 @@ pub fn serve( .pubkeys .iter() .map(|pubkey| { - state - .get_validator_index_readonly(pubkey) + Ok(state + .get_validator_index_readonly(pubkey)? .map(|index| index as u64) .ok_or(BeaconChainError::ValidatorPubkeyUnknown( *pubkey, - )) + ))?) }) - .collect::, _>>() + .collect::, BeaconChainError>>() .map_err(warp_utils::reject::beacon_chain_error)?; Ok((validators, execution_optimistic, finalized)) @@ -3583,7 +3583,9 @@ pub fn serve( .filter_map(|register_data| { head_snapshot .beacon_state - .get_validator_index_readonly(®ister_data.message.pubkey) + .get_validator_index_readonly_unchecked( + ®ister_data.message.pubkey, + ) .and_then(|validator_index| { let validator = head_snapshot .beacon_state diff --git a/beacon_node/operation_pool/src/lib.rs b/beacon_node/operation_pool/src/lib.rs index 7e1ddb1fd2f..6f4ac374b2a 100644 --- a/beacon_node/operation_pool/src/lib.rs +++ b/beacon_node/operation_pool/src/lib.rs @@ -1529,7 +1529,8 @@ mod release_tests { let (harness, _) = sync_contribution_test_state::(1).await; let op_pool = OperationPool::::new(); - let state = harness.get_current_state(); + let mut state = harness.get_current_state(); + state.update_pubkey_cache().unwrap(); let block_root = *state .get_block_root(state.slot() - Slot::new(1)) @@ -1589,7 +1590,8 @@ mod release_tests { let (harness, _) = sync_contribution_test_state::(1).await; let op_pool = OperationPool::::new(); - let state = harness.get_current_state(); + let mut state = harness.get_current_state(); + state.update_pubkey_cache().unwrap(); let block_root = *state .get_block_root(state.slot() - Slot::new(1)) .ok() @@ -1626,7 +1628,8 @@ mod release_tests { let (harness, _) = sync_contribution_test_state::(1).await; let op_pool = OperationPool::::new(); - let state = harness.get_current_state(); + let mut state = harness.get_current_state(); + state.update_pubkey_cache().unwrap(); let block_root = *state .get_block_root(state.slot() - Slot::new(1)) .ok() @@ -1706,7 +1709,8 @@ mod release_tests { let (harness, _) = sync_contribution_test_state::(1).await; let op_pool = OperationPool::::new(); - let state = harness.get_current_state(); + let mut state = harness.get_current_state(); + state.update_pubkey_cache().unwrap(); let block_root = *state .get_block_root(state.slot() - Slot::new(1)) .ok() diff --git a/consensus/types/src/beacon_state.rs b/consensus/types/src/beacon_state.rs index 0f542f2d2a0..73b243ee7d2 100644 --- a/consensus/types/src/beacon_state.rs +++ b/consensus/types/src/beacon_state.rs @@ -85,7 +85,7 @@ pub enum Error { index: CommitteeIndex, }, ZeroSlotsPerEpoch, - PubkeyCacheInconsistent, + PubkeyCacheMissingKey(PublicKeyBytes), PubkeyCacheIncomplete { cache_len: usize, registry_len: usize, @@ -497,14 +497,40 @@ impl BeaconState { /// otherwise returns `None`. pub fn get_validator_index(&mut self, pubkey: &PublicKeyBytes) -> Result, Error> { self.update_pubkey_cache()?; - Ok(self.get_validator_index_readonly(pubkey)) + Ok(self.get_validator_index_readonly_unchecked(pubkey)) } /// pubkey_cache may contain pubkeys from future blocks not yet known to this state. Ignore /// pubkeys that resolve to indexes beyond the current validators list. - pub fn get_validator_index_readonly(&self, pubkey: &PublicKeyBytes) -> Option { + pub fn get_validator_index_readonly( + &self, + pubkey: &PublicKeyBytes, + ) -> Result, Error> { + if self.pubkey_cache().len() < self.validators().len() { + return Err(Error::PubkeyCacheIncomplete { + cache_len: self.pubkey_cache().len(), + registry_len: self.validators().len(), + }); + } + + Ok(self.get_validator_index_readonly_unchecked(pubkey)) + } + + pub fn assert_pubkey_cache(&self) -> Result<(), Error> { + if self.pubkey_cache().len() < self.validators().len() { + return Err(Error::PubkeyCacheIncomplete { + cache_len: self.pubkey_cache().len(), + registry_len: self.validators().len(), + }); + } + Ok(()) + } + + /// pubkey_cache may contain pubkeys from future blocks not yet known to this state. Ignore + /// pubkeys that resolve to indexes beyond the current validators list. + pub fn get_validator_index_readonly_unchecked(&self, pubkey: &PublicKeyBytes) -> Option { self.pubkey_cache().get(pubkey).and_then(|index| { - if index >= self.validators().len() { + if index < self.validators().len() { Some(index) } else { None @@ -898,7 +924,7 @@ impl BeaconState { for pubkey in sync_committee.pubkeys.iter() { indices.push( self.get_validator_index(pubkey)? - .ok_or(Error::PubkeyCacheInconsistent)?, + .ok_or(Error::PubkeyCacheMissingKey(*pubkey))?, ) } Ok(indices) @@ -1689,8 +1715,7 @@ impl BeaconState { /// Updates the pubkey cache, if required. /// - /// Adds all `pubkeys` from the `validators` which are not already in the cache. Will - /// never re-add a pubkey. + /// Adds all `pubkeys` from the `validators`. May re-add the same pubkey pub fn update_pubkey_cache(&mut self) -> Result<(), Error> { let mut pubkey_cache = mem::take(self.pubkey_cache_mut()); for (i, validator) in self @@ -1699,10 +1724,8 @@ impl BeaconState { .enumerate() .skip(pubkey_cache.len()) { - let success = pubkey_cache.insert(validator.pubkey, i); - if !success { - return Err(Error::PubkeyCacheInconsistent); - } + // Ok to read pubkeys before EIP-6110 + pubkey_cache.insert(validator.pubkey, i); } *self.pubkey_cache_mut() = pubkey_cache; diff --git a/consensus/types/src/beacon_state/pubkey_cache.rs b/consensus/types/src/beacon_state/pubkey_cache.rs index 27b7b86a623..dfe3feee885 100644 --- a/consensus/types/src/beacon_state/pubkey_cache.rs +++ b/consensus/types/src/beacon_state/pubkey_cache.rs @@ -1,5 +1,6 @@ use crate::*; use rpds::HashTrieMapSync; +use std::cmp::Ordering; type ValidatorIndex = usize; @@ -24,15 +25,18 @@ impl PubkeyCache { /// The added index must equal the number of validators already added to the map. This ensures /// that an index is never skipped. pub fn insert(&mut self, pubkey: PublicKeyBytes, index: ValidatorIndex) -> bool { - if index == self.len { - self.map.insert_mut(pubkey, index); - self.len = self - .len - .checked_add(1) - .expect("map length cannot exceed usize"); - true - } else { - false + match index.cmp(&self.len) { + Ordering::Greater => false, + Ordering::Equal => { + self.map.insert_mut(pubkey, index); + self.len = self + .len + .checked_add(1) + .expect("map length cannot exceed usize"); + true + } + // Ignore inserts for already known keys + Ordering::Less => true, } } diff --git a/lcli/src/transition_blocks.rs b/lcli/src/transition_blocks.rs index 434e6c06118..6bd53e879aa 100644 --- a/lcli/src/transition_blocks.rs +++ b/lcli/src/transition_blocks.rs @@ -358,7 +358,7 @@ fn do_transition( let state = pre_state.clone(); let decompressor = move |pk_bytes| { // Map compressed pubkey to validator index. - let validator_index = state.get_validator_index_readonly(pk_bytes)?; + let validator_index = state.get_validator_index_readonly_unchecked(pk_bytes)?; // Map validator index to pubkey (respecting guard on unknown validators). get_pubkey(validator_index) };