feat: epoch stake syscalls (#141)

Author: buffalojoec

Cargo.lock

@@ -29,6 +29,7 @@ audit:
 build-test-programs:
 	@cargo build-sbf --manifest-path test-programs/cpi-target/Cargo.toml
 	@cargo build-sbf --manifest-path test-programs/custom-syscall/Cargo.toml
+	@cargo build-sbf --manifest-path test-programs/epoch-stake/Cargo.toml
 	@cargo build-sbf --manifest-path test-programs/primary/Cargo.toml
 
 # Pre-publish checks

Makefile

@@ -0,0 +1,31 @@
+use {solana_pubkey::Pubkey, std::collections::HashMap};
+
+/// A simple map of vote accounts to their epoch stake.
+///
+/// Developers can work with this map directly to configure stake for testing.
+/// The total epoch stake is calculated by summing all vote account stakes.
+pub type EpochStake = HashMap<Pubkey, u64>;
+
+/// Create an `EpochStake` instance with a few mocked-out vote accounts to
+/// achieve the provided total stake.
+pub fn create_mock_epoch_stake(target_total: u64) -> EpochStake {
+    let mut epoch_stake = HashMap::new();
+
+    if target_total == 0 {
+        return epoch_stake;
+    }
+
+    let num_accounts = target_total.div_ceil(1_000_000_000);
+
+    let base_stake = target_total / num_accounts;
+    let remainder = target_total % num_accounts;
+
+    std::iter::repeat(base_stake)
+        .take(num_accounts as usize - 1)
+        .chain(std::iter::once(base_stake + remainder))
+        .for_each(|stake| {
+            epoch_stake.insert(Pubkey::new_unique(), stake);
+        });
+
+    epoch_stake
+}

harness/src/epoch_stake.rs

@@ -441,6 +441,7 @@
 
 pub mod account_store;
 mod compile_accounts;
+pub mod epoch_stake;
 pub mod file;
 #[cfg(any(feature = "fuzz", feature = "fuzz-fd"))]
 pub mod fuzz;
@@ -453,8 +454,8 @@ pub use mollusk_svm_result as result;
 use mollusk_svm_result::Compare;
 use {
     crate::{
-        account_store::AccountStore, compile_accounts::CompiledAccounts, program::ProgramCache,
-        sysvar::Sysvars,
+        account_store::AccountStore, compile_accounts::CompiledAccounts, epoch_stake::EpochStake,
+        program::ProgramCache, sysvar::Sysvars,
     },
     agave_feature_set::FeatureSet,
     mollusk_svm_error::error::{MolluskError, MolluskPanic},
@@ -482,6 +483,7 @@ pub(crate) const DEFAULT_LOADER_KEY: Pubkey = solana_sdk_ids::bpf_loader_upgrade
 pub struct Mollusk {
     pub config: Config,
     pub compute_budget: ComputeBudget,
+    pub epoch_stake: EpochStake,
     pub feature_set: FeatureSet,
     pub logger: Option<Rc<RefCell<LogCollector>>>,
     pub program_cache: ProgramCache,
@@ -521,6 +523,7 @@ impl Default for Mollusk {
         Self {
             config: Config::default(),
             compute_budget,
+            epoch_stake: EpochStake::default(),
             feature_set,
             logger: None,
             program_cache,
@@ -539,9 +542,18 @@ impl CheckContext for Mollusk {
 
 struct MolluskInvokeContextCallback<'a> {
     feature_set: &'a FeatureSet,
+    epoch_stake: &'a EpochStake,
 }
 
 impl InvokeContextCallback for MolluskInvokeContextCallback<'_> {
+    fn get_epoch_stake(&self) -> u64 {
+        self.epoch_stake.values().sum()
+    }
+
+    fn get_epoch_stake_for_vote_account(&self, vote_address: &Pubkey) -> u64 {
+        self.epoch_stake.get(vote_address).copied().unwrap_or(0)
+    }
+
     fn is_precompile(&self, program_id: &Pubkey) -> bool {
         agave_precompiles::is_precompile(program_id, |feature_id| {
             self.feature_set.is_active(feature_id)
@@ -647,6 +659,7 @@ impl Mollusk {
         let invoke_result = {
             let mut program_cache = self.program_cache.cache();
             let callback = MolluskInvokeContextCallback {
+                epoch_stake: &self.epoch_stake,
                 feature_set: &self.feature_set,
             };
             let runtime_features = self.feature_set.runtime_features();

harness/src/lib.rs

@@ -44,7 +44,7 @@ fn test_custom_syscall() {
             .unwrap();
         mollusk.add_program(
             &program_id,
-            "custom_syscall_program",
+            "test_program_custom_syscall",
             &mollusk_svm::program::loader_keys::LOADER_V3,
         );
         mollusk

harness/tests/custom_syscall.rs

@@ -0,0 +1,50 @@
+use {
+    mollusk_svm::{result::Check, Mollusk},
+    solana_account::Account,
+    solana_instruction::{AccountMeta, Instruction},
+    solana_pubkey::Pubkey,
+};
+
+#[test]
+fn test_epoch_stake() {
+    std::env::set_var("SBF_OUT_DIR", "../target/deploy");
+
+    let program_id = Pubkey::new_unique();
+    let mut mollusk = Mollusk::new(&program_id, "test_program_epoch_stake");
+
+    let key = Pubkey::new_unique();
+
+    let mut total_stake: u64 = 0;
+
+    for i in 1..=3 {
+        let stake = 1_000_000_000_000_000 * i;
+        total_stake += stake;
+
+        let vote_address = Pubkey::new_unique();
+
+        mollusk.epoch_stake.insert(vote_address, stake);
+        assert_eq!(total_stake, mollusk.epoch_stake.values().sum::<u64>());
+
+        let instruction = Instruction::new_with_bytes(
+            program_id,
+            &vote_address.to_bytes(),
+            vec![AccountMeta::new(key, false)],
+        );
+
+        mollusk.process_and_validate_instruction(
+            &instruction,
+            &[(key, Account::new(1_000, 16, &program_id))],
+            &[
+                Check::success(),
+                Check::account(&key)
+                    .data(&{
+                        let mut data = vec![0; 16];
+                        data[0..8].copy_from_slice(&total_stake.to_le_bytes());
+                        data[8..16].copy_from_slice(&stake.to_le_bytes());
+                        data
+                    })
+                    .build(),
+            ],
+        );
+    }
+}

harness/tests/epoch_stake.rs

@@ -1,5 +1,5 @@
 [package]
-name = "custom-syscall-program"
+name = "test-program-custom-syscall"
 version = "0.1.0"
 edition = "2021"
 
@@ -18,4 +18,3 @@ check-cfg = [
     'cfg(feature, values("custom-heap", "custom-panic"))',
     'cfg(target_os, values("solana"))',
 ]
-

test-programs/custom-syscall/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "test-program-epoch-stake"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+solana-account-info = { workspace = true }
+solana-program-entrypoint = { workspace = true }
+solana-program-error = { workspace = true }
+solana-pubkey = { workspace = true }
+
+[lib]
+crate-type = ["cdylib", "lib"]
+
+[lints.rust.unexpected_cfgs]
+level = "warn"
+check-cfg = [
+    'cfg(feature, values("custom-heap", "custom-panic"))',
+    'cfg(target_os, values("solana"))',
+]

test-programs/epoch-stake/Cargo.toml

@@ -0,0 +1,34 @@
+#![cfg(target_os = "solana")]
+
+use {solana_account_info::AccountInfo, solana_program_error::ProgramError, solana_pubkey::Pubkey};
+
+extern "C" {
+    fn sol_get_epoch_stake(vote_address: *const u8) -> u64;
+}
+
+unsafe fn get_epoch_total_stake() -> u64 {
+    sol_get_epoch_stake(std::ptr::null::<Pubkey>() as *const u8)
+}
+
+unsafe fn get_epoch_stake_for_vote_account(vote_address: &Pubkey) -> u64 {
+    sol_get_epoch_stake(vote_address as *const _ as *const u8)
+}
+
+solana_program_entrypoint::entrypoint!(process_instruction);
+
+fn process_instruction(
+    _program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    input: &[u8],
+) -> Result<(), ProgramError> {
+    let vote_address = Pubkey::new_from_array(input.try_into().unwrap());
+
+    let total_stake = unsafe { get_epoch_total_stake() };
+    let vote_account_stake = unsafe { get_epoch_stake_for_vote_account(&vote_address) };
+
+    let mut data = accounts[0].try_borrow_mut_data()?;
+    data[0..8].copy_from_slice(&total_stake.to_le_bytes());
+    data[8..16].copy_from_slice(&vote_account_stake.to_le_bytes());
+
+    Ok(())
+}