From 00b03ea438f4c5d6d3d338e41640cd31ead70fed Mon Sep 17 00:00:00 2001 From: Joe Caulfield Date: Thu, 6 Jun 2024 15:10:09 -0500 Subject: [PATCH] refactor for total stake --- program-runtime/src/invoke_context.rs | 23 ++- programs/bpf_loader/src/syscalls/mod.rs | 146 +++++++++++++----- runtime/src/bank.rs | 13 +- .../bank/builtins/core_bpf_migration/mod.rs | 3 +- svm/src/account_loader.rs | 6 +- svm/src/message_processor.rs | 21 ++- svm/src/program_loader.rs | 6 +- svm/src/transaction_processing_callback.rs | 4 +- svm/src/transaction_processor.rs | 9 +- svm/tests/conformance.rs | 3 +- svm/tests/mock_bank.rs | 6 +- 11 files changed, 178 insertions(+), 62 deletions(-) diff --git a/program-runtime/src/invoke_context.rs b/program-runtime/src/invoke_context.rs index 28c4d82126b480..293214a870f13a 100644 --- a/program-runtime/src/invoke_context.rs +++ b/program-runtime/src/invoke_context.rs @@ -147,24 +147,27 @@ impl BpfAllocator { pub struct EnvironmentConfig<'a> { pub blockhash: Hash, + epoch_total_stake: Option, + epoch_vote_accounts: Option<&'a VoteAccountsHashMap>, pub feature_set: Arc, pub lamports_per_signature: u64, - vote_accounts: Option<&'a VoteAccountsHashMap>, sysvar_cache: &'a SysvarCache, } impl<'a> EnvironmentConfig<'a> { pub fn new( blockhash: Hash, + epoch_total_stake: Option, + epoch_vote_accounts: Option<&'a VoteAccountsHashMap>, feature_set: Arc, lamports_per_signature: u64, - vote_accounts: Option<&'a VoteAccountsHashMap>, sysvar_cache: &'a SysvarCache, ) -> Self { Self { blockhash, + epoch_total_stake, + epoch_vote_accounts, feature_set, lamports_per_signature, - vote_accounts, sysvar_cache, } } @@ -618,9 +621,14 @@ impl<'a> InvokeContext<'a> { self.environment_config.sysvar_cache } - /// Get cached vote accounts. - pub fn get_vote_accounts(&self) -> Option<&VoteAccountsHashMap> { - self.environment_config.vote_accounts + /// Get cached epoch total stake. + pub fn get_epoch_total_stake(&self) -> Option { + self.environment_config.epoch_total_stake + } + + /// Get cached epoch vote accounts. + pub fn get_epoch_vote_accounts(&self) -> Option<&VoteAccountsHashMap> { + self.environment_config.epoch_vote_accounts } // Should alignment be enforced during user pointer translation @@ -719,9 +727,10 @@ macro_rules! with_mock_invoke_context { }); let environment_config = EnvironmentConfig::new( Hash::default(), + None, + None, Arc::new(FeatureSet::all_enabled()), 0, - /* vote_accounts */ None, &sysvar_cache, ); let program_cache_for_tx_batch = ProgramCacheForTxBatch::default(); diff --git a/programs/bpf_loader/src/syscalls/mod.rs b/programs/bpf_loader/src/syscalls/mod.rs index d62decf68d4d09..3d530cccc58848 100644 --- a/programs/bpf_loader/src/syscalls/mod.rs +++ b/programs/bpf_loader/src/syscalls/mod.rs @@ -2025,51 +2025,75 @@ declare_builtin_function!( SyscallGetEpochStake, fn rust( invoke_context: &mut InvokeContext, - vote_address: u64, + var_addr: u64, _arg2: u64, _arg3: u64, _arg4: u64, _arg5: u64, memory_mapping: &mut MemoryMapping, ) -> Result { - // Compute units, as specified by SIMD-0133. - // cu = syscall_base_cost - // + floor(32/cpi_bytes_per_unit) - // + mem_op_base_cost let compute_budget = invoke_context.get_compute_budget(); - let compute_units = compute_budget - .syscall_base_cost - .saturating_add( - (PUBKEY_BYTES as u64) - .checked_div(compute_budget.cpi_bytes_per_unit) - .unwrap_or(u64::MAX), + + if var_addr == 0 { + // As specified by SIMD-0133: If `var_addr` is a null pointer: + // + // Compute units: + // + // ``` + // syscall_base + // ``` + let compute_units = compute_budget.syscall_base_cost; + consume_compute_meter(invoke_context, compute_units)?; + // + // Control flow: + // + // - The syscall aborts the virtual machine if: + // - Compute budget is exceeded. + // - Otherwise, the syscall returns a `u64` integer representing the total active + // stake on the cluster for the current epoch. + Ok(invoke_context.get_epoch_total_stake().unwrap_or(0)) + } else { + // As specified by SIMD-0133: If `var_addr` is _not_ a null pointer: + // + // Compute units: + // + // ``` + // syscall_base + floor(PUBKEY_BYTES/cpi_bytes_per_unit) + mem_op_base + // ``` + let compute_units = compute_budget + .syscall_base_cost + .saturating_add( + (PUBKEY_BYTES as u64) + .checked_div(compute_budget.cpi_bytes_per_unit) + .unwrap_or(u64::MAX), + ) + .saturating_add(compute_budget.mem_op_base_cost); + consume_compute_meter(invoke_context, compute_units)?; + // + // Control flow: + // + // - The syscall aborts the virtual machine if: + // - Not all bytes in VM memory range `[vote_addr, vote_addr + 32)` are + // readable. + // - Compute budget is exceeded. + // - Otherwise, the syscall returns a `u64` integer representing the total active + // stake delegated to the vote account at the provided address. + // If the provided vote address corresponds to an account that is not a vote + // account or does not exist, the syscall will return `0` for active stake. + let check_aligned = invoke_context.get_check_aligned(); + let vote_address = translate_type::(memory_mapping, var_addr, check_aligned)?; + + Ok( + if let Some(vote_accounts) = invoke_context.get_epoch_vote_accounts() { + vote_accounts + .get(vote_address) + .map(|(stake, _)| *stake) + .unwrap_or(0) + } else { + 0 + }, ) - .saturating_add(compute_budget.mem_op_base_cost); - - consume_compute_meter(invoke_context, compute_units)?; - - // Control flow, as specified by SIMD-0133. - // * The syscall aborts the virtual machine if not all bytes in VM - // memory range `[vote_addr, vote_addr + 32)` are readable. - // * Otherwise, the syscall returns a `u64` integer representing the - // total active stake delegated to the vote account at the provided - // address. - // * If the provided vote address corresponds to an account that is - // not a vote account or does not exist, the syscall will return - // `0` for active stake. - let check_aligned = invoke_context.get_check_aligned(); - let vote_address = translate_type::(memory_mapping, vote_address, check_aligned)?; - - Ok( - if let Some(vote_accounts) = invoke_context.get_vote_accounts() { - vote_accounts - .get(vote_address) - .map(|(stake, _)| *stake) - .unwrap_or(0) - } else { - 0 - }, - ) + } } ); @@ -4745,7 +4769,50 @@ mod tests { } #[test] - fn test_syscall_get_epoch_stake() { + fn test_syscall_get_epoch_stake_total_stake() { + let config = Config::default(); + let mut compute_budget = ComputeBudget::default(); + let sysvar_cache = Arc::::default(); + + let expected_total_stake = 200_000_000_000_000u64; + // Compute units, as specified by SIMD-0133. + // cu = syscall_base_cost + let expected_cus = compute_budget.syscall_base_cost; + + // Set the compute budget to the expected CUs to ensure the syscall + // doesn't exceed the expected usage. + compute_budget.compute_unit_limit = expected_cus; + + with_mock_invoke_context!(invoke_context, transaction_context, vec![]); + invoke_context.environment_config = EnvironmentConfig::new( + Hash::default(), + Some(expected_total_stake), + None, // Vote accounts are not needed for this test. + Arc::::default(), + 0, + &sysvar_cache, + ); + + let null_pointer_var = std::ptr::null::() as u64; + + let mut memory_mapping = MemoryMapping::new(vec![], &config, &SBPFVersion::V2).unwrap(); + + let result = SyscallGetEpochStake::rust( + &mut invoke_context, + null_pointer_var, + 0, + 0, + 0, + 0, + &mut memory_mapping, + ) + .unwrap(); + + assert_eq!(result, expected_total_stake); + } + + #[test] + fn test_syscall_get_epoch_stake_vote_account_stake() { let config = Config::default(); let mut compute_budget = ComputeBudget::default(); let sysvar_cache = Arc::::default(); @@ -4781,9 +4848,10 @@ mod tests { with_mock_invoke_context!(invoke_context, transaction_context, vec![]); invoke_context.environment_config = EnvironmentConfig::new( Hash::default(), + None, // Total stake is not needed for this test. + Some(&vote_accounts_map), Arc::::default(), 0, - Some(&vote_accounts_map), &sysvar_cache, ); diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index 3be1fa4503f31a..e39bc84948dd9b 100644 --- a/runtime/src/bank.rs +++ b/runtime/src/bank.rs @@ -6155,6 +6155,13 @@ impl Bank { Some(self.epoch_stakes.get(&epoch)?.stakes().staked_nodes()) } + /// Get the total epoch stake for the given epoch. + pub fn epoch_total_stake(&self, epoch: Epoch) -> Option { + self.epoch_stakes + .get(&epoch) + .map(|epoch_stakes| epoch_stakes.total_stake()) + } + /// vote accounts for the specific epoch along with the stake /// attributed to each account pub fn epoch_vote_accounts(&self, epoch: Epoch) -> Option<&VoteAccountsHashMap> { @@ -6838,7 +6845,11 @@ impl TransactionProcessingCallback for Bank { self.feature_set.clone() } - fn get_vote_accounts(&self) -> Option<&VoteAccountsHashMap> { + fn get_epoch_total_stake(&self) -> Option { + self.epoch_total_stake(self.epoch()) + } + + fn get_epoch_vote_accounts(&self) -> Option<&VoteAccountsHashMap> { self.epoch_vote_accounts(self.epoch()) } diff --git a/runtime/src/bank/builtins/core_bpf_migration/mod.rs b/runtime/src/bank/builtins/core_bpf_migration/mod.rs index b174f2e1531ded..69f1fbdbd50c70 100644 --- a/runtime/src/bank/builtins/core_bpf_migration/mod.rs +++ b/runtime/src/bank/builtins/core_bpf_migration/mod.rs @@ -205,9 +205,10 @@ impl Bank { &program_cache_for_tx_batch, EnvironmentConfig::new( Hash::default(), + None, + None, self.feature_set.clone(), 0, - None, &sysvar_cache, ), None, diff --git a/svm/src/account_loader.rs b/svm/src/account_loader.rs index 1a2154122a0965..a6af292e94c88d 100644 --- a/svm/src/account_loader.rs +++ b/svm/src/account_loader.rs @@ -545,7 +545,11 @@ mod tests { self.feature_set.clone() } - fn get_vote_accounts(&self) -> Option<&VoteAccountsHashMap> { + fn get_epoch_total_stake(&self) -> Option { + None + } + + fn get_epoch_vote_accounts(&self) -> Option<&VoteAccountsHashMap> { None } } diff --git a/svm/src/message_processor.rs b/svm/src/message_processor.rs index 59f6119aa2e1b2..0505a730bd08d8 100644 --- a/svm/src/message_processor.rs +++ b/svm/src/message_processor.rs @@ -275,9 +275,10 @@ mod tests { let mut programs_modified_by_tx = ProgramCacheForTxBatch::default(); let environment_config = EnvironmentConfig::new( Hash::default(), + None, + None, Arc::new(FeatureSet::all_enabled()), 0, - None, &sysvar_cache, ); let mut invoke_context = InvokeContext::new( @@ -330,9 +331,10 @@ mod tests { let mut programs_modified_by_tx = ProgramCacheForTxBatch::default(); let environment_config = EnvironmentConfig::new( Hash::default(), + None, + None, Arc::new(FeatureSet::all_enabled()), 0, - None, &sysvar_cache, ); let mut invoke_context = InvokeContext::new( @@ -375,9 +377,10 @@ mod tests { let mut programs_modified_by_tx = ProgramCacheForTxBatch::default(); let environment_config = EnvironmentConfig::new( Hash::default(), + None, + None, Arc::new(FeatureSet::all_enabled()), 0, - None, &sysvar_cache, ); let mut invoke_context = InvokeContext::new( @@ -511,9 +514,10 @@ mod tests { let mut programs_modified_by_tx = ProgramCacheForTxBatch::default(); let environment_config = EnvironmentConfig::new( Hash::default(), + None, + None, Arc::new(FeatureSet::all_enabled()), 0, - None, &sysvar_cache, ); let mut invoke_context = InvokeContext::new( @@ -551,9 +555,10 @@ mod tests { let mut programs_modified_by_tx = ProgramCacheForTxBatch::default(); let environment_config = EnvironmentConfig::new( Hash::default(), + None, + None, Arc::new(FeatureSet::all_enabled()), 0, - None, &sysvar_cache, ); let mut invoke_context = InvokeContext::new( @@ -588,9 +593,10 @@ mod tests { let mut programs_modified_by_tx = ProgramCacheForTxBatch::default(); let environment_config = EnvironmentConfig::new( Hash::default(), + None, + None, Arc::new(FeatureSet::all_enabled()), 0, - None, &sysvar_cache, ); let mut invoke_context = InvokeContext::new( @@ -686,9 +692,10 @@ mod tests { let mut programs_modified_by_tx = ProgramCacheForTxBatch::default(); let environment_config = EnvironmentConfig::new( Hash::default(), + None, + None, Arc::new(FeatureSet::all_enabled()), 0, - None, &sysvar_cache, ); let mut invoke_context = InvokeContext::new( diff --git a/svm/src/program_loader.rs b/svm/src/program_loader.rs index 3d8b5e77eceee6..3b8f35923b5055 100644 --- a/svm/src/program_loader.rs +++ b/svm/src/program_loader.rs @@ -286,7 +286,11 @@ mod tests { self.feature_set.clone() } - fn get_vote_accounts(&self) -> Option<&VoteAccountsHashMap> { + fn get_epoch_total_stake(&self) -> Option { + None + } + + fn get_epoch_vote_accounts(&self) -> Option<&VoteAccountsHashMap> { None } diff --git a/svm/src/transaction_processing_callback.rs b/svm/src/transaction_processing_callback.rs index 83bb749ee15c3a..f92dc6e2299b39 100644 --- a/svm/src/transaction_processing_callback.rs +++ b/svm/src/transaction_processing_callback.rs @@ -20,7 +20,9 @@ pub trait TransactionProcessingCallback { fn get_feature_set(&self) -> Arc; - fn get_vote_accounts(&self) -> Option<&VoteAccountsHashMap>; + fn get_epoch_total_stake(&self) -> Option; + + fn get_epoch_vote_accounts(&self) -> Option<&VoteAccountsHashMap>; fn get_program_match_criteria(&self, _program: &Pubkey) -> ProgramCacheMatchCriteria { ProgramCacheMatchCriteria::NoCriteria diff --git a/svm/src/transaction_processor.rs b/svm/src/transaction_processor.rs index f980c2ad5d6113..a439ca8d9f0b40 100644 --- a/svm/src/transaction_processor.rs +++ b/svm/src/transaction_processor.rs @@ -626,9 +626,10 @@ impl TransactionBatchProcessor { program_cache_for_tx_batch, EnvironmentConfig::new( blockhash, + callback.get_epoch_total_stake(), + callback.get_epoch_vote_accounts(), callback.get_feature_set(), lamports_per_signature, - callback.get_vote_accounts(), sysvar_cache, ), log_collector.clone(), @@ -923,7 +924,11 @@ mod tests { self.feature_set.clone() } - fn get_vote_accounts(&self) -> Option<&VoteAccountsHashMap> { + fn get_epoch_total_stake(&self) -> Option { + None + } + + fn get_epoch_vote_accounts(&self) -> Option<&VoteAccountsHashMap> { None } diff --git a/svm/tests/conformance.rs b/svm/tests/conformance.rs index 65e4fc03ff7ca2..e1c28b75d1c87f 100644 --- a/svm/tests/conformance.rs +++ b/svm/tests/conformance.rs @@ -465,9 +465,10 @@ fn execute_fixture_as_instr( let sysvar_cache = &batch_processor.sysvar_cache.read().unwrap(); let env_config = EnvironmentConfig::new( mock_bank.blockhash, + None, + None, mock_bank.feature_set.clone(), mock_bank.lamports_per_sginature, - None, sysvar_cache, ); diff --git a/svm/tests/mock_bank.rs b/svm/tests/mock_bank.rs index f8fd89d1453117..7ee17547577957 100644 --- a/svm/tests/mock_bank.rs +++ b/svm/tests/mock_bank.rs @@ -70,7 +70,11 @@ impl TransactionProcessingCallback for MockBankCallback { self.feature_set.clone() } - fn get_vote_accounts(&self) -> Option<&VoteAccountsHashMap> { + fn get_epoch_total_stake(&self) -> Option { + None + } + + fn get_epoch_vote_accounts(&self) -> Option<&VoteAccountsHashMap> { None }