solana_program_test/
lib.rs

1#![cfg_attr(
2    not(feature = "agave-unstable-api"),
3    deprecated(
4        since = "3.1.0",
5        note = "This crate has been marked for formal inclusion in the Agave Unstable API. From \
6                v4.0.0 onward, the `agave-unstable-api` crate feature must be specified to \
7                acknowledge use of an interface that may break without warning."
8    )
9)]
10//! The solana-program-test provides a BanksClient-based test framework SBF programs
11#![allow(clippy::arithmetic_side_effects)]
12
13// Export tokio for test clients
14pub use tokio;
15use {
16    agave_feature_set::{
17        increase_cpi_account_info_limit, raise_cpi_nesting_limit_to_8, FEATURE_NAMES,
18    },
19    async_trait::async_trait,
20    base64::{prelude::BASE64_STANDARD, Engine},
21    chrono_humanize::{Accuracy, HumanTime, Tense},
22    log::*,
23    solana_account::{
24        create_account_shared_data_for_test, state_traits::StateMut, Account, AccountSharedData,
25        ReadableAccount,
26    },
27    solana_account_info::AccountInfo,
28    solana_accounts_db::accounts_db::ACCOUNTS_DB_CONFIG_FOR_TESTING,
29    solana_banks_client::start_client,
30    solana_banks_server::banks_server::start_local_server,
31    solana_clock::{Epoch, Slot},
32    solana_cluster_type::ClusterType,
33    solana_compute_budget::compute_budget::ComputeBudget,
34    solana_fee_calculator::{FeeRateGovernor, DEFAULT_TARGET_LAMPORTS_PER_SIGNATURE},
35    solana_genesis_config::GenesisConfig,
36    solana_hash::Hash,
37    solana_instruction::{
38        error::{InstructionError, UNSUPPORTED_SYSVAR},
39        Instruction,
40    },
41    solana_keypair::Keypair,
42    solana_native_token::LAMPORTS_PER_SOL,
43    solana_poh_config::PohConfig,
44    solana_program_binaries as programs,
45    solana_program_entrypoint::{deserialize, SUCCESS},
46    solana_program_error::{ProgramError, ProgramResult},
47    solana_program_runtime::{
48        invoke_context::BuiltinFunctionWithContext, loaded_programs::ProgramCacheEntry,
49        serialization::serialize_parameters, stable_log,
50    },
51    solana_pubkey::Pubkey,
52    solana_rent::Rent,
53    solana_runtime::{
54        bank::Bank,
55        bank_forks::BankForks,
56        commitment::BlockCommitmentCache,
57        genesis_utils::{create_genesis_config_with_leader_ex, GenesisConfigInfo},
58        runtime_config::RuntimeConfig,
59    },
60    solana_signer::Signer,
61    solana_svm_log_collector::ic_msg,
62    solana_svm_timings::ExecuteTimings,
63    solana_sysvar::SysvarSerialize,
64    solana_sysvar_id::SysvarId,
65    solana_vote_program::vote_state::{VoteStateV4, VoteStateVersions},
66    std::{
67        cell::RefCell,
68        collections::{HashMap, HashSet},
69        convert::TryFrom,
70        fs::File,
71        io::{self, Read},
72        mem::transmute,
73        panic::AssertUnwindSafe,
74        path::{Path, PathBuf},
75        sync::{
76            atomic::{AtomicBool, Ordering},
77            Arc, RwLock,
78        },
79        time::{Duration, Instant},
80    },
81    thiserror::Error,
82    tokio::task::JoinHandle,
83};
84// Export types so test clients can limit their solana crate dependencies
85pub use {
86    solana_banks_client::{BanksClient, BanksClientError},
87    solana_banks_interface::BanksTransactionResultWithMetadata,
88    solana_program_runtime::invoke_context::InvokeContext,
89    solana_sbpf::{
90        error::EbpfError,
91        vm::{get_runtime_environment_key, EbpfVm},
92    },
93    solana_transaction_context::IndexOfAccount,
94};
95
96/// Errors from the program test environment
97#[derive(Error, Debug, PartialEq, Eq)]
98pub enum ProgramTestError {
99    /// The chosen warp slot is not in the future, so warp is not performed
100    #[error("Warp slot not in the future")]
101    InvalidWarpSlot,
102}
103
104thread_local! {
105    static INVOKE_CONTEXT: RefCell<Option<usize>> = const { RefCell::new(None) };
106}
107fn set_invoke_context(new: &mut InvokeContext) {
108    INVOKE_CONTEXT.with(|invoke_context| unsafe {
109        invoke_context.replace(Some(transmute::<&mut InvokeContext, usize>(new)))
110    });
111}
112fn get_invoke_context<'a, 'b>() -> &'a mut InvokeContext<'b, 'b> {
113    let ptr = INVOKE_CONTEXT.with(|invoke_context| match *invoke_context.borrow() {
114        Some(val) => val,
115        None => panic!("Invoke context not set!"),
116    });
117    unsafe { transmute::<usize, &mut InvokeContext>(ptr) }
118}
119
120pub fn invoke_builtin_function(
121    builtin_function: solana_program_entrypoint::ProcessInstruction,
122    invoke_context: &mut InvokeContext,
123) -> Result<u64, Box<dyn std::error::Error>> {
124    set_invoke_context(invoke_context);
125
126    let transaction_context = &invoke_context.transaction_context;
127    let instruction_context = transaction_context.get_current_instruction_context()?;
128    let instruction_account_indices = 0..instruction_context.get_number_of_instruction_accounts();
129
130    // mock builtin program must consume units
131    invoke_context.consume_checked(1)?;
132
133    let log_collector = invoke_context.get_log_collector();
134    let program_id = instruction_context.get_program_key()?;
135    stable_log::program_invoke(
136        &log_collector,
137        program_id,
138        invoke_context.get_stack_height(),
139    );
140
141    // Copy indices_in_instruction into a HashSet to ensure there are no duplicates
142    let deduplicated_indices: HashSet<IndexOfAccount> = instruction_account_indices.collect();
143
144    // Serialize entrypoint parameters with SBF ABI
145    let mask_out_rent_epoch_in_vm_serialization = invoke_context
146        .get_feature_set()
147        .mask_out_rent_epoch_in_vm_serialization;
148    let (mut parameter_bytes, _regions, _account_lengths, _instruction_data_offset) =
149        serialize_parameters(
150            &instruction_context,
151            false, // There is no VM so stricter_abi_and_runtime_constraints can not be implemented here
152            false, // There is no VM so account_data_direct_mapping can not be implemented here
153            mask_out_rent_epoch_in_vm_serialization,
154        )?;
155
156    // Deserialize data back into instruction params
157    let (program_id, account_infos, input) =
158        unsafe { deserialize(&mut parameter_bytes.as_slice_mut()[0] as *mut u8) };
159
160    // Execute the program
161    match std::panic::catch_unwind(AssertUnwindSafe(|| {
162        builtin_function(program_id, &account_infos, input)
163    })) {
164        Ok(program_result) => {
165            program_result.map_err(|program_error| {
166                let err = InstructionError::from(u64::from(program_error));
167                stable_log::program_failure(&log_collector, program_id, &err);
168                let err: Box<dyn std::error::Error> = Box::new(err);
169                err
170            })?;
171        }
172        Err(_panic_error) => {
173            let err = InstructionError::ProgramFailedToComplete;
174            stable_log::program_failure(&log_collector, program_id, &err);
175            let err: Box<dyn std::error::Error> = Box::new(err);
176            Err(err)?;
177        }
178    };
179
180    stable_log::program_success(&log_collector, program_id);
181
182    // Lookup table for AccountInfo
183    let account_info_map: HashMap<_, _> = account_infos.into_iter().map(|a| (a.key, a)).collect();
184
185    // Re-fetch the instruction context. The previous reference may have been
186    // invalidated due to the `set_invoke_context` in a CPI.
187    let transaction_context = &invoke_context.transaction_context;
188    let instruction_context = transaction_context.get_current_instruction_context()?;
189
190    // Commit AccountInfo changes back into KeyedAccounts
191    for i in deduplicated_indices.into_iter() {
192        let mut borrowed_account = instruction_context.try_borrow_instruction_account(i)?;
193        if borrowed_account.is_writable() {
194            if let Some(account_info) = account_info_map.get(borrowed_account.get_key()) {
195                if borrowed_account.get_lamports() != account_info.lamports() {
196                    borrowed_account.set_lamports(account_info.lamports())?;
197                }
198
199                if borrowed_account
200                    .can_data_be_resized(account_info.data_len())
201                    .is_ok()
202                {
203                    borrowed_account.set_data_from_slice(&account_info.data.borrow())?;
204                }
205                if borrowed_account.get_owner() != account_info.owner {
206                    borrowed_account.set_owner(account_info.owner.as_ref())?;
207                }
208            }
209        }
210    }
211
212    Ok(0)
213}
214
215/// Converts a `solana-program`-style entrypoint into the runtime's entrypoint style, for
216/// use with `ProgramTest::add_program`
217#[macro_export]
218macro_rules! processor {
219    ($builtin_function:expr) => {
220        Some(|vm, _arg0, _arg1, _arg2, _arg3, _arg4| {
221            let vm = unsafe {
222                &mut *((vm as *mut u64).offset(-($crate::get_runtime_environment_key() as isize))
223                    as *mut $crate::EbpfVm<$crate::InvokeContext>)
224            };
225            vm.program_result =
226                $crate::invoke_builtin_function($builtin_function, vm.context_object_pointer)
227                    .map_err(|err| $crate::EbpfError::SyscallError(err))
228                    .into();
229        })
230    };
231}
232
233fn get_sysvar<T: Default + SysvarSerialize + Sized + serde::de::DeserializeOwned + Clone>(
234    sysvar: Result<Arc<T>, InstructionError>,
235    var_addr: *mut u8,
236) -> u64 {
237    let invoke_context = get_invoke_context();
238    if invoke_context
239        .consume_checked(invoke_context.get_execution_cost().sysvar_base_cost + T::size_of() as u64)
240        .is_err()
241    {
242        panic!("Exceeded compute budget");
243    }
244
245    match sysvar {
246        Ok(sysvar_data) => unsafe {
247            *(var_addr as *mut _ as *mut T) = T::clone(&sysvar_data);
248            SUCCESS
249        },
250        Err(_) => UNSUPPORTED_SYSVAR,
251    }
252}
253
254struct SyscallStubs {}
255impl solana_sysvar::program_stubs::SyscallStubs for SyscallStubs {
256    fn sol_log(&self, message: &str) {
257        let invoke_context = get_invoke_context();
258        ic_msg!(invoke_context, "Program log: {}", message);
259    }
260
261    fn sol_invoke_signed(
262        &self,
263        instruction: &Instruction,
264        account_infos: &[AccountInfo],
265        signers_seeds: &[&[&[u8]]],
266    ) -> ProgramResult {
267        let invoke_context = get_invoke_context();
268        let log_collector = invoke_context.get_log_collector();
269        let transaction_context = &invoke_context.transaction_context;
270        let instruction_context = transaction_context
271            .get_current_instruction_context()
272            .unwrap();
273        let caller = instruction_context.get_program_key().unwrap();
274
275        stable_log::program_invoke(
276            &log_collector,
277            &instruction.program_id,
278            invoke_context.get_stack_height(),
279        );
280
281        let signers = signers_seeds
282            .iter()
283            .map(|seeds| Pubkey::create_program_address(seeds, caller).unwrap())
284            .collect::<Vec<_>>();
285
286        invoke_context
287            .prepare_next_instruction(instruction.clone(), &signers)
288            .unwrap();
289
290        // Copy caller's account_info modifications into invoke_context accounts
291        let transaction_context = &invoke_context.transaction_context;
292        let instruction_context = transaction_context
293            .get_current_instruction_context()
294            .unwrap();
295        let next_instruction_context = transaction_context.get_next_instruction_context().unwrap();
296        let next_instruction_accounts = next_instruction_context.instruction_accounts();
297        let mut account_indices = Vec::with_capacity(next_instruction_accounts.len());
298        for instruction_account in next_instruction_accounts.iter() {
299            let account_key = transaction_context
300                .get_key_of_account_at_index(instruction_account.index_in_transaction)
301                .unwrap();
302            let account_info_index = account_infos
303                .iter()
304                .position(|account_info| account_info.unsigned_key() == account_key)
305                .ok_or(InstructionError::MissingAccount)
306                .unwrap();
307            let account_info = &account_infos[account_info_index];
308            let index_in_caller = instruction_context
309                .get_index_of_account_in_instruction(instruction_account.index_in_transaction)
310                .unwrap();
311            let mut borrowed_account = instruction_context
312                .try_borrow_instruction_account(index_in_caller)
313                .unwrap();
314            if borrowed_account.get_lamports() != account_info.lamports() {
315                borrowed_account
316                    .set_lamports(account_info.lamports())
317                    .unwrap();
318            }
319            let account_info_data = account_info.try_borrow_data().unwrap();
320            // The redundant check helps to avoid the expensive data comparison if we can
321            match borrowed_account.can_data_be_resized(account_info_data.len()) {
322                Ok(()) => borrowed_account
323                    .set_data_from_slice(&account_info_data)
324                    .unwrap(),
325                Err(err) if borrowed_account.get_data() != *account_info_data => {
326                    panic!("{err:?}");
327                }
328                _ => {}
329            }
330            // Change the owner at the end so that we are allowed to change the lamports and data before
331            if borrowed_account.get_owner() != account_info.owner {
332                borrowed_account
333                    .set_owner(account_info.owner.as_ref())
334                    .unwrap();
335            }
336            if instruction_account.is_writable() {
337                account_indices
338                    .push((instruction_account.index_in_transaction, account_info_index));
339            }
340        }
341
342        let mut compute_units_consumed = 0;
343        invoke_context
344            .process_instruction(&mut compute_units_consumed, &mut ExecuteTimings::default())
345            .map_err(|err| ProgramError::try_from(err).unwrap_or_else(|err| panic!("{}", err)))?;
346
347        // Copy invoke_context accounts modifications into caller's account_info
348        let transaction_context = &invoke_context.transaction_context;
349        let instruction_context = transaction_context
350            .get_current_instruction_context()
351            .unwrap();
352        for (index_in_transaction, account_info_index) in account_indices.into_iter() {
353            let index_in_caller = instruction_context
354                .get_index_of_account_in_instruction(index_in_transaction)
355                .unwrap();
356            let borrowed_account = instruction_context
357                .try_borrow_instruction_account(index_in_caller)
358                .unwrap();
359            let account_info = &account_infos[account_info_index];
360            **account_info.try_borrow_mut_lamports().unwrap() = borrowed_account.get_lamports();
361            if account_info.owner != borrowed_account.get_owner() {
362                // TODO Figure out a better way to allow the System Program to set the account owner
363                #[allow(clippy::transmute_ptr_to_ptr)]
364                #[allow(mutable_transmutes)]
365                let account_info_mut =
366                    unsafe { transmute::<&Pubkey, &mut Pubkey>(account_info.owner) };
367                *account_info_mut = *borrowed_account.get_owner();
368            }
369
370            let new_data = borrowed_account.get_data();
371            let new_len = new_data.len();
372
373            // Resize account_info data
374            if account_info.data_len() != new_len {
375                account_info.resize(new_len)?;
376            }
377
378            // Clone the data
379            let mut data = account_info.try_borrow_mut_data()?;
380            data.clone_from_slice(new_data);
381        }
382
383        stable_log::program_success(&log_collector, &instruction.program_id);
384        Ok(())
385    }
386
387    fn sol_get_clock_sysvar(&self, var_addr: *mut u8) -> u64 {
388        get_sysvar(
389            get_invoke_context().get_sysvar_cache().get_clock(),
390            var_addr,
391        )
392    }
393
394    fn sol_get_epoch_schedule_sysvar(&self, var_addr: *mut u8) -> u64 {
395        get_sysvar(
396            get_invoke_context().get_sysvar_cache().get_epoch_schedule(),
397            var_addr,
398        )
399    }
400
401    fn sol_get_epoch_rewards_sysvar(&self, var_addr: *mut u8) -> u64 {
402        get_sysvar(
403            get_invoke_context().get_sysvar_cache().get_epoch_rewards(),
404            var_addr,
405        )
406    }
407
408    #[allow(deprecated)]
409    fn sol_get_fees_sysvar(&self, var_addr: *mut u8) -> u64 {
410        get_sysvar(get_invoke_context().get_sysvar_cache().get_fees(), var_addr)
411    }
412
413    fn sol_get_rent_sysvar(&self, var_addr: *mut u8) -> u64 {
414        get_sysvar(get_invoke_context().get_sysvar_cache().get_rent(), var_addr)
415    }
416
417    fn sol_get_last_restart_slot(&self, var_addr: *mut u8) -> u64 {
418        get_sysvar(
419            get_invoke_context()
420                .get_sysvar_cache()
421                .get_last_restart_slot(),
422            var_addr,
423        )
424    }
425
426    fn sol_get_return_data(&self) -> Option<(Pubkey, Vec<u8>)> {
427        let (program_id, data) = get_invoke_context().transaction_context.get_return_data();
428        Some((*program_id, data.to_vec()))
429    }
430
431    fn sol_set_return_data(&self, data: &[u8]) {
432        let invoke_context = get_invoke_context();
433        let transaction_context = &mut invoke_context.transaction_context;
434        let instruction_context = transaction_context
435            .get_current_instruction_context()
436            .unwrap();
437        let caller = *instruction_context.get_program_key().unwrap();
438        transaction_context
439            .set_return_data(caller, data.to_vec())
440            .unwrap();
441    }
442
443    fn sol_get_stack_height(&self) -> u64 {
444        let invoke_context = get_invoke_context();
445        invoke_context.get_stack_height().try_into().unwrap()
446    }
447}
448
449pub fn find_file(filename: &str) -> Option<PathBuf> {
450    for dir in default_shared_object_dirs() {
451        let candidate = dir.join(filename);
452        if candidate.exists() {
453            return Some(candidate);
454        }
455    }
456    None
457}
458
459fn default_shared_object_dirs() -> Vec<PathBuf> {
460    let mut search_path = vec![];
461    if let Ok(bpf_out_dir) = std::env::var("BPF_OUT_DIR") {
462        search_path.push(PathBuf::from(bpf_out_dir));
463    } else if let Ok(bpf_out_dir) = std::env::var("SBF_OUT_DIR") {
464        search_path.push(PathBuf::from(bpf_out_dir));
465    }
466    search_path.push(PathBuf::from("tests/fixtures"));
467    if let Ok(dir) = std::env::current_dir() {
468        search_path.push(dir);
469    }
470    trace!("SBF .so search path: {search_path:?}");
471    search_path
472}
473
474pub fn read_file<P: AsRef<Path>>(path: P) -> Vec<u8> {
475    let path = path.as_ref();
476    let mut file = File::open(path)
477        .unwrap_or_else(|err| panic!("Failed to open \"{}\": {}", path.display(), err));
478
479    let mut file_data = Vec::new();
480    file.read_to_end(&mut file_data)
481        .unwrap_or_else(|err| panic!("Failed to read \"{}\": {}", path.display(), err));
482    file_data
483}
484
485pub struct ProgramTest {
486    accounts: Vec<(Pubkey, AccountSharedData)>,
487    genesis_accounts: Vec<(Pubkey, AccountSharedData)>,
488    builtin_programs: Vec<(Pubkey, &'static str, ProgramCacheEntry)>,
489    compute_max_units: Option<u64>,
490    prefer_bpf: bool,
491    deactivate_feature_set: HashSet<Pubkey>,
492    transaction_account_lock_limit: Option<usize>,
493}
494
495impl Default for ProgramTest {
496    /// Initialize a new ProgramTest
497    ///
498    /// If the `BPF_OUT_DIR` environment variable is defined, BPF programs will be preferred over
499    /// over a native instruction processor.  The `ProgramTest::prefer_bpf()` method may be
500    /// used to override this preference at runtime.  `cargo test-bpf` will set `BPF_OUT_DIR`
501    /// automatically.
502    ///
503    /// SBF program shared objects and account data files are searched for in
504    /// * the value of the `BPF_OUT_DIR` environment variable
505    /// * the `tests/fixtures` sub-directory
506    /// * the current working directory
507    ///
508    fn default() -> Self {
509        agave_logger::setup_with_default(
510            "solana_sbpf::vm=debug,solana_runtime::message_processor=debug,\
511             solana_runtime::system_instruction_processor=trace,solana_program_test=info",
512        );
513        let prefer_bpf =
514            std::env::var("BPF_OUT_DIR").is_ok() || std::env::var("SBF_OUT_DIR").is_ok();
515
516        Self {
517            accounts: vec![],
518            genesis_accounts: vec![],
519            builtin_programs: vec![],
520            compute_max_units: None,
521            prefer_bpf,
522            deactivate_feature_set: HashSet::default(),
523            transaction_account_lock_limit: None,
524        }
525    }
526}
527
528impl ProgramTest {
529    /// Create a `ProgramTest`.
530    ///
531    /// This is a wrapper around [`default`] and [`add_program`]. See their documentation for more
532    /// details.
533    ///
534    /// [`default`]: #method.default
535    /// [`add_program`]: #method.add_program
536    pub fn new(
537        program_name: &'static str,
538        program_id: Pubkey,
539        builtin_function: Option<BuiltinFunctionWithContext>,
540    ) -> Self {
541        let mut me = Self::default();
542        me.add_program(program_name, program_id, builtin_function);
543        me
544    }
545
546    /// Override default SBF program selection
547    pub fn prefer_bpf(&mut self, prefer_bpf: bool) {
548        self.prefer_bpf = prefer_bpf;
549    }
550
551    /// Override the default maximum compute units
552    pub fn set_compute_max_units(&mut self, compute_max_units: u64) {
553        debug_assert!(
554            compute_max_units <= i64::MAX as u64,
555            "Compute unit limit must fit in `i64::MAX`"
556        );
557        self.compute_max_units = Some(compute_max_units);
558    }
559
560    /// Override the default transaction account lock limit
561    pub fn set_transaction_account_lock_limit(&mut self, transaction_account_lock_limit: usize) {
562        self.transaction_account_lock_limit = Some(transaction_account_lock_limit);
563    }
564
565    /// Add an account to the test environment's genesis config.
566    pub fn add_genesis_account(&mut self, address: Pubkey, account: Account) {
567        self.genesis_accounts
568            .push((address, AccountSharedData::from(account)));
569    }
570
571    /// Add an account to the test environment
572    pub fn add_account(&mut self, address: Pubkey, account: Account) {
573        self.accounts
574            .push((address, AccountSharedData::from(account)));
575    }
576
577    /// Add an account to the test environment with the account data in the provided `filename`
578    pub fn add_account_with_file_data(
579        &mut self,
580        address: Pubkey,
581        lamports: u64,
582        owner: Pubkey,
583        filename: &str,
584    ) {
585        self.add_account(
586            address,
587            Account {
588                lamports,
589                data: read_file(find_file(filename).unwrap_or_else(|| {
590                    panic!("Unable to locate {filename}");
591                })),
592                owner,
593                executable: false,
594                rent_epoch: 0,
595            },
596        );
597    }
598
599    /// Add an account to the test environment with the account data in the provided as a base 64
600    /// string
601    pub fn add_account_with_base64_data(
602        &mut self,
603        address: Pubkey,
604        lamports: u64,
605        owner: Pubkey,
606        data_base64: &str,
607    ) {
608        self.add_account(
609            address,
610            Account {
611                lamports,
612                data: BASE64_STANDARD
613                    .decode(data_base64)
614                    .unwrap_or_else(|err| panic!("Failed to base64 decode: {err}")),
615                owner,
616                executable: false,
617                rent_epoch: 0,
618            },
619        );
620    }
621
622    pub fn add_sysvar_account<S: SysvarSerialize>(&mut self, address: Pubkey, sysvar: &S) {
623        let account = create_account_shared_data_for_test(sysvar);
624        self.add_account(address, account.into());
625    }
626
627    /// Add a BPF Upgradeable program to the test environment's genesis config.
628    ///
629    /// When testing BPF programs using the program ID of a runtime builtin
630    /// program - such as Core BPF programs - the program accounts must be
631    /// added to the genesis config in order to make them available to the new
632    /// Bank as it's being initialized.
633    ///
634    /// The presence of these program accounts will cause Bank to skip adding
635    /// the builtin version of the program, allowing the provided BPF program
636    /// to be used at the designated program ID instead.
637    ///
638    /// See https://github.com/anza-xyz/agave/blob/c038908600b8a1b0080229dea015d7fc9939c418/runtime/src/bank.rs#L5109-L5126.
639    pub fn add_upgradeable_program_to_genesis(
640        &mut self,
641        program_name: &'static str,
642        program_id: &Pubkey,
643    ) {
644        let program_file = find_file(&format!("{program_name}.so")).unwrap_or_else(|| {
645            panic!("Program file data not available for {program_name} ({program_id})")
646        });
647        let elf = read_file(program_file);
648        let program_accounts =
649            programs::bpf_loader_upgradeable_program_accounts(program_id, &elf, &Rent::default());
650        for (address, account) in program_accounts {
651            self.add_genesis_account(address, account);
652        }
653    }
654
655    /// Add a SBF program to the test environment.
656    ///
657    /// `program_name` will also be used to locate the SBF shared object in the current or fixtures
658    /// directory.
659    ///
660    /// If `builtin_function` is provided, the natively built-program may be used instead of the
661    /// SBF shared object depending on the `BPF_OUT_DIR` environment variable.
662    pub fn add_program(
663        &mut self,
664        program_name: &'static str,
665        program_id: Pubkey,
666        builtin_function: Option<BuiltinFunctionWithContext>,
667    ) {
668        let add_bpf = |this: &mut ProgramTest, program_file: PathBuf| {
669            let data = read_file(&program_file);
670            info!(
671                "\"{}\" SBF program from {}{}",
672                program_name,
673                program_file.display(),
674                std::fs::metadata(&program_file)
675                    .map(|metadata| {
676                        metadata
677                            .modified()
678                            .map(|time| {
679                                format!(
680                                    ", modified {}",
681                                    HumanTime::from(time)
682                                        .to_text_en(Accuracy::Precise, Tense::Past)
683                                )
684                            })
685                            .ok()
686                    })
687                    .ok()
688                    .flatten()
689                    .unwrap_or_default()
690            );
691
692            this.add_account(
693                program_id,
694                Account {
695                    lamports: Rent::default().minimum_balance(data.len()).max(1),
696                    data,
697                    owner: solana_sdk_ids::bpf_loader::id(),
698                    executable: true,
699                    rent_epoch: 0,
700                },
701            );
702        };
703
704        let warn_invalid_program_name = || {
705            let valid_program_names = default_shared_object_dirs()
706                .iter()
707                .filter_map(|dir| dir.read_dir().ok())
708                .flat_map(|read_dir| {
709                    read_dir.filter_map(|entry| {
710                        let path = entry.ok()?.path();
711                        if !path.is_file() {
712                            return None;
713                        }
714                        match path.extension()?.to_str()? {
715                            "so" => Some(path.file_stem()?.to_os_string()),
716                            _ => None,
717                        }
718                    })
719                })
720                .collect::<Vec<_>>();
721
722            if valid_program_names.is_empty() {
723                // This should be unreachable as `test-bpf` should guarantee at least one shared
724                // object exists somewhere.
725                warn!("No SBF shared objects found.");
726                return;
727            }
728
729            warn!(
730                "Possible bogus program name. Ensure the program name ({program_name}) matches \
731                 one of the following recognizable program names:",
732            );
733            for name in valid_program_names {
734                warn!(" - {}", name.to_str().unwrap());
735            }
736        };
737
738        let program_file = find_file(&format!("{program_name}.so"));
739        match (self.prefer_bpf, program_file, builtin_function) {
740            // If SBF is preferred (i.e., `test-sbf` is invoked) and a BPF shared object exists,
741            // use that as the program data.
742            (true, Some(file), _) => add_bpf(self, file),
743
744            // If SBF is not required (i.e., we were invoked with `test`), use the provided
745            // processor function as is.
746            (false, _, Some(builtin_function)) => {
747                self.add_builtin_program(program_name, program_id, builtin_function)
748            }
749
750            // Invalid: `test-sbf` invocation with no matching SBF shared object.
751            (true, None, _) => {
752                warn_invalid_program_name();
753                panic!("Program file data not available for {program_name} ({program_id})");
754            }
755
756            // Invalid: regular `test` invocation without a processor.
757            (false, _, None) => {
758                panic!("Program processor not available for {program_name} ({program_id})");
759            }
760        }
761    }
762
763    /// Add a builtin program to the test environment.
764    ///
765    /// Note that builtin programs are responsible for their own `stable_log` output.
766    pub fn add_builtin_program(
767        &mut self,
768        program_name: &'static str,
769        program_id: Pubkey,
770        builtin_function: BuiltinFunctionWithContext,
771    ) {
772        info!("\"{program_name}\" builtin program");
773        self.builtin_programs.push((
774            program_id,
775            program_name,
776            ProgramCacheEntry::new_builtin(0, program_name.len(), builtin_function),
777        ));
778    }
779
780    /// Deactivate a runtime feature.
781    ///
782    /// Note that all features are activated by default.
783    pub fn deactivate_feature(&mut self, feature_id: Pubkey) {
784        self.deactivate_feature_set.insert(feature_id);
785    }
786
787    fn setup_bank(
788        &mut self,
789    ) -> (
790        Arc<RwLock<BankForks>>,
791        Arc<RwLock<BlockCommitmentCache>>,
792        Hash,
793        GenesisConfigInfo,
794    ) {
795        {
796            use std::sync::Once;
797            static ONCE: Once = Once::new();
798
799            ONCE.call_once(|| {
800                solana_sysvar::program_stubs::set_syscall_stubs(Box::new(SyscallStubs {}));
801            });
802        }
803
804        let rent = Rent::default();
805        let fee_rate_governor = FeeRateGovernor {
806            // Initialize with a non-zero fee
807            lamports_per_signature: DEFAULT_TARGET_LAMPORTS_PER_SIGNATURE / 2,
808            ..FeeRateGovernor::default()
809        };
810        let bootstrap_validator_pubkey = Pubkey::new_unique();
811        let bootstrap_validator_stake_lamports =
812            rent.minimum_balance(VoteStateV4::size_of()) + 1_000_000 * LAMPORTS_PER_SOL;
813
814        let mint_keypair = Keypair::new();
815        let voting_keypair = Keypair::new();
816
817        let mut genesis_config = create_genesis_config_with_leader_ex(
818            1_000_000 * LAMPORTS_PER_SOL,
819            &mint_keypair.pubkey(),
820            &bootstrap_validator_pubkey,
821            &voting_keypair.pubkey(),
822            &Pubkey::new_unique(),
823            None,
824            bootstrap_validator_stake_lamports,
825            42,
826            fee_rate_governor,
827            rent.clone(),
828            ClusterType::Development,
829            std::mem::take(&mut self.genesis_accounts),
830        );
831
832        // Remove features tagged to deactivate
833        for deactivate_feature_pk in &self.deactivate_feature_set {
834            if FEATURE_NAMES.contains_key(deactivate_feature_pk) {
835                match genesis_config.accounts.remove(deactivate_feature_pk) {
836                    Some(_) => debug!("Feature for {deactivate_feature_pk:?} deactivated"),
837                    None => warn!(
838                        "Feature {deactivate_feature_pk:?} set for deactivation not found in \
839                         genesis_config account list, ignored."
840                    ),
841                }
842            } else {
843                warn!(
844                    "Feature {deactivate_feature_pk:?} set for deactivation is not a known \
845                     Feature public key"
846                );
847            }
848        }
849
850        let target_tick_duration = Duration::from_micros(100);
851        genesis_config.poh_config = PohConfig::new_sleep(target_tick_duration);
852        debug!("Payer address: {}", mint_keypair.pubkey());
853        debug!("Genesis config: {genesis_config}");
854
855        let bank = Bank::new_from_genesis(
856            &genesis_config,
857            Arc::new(RuntimeConfig {
858                compute_budget: self.compute_max_units.map(|max_units| ComputeBudget {
859                    compute_unit_limit: max_units,
860                    ..ComputeBudget::new_with_defaults(
861                        genesis_config
862                            .accounts
863                            .contains_key(&raise_cpi_nesting_limit_to_8::id()),
864                        genesis_config
865                            .accounts
866                            .contains_key(&increase_cpi_account_info_limit::id()),
867                    )
868                }),
869                transaction_account_lock_limit: self.transaction_account_lock_limit,
870                ..RuntimeConfig::default()
871            }),
872            Vec::default(),
873            None,
874            ACCOUNTS_DB_CONFIG_FOR_TESTING,
875            None,
876            None,
877            Arc::default(),
878            None,
879            None,
880        );
881
882        // Add commonly-used SPL programs as a convenience to the user
883        for (program_id, account) in programs::spl_programs(&rent).iter() {
884            bank.store_account(program_id, account);
885        }
886
887        // Add migrated Core BPF programs.
888        for (program_id, account) in programs::core_bpf_programs(&rent, |feature_id| {
889            genesis_config.accounts.contains_key(feature_id)
890        })
891        .iter()
892        {
893            bank.store_account(program_id, account);
894        }
895
896        // User-supplied additional builtins
897        let mut builtin_programs = Vec::new();
898        std::mem::swap(&mut self.builtin_programs, &mut builtin_programs);
899        for (program_id, name, builtin) in builtin_programs.into_iter() {
900            bank.add_builtin(program_id, name, builtin);
901        }
902
903        for (address, account) in self.accounts.iter() {
904            if bank.get_account(address).is_some() {
905                info!("Overriding account at {address}");
906            }
907            bank.store_account(address, account);
908        }
909        bank.set_capitalization_for_tests(bank.calculate_capitalization_for_tests());
910        // Advance beyond slot 0 for a slightly more realistic test environment
911        let bank = {
912            let bank = Arc::new(bank);
913            bank.fill_bank_with_ticks_for_tests();
914            let bank = Bank::new_from_parent(bank.clone(), bank.collector_id(), bank.slot() + 1);
915            debug!("Bank slot: {}", bank.slot());
916            bank
917        };
918        let slot = bank.slot();
919        let last_blockhash = bank.last_blockhash();
920        let bank_forks = BankForks::new_rw_arc(bank);
921        let block_commitment_cache = Arc::new(RwLock::new(
922            BlockCommitmentCache::new_for_tests_with_slots(slot, slot),
923        ));
924
925        (
926            bank_forks,
927            block_commitment_cache,
928            last_blockhash,
929            GenesisConfigInfo {
930                genesis_config,
931                mint_keypair,
932                voting_keypair,
933                validator_pubkey: bootstrap_validator_pubkey,
934            },
935        )
936    }
937
938    pub async fn start(mut self) -> (BanksClient, Keypair, Hash) {
939        let (bank_forks, block_commitment_cache, last_blockhash, gci) = self.setup_bank();
940        let target_tick_duration = gci.genesis_config.poh_config.target_tick_duration;
941        let target_slot_duration = target_tick_duration * gci.genesis_config.ticks_per_slot as u32;
942        let transport = start_local_server(
943            bank_forks.clone(),
944            block_commitment_cache.clone(),
945            target_tick_duration,
946        )
947        .await;
948        let banks_client = start_client(transport)
949            .await
950            .unwrap_or_else(|err| panic!("Failed to start banks client: {err}"));
951
952        // Run a simulated PohService to provide the client with new blockhashes.  New blockhashes
953        // are required when sending multiple otherwise identical transactions in series from a
954        // test
955        tokio::spawn(async move {
956            loop {
957                tokio::time::sleep(target_slot_duration).await;
958                bank_forks
959                    .read()
960                    .unwrap()
961                    .working_bank()
962                    .register_unique_recent_blockhash_for_test();
963            }
964        });
965
966        (banks_client, gci.mint_keypair, last_blockhash)
967    }
968
969    /// Start the test client
970    ///
971    /// Returns a `BanksClient` interface into the test environment as well as a payer `Keypair`
972    /// with SOL for sending transactions
973    pub async fn start_with_context(mut self) -> ProgramTestContext {
974        let (bank_forks, block_commitment_cache, last_blockhash, gci) = self.setup_bank();
975        let target_tick_duration = gci.genesis_config.poh_config.target_tick_duration;
976        let transport = start_local_server(
977            bank_forks.clone(),
978            block_commitment_cache.clone(),
979            target_tick_duration,
980        )
981        .await;
982        let banks_client = start_client(transport)
983            .await
984            .unwrap_or_else(|err| panic!("Failed to start banks client: {err}"));
985
986        ProgramTestContext::new(
987            bank_forks,
988            block_commitment_cache,
989            banks_client,
990            last_blockhash,
991            gci,
992        )
993    }
994}
995
996#[async_trait]
997pub trait ProgramTestBanksClientExt {
998    /// Get a new latest blockhash, similar in spirit to RpcClient::get_latest_blockhash()
999    async fn get_new_latest_blockhash(&mut self, blockhash: &Hash) -> io::Result<Hash>;
1000}
1001
1002#[async_trait]
1003impl ProgramTestBanksClientExt for BanksClient {
1004    async fn get_new_latest_blockhash(&mut self, blockhash: &Hash) -> io::Result<Hash> {
1005        let mut num_retries = 0;
1006        let start = Instant::now();
1007        while start.elapsed().as_secs() < 5 {
1008            let new_blockhash = self.get_latest_blockhash().await?;
1009            if new_blockhash != *blockhash {
1010                return Ok(new_blockhash);
1011            }
1012            debug!("Got same blockhash ({blockhash:?}), will retry...");
1013
1014            tokio::time::sleep(Duration::from_millis(200)).await;
1015            num_retries += 1;
1016        }
1017
1018        Err(io::Error::other(format!(
1019            "Unable to get new blockhash after {}ms (retried {} times), stuck at {}",
1020            start.elapsed().as_millis(),
1021            num_retries,
1022            blockhash
1023        )))
1024    }
1025}
1026
1027struct DroppableTask<T>(Arc<AtomicBool>, JoinHandle<T>);
1028
1029impl<T> Drop for DroppableTask<T> {
1030    fn drop(&mut self) {
1031        self.0.store(true, Ordering::Relaxed);
1032        trace!(
1033            "stopping task, which is currently {}",
1034            if self.1.is_finished() {
1035                "finished"
1036            } else {
1037                "running"
1038            }
1039        );
1040    }
1041}
1042
1043pub struct ProgramTestContext {
1044    pub banks_client: BanksClient,
1045    pub last_blockhash: Hash,
1046    pub payer: Keypair,
1047    genesis_config: GenesisConfig,
1048    bank_forks: Arc<RwLock<BankForks>>,
1049    block_commitment_cache: Arc<RwLock<BlockCommitmentCache>>,
1050    _bank_task: DroppableTask<()>,
1051}
1052
1053impl ProgramTestContext {
1054    fn new(
1055        bank_forks: Arc<RwLock<BankForks>>,
1056        block_commitment_cache: Arc<RwLock<BlockCommitmentCache>>,
1057        banks_client: BanksClient,
1058        last_blockhash: Hash,
1059        genesis_config_info: GenesisConfigInfo,
1060    ) -> Self {
1061        // Run a simulated PohService to provide the client with new blockhashes.  New blockhashes
1062        // are required when sending multiple otherwise identical transactions in series from a
1063        // test
1064        let running_bank_forks = bank_forks.clone();
1065        let target_tick_duration = genesis_config_info
1066            .genesis_config
1067            .poh_config
1068            .target_tick_duration;
1069        let target_slot_duration =
1070            target_tick_duration * genesis_config_info.genesis_config.ticks_per_slot as u32;
1071        let exit = Arc::new(AtomicBool::new(false));
1072        let bank_task = DroppableTask(
1073            exit.clone(),
1074            tokio::spawn(async move {
1075                loop {
1076                    if exit.load(Ordering::Relaxed) {
1077                        break;
1078                    }
1079                    tokio::time::sleep(target_slot_duration).await;
1080                    running_bank_forks
1081                        .read()
1082                        .unwrap()
1083                        .working_bank()
1084                        .register_unique_recent_blockhash_for_test();
1085                }
1086            }),
1087        );
1088
1089        Self {
1090            banks_client,
1091            last_blockhash,
1092            payer: genesis_config_info.mint_keypair,
1093            genesis_config: genesis_config_info.genesis_config,
1094            bank_forks,
1095            block_commitment_cache,
1096            _bank_task: bank_task,
1097        }
1098    }
1099
1100    pub fn genesis_config(&self) -> &GenesisConfig {
1101        &self.genesis_config
1102    }
1103
1104    /// Manually increment vote credits for the current epoch in the specified vote account to simulate validator voting activity
1105    pub fn increment_vote_account_credits(
1106        &mut self,
1107        vote_account_address: &Pubkey,
1108        number_of_credits: u64,
1109    ) {
1110        let bank_forks = self.bank_forks.read().unwrap();
1111        let bank = bank_forks.working_bank();
1112
1113        // generate some vote activity for rewards
1114        let mut vote_account = bank.get_account(vote_account_address).unwrap();
1115        let mut vote_state =
1116            VoteStateV4::deserialize(vote_account.data(), vote_account_address).unwrap();
1117
1118        let epoch = bank.epoch();
1119        // Inlined from vote program - maximum number of epoch credits to keep in history
1120        const MAX_EPOCH_CREDITS_HISTORY: usize = 64;
1121        for _ in 0..number_of_credits {
1122            // Inline increment_credits logic from vote program.
1123            let credits = 1;
1124
1125            // never seen a credit
1126            if vote_state.epoch_credits.is_empty() {
1127                vote_state.epoch_credits.push((epoch, 0, 0));
1128            } else if epoch != vote_state.epoch_credits.last().unwrap().0 {
1129                let (_, credits_val, prev_credits) = *vote_state.epoch_credits.last().unwrap();
1130
1131                if credits_val != prev_credits {
1132                    // if credits were earned previous epoch
1133                    // append entry at end of list for the new epoch
1134                    vote_state
1135                        .epoch_credits
1136                        .push((epoch, credits_val, credits_val));
1137                } else {
1138                    // else just move the current epoch
1139                    vote_state.epoch_credits.last_mut().unwrap().0 = epoch;
1140                }
1141
1142                // Remove too old epoch_credits
1143                if vote_state.epoch_credits.len() > MAX_EPOCH_CREDITS_HISTORY {
1144                    vote_state.epoch_credits.remove(0);
1145                }
1146            }
1147
1148            vote_state.epoch_credits.last_mut().unwrap().1 = vote_state
1149                .epoch_credits
1150                .last()
1151                .unwrap()
1152                .1
1153                .saturating_add(credits);
1154        }
1155        let versioned = VoteStateVersions::new_v4(vote_state);
1156        vote_account.set_state(&versioned).unwrap();
1157        bank.store_account(vote_account_address, &vote_account);
1158    }
1159
1160    /// Create or overwrite an account, subverting normal runtime checks.
1161    ///
1162    /// This method exists to make it easier to set up artificial situations
1163    /// that would be difficult to replicate by sending individual transactions.
1164    /// Beware that it can be used to create states that would not be reachable
1165    /// by sending transactions!
1166    pub fn set_account(&mut self, address: &Pubkey, account: &AccountSharedData) {
1167        let bank_forks = self.bank_forks.read().unwrap();
1168        let bank = bank_forks.working_bank();
1169        bank.store_account(address, account);
1170    }
1171
1172    /// Create or overwrite a sysvar, subverting normal runtime checks.
1173    ///
1174    /// This method exists to make it easier to set up artificial situations
1175    /// that would be difficult to replicate on a new test cluster. Beware
1176    /// that it can be used to create states that would not be reachable
1177    /// under normal conditions!
1178    pub fn set_sysvar<T: SysvarId + SysvarSerialize>(&self, sysvar: &T) {
1179        let bank_forks = self.bank_forks.read().unwrap();
1180        let bank = bank_forks.working_bank();
1181        bank.set_sysvar_for_tests(sysvar);
1182    }
1183
1184    /// Force the working bank ahead to a new slot
1185    pub fn warp_to_slot(&mut self, warp_slot: Slot) -> Result<(), ProgramTestError> {
1186        let mut bank_forks = self.bank_forks.write().unwrap();
1187        let bank = bank_forks.working_bank();
1188
1189        // Fill ticks until a new blockhash is recorded, otherwise retried transactions will have
1190        // the same signature
1191        bank.fill_bank_with_ticks_for_tests();
1192
1193        // Ensure that we are actually progressing forward
1194        let working_slot = bank.slot();
1195        if warp_slot <= working_slot {
1196            return Err(ProgramTestError::InvalidWarpSlot);
1197        }
1198
1199        // Warp ahead to one slot *before* the desired slot because the bank
1200        // from Bank::warp_from_parent() is frozen. If the desired slot is one
1201        // slot *after* the working_slot, no need to warp at all.
1202        let pre_warp_slot = warp_slot - 1;
1203        let warp_bank = if pre_warp_slot == working_slot {
1204            bank.freeze();
1205            bank
1206        } else {
1207            bank_forks
1208                .insert(Bank::warp_from_parent(
1209                    bank,
1210                    &Pubkey::default(),
1211                    pre_warp_slot,
1212                ))
1213                .clone_without_scheduler()
1214        };
1215
1216        bank_forks.set_root(
1217            pre_warp_slot,
1218            None, // snapshots are disabled
1219            Some(pre_warp_slot),
1220        );
1221
1222        // warp_bank is frozen so go forward to get unfrozen bank at warp_slot
1223        bank_forks.insert(Bank::new_from_parent(
1224            warp_bank,
1225            &Pubkey::default(),
1226            warp_slot,
1227        ));
1228
1229        // Update block commitment cache, otherwise banks server will poll at
1230        // the wrong slot
1231        let mut w_block_commitment_cache = self.block_commitment_cache.write().unwrap();
1232        // HACK: The root set here should be `pre_warp_slot`, but since we're
1233        // in a testing environment, the root bank never updates after a warp.
1234        // The ticking thread only updates the working bank, and never the root
1235        // bank.
1236        w_block_commitment_cache.set_all_slots(warp_slot, warp_slot);
1237
1238        let bank = bank_forks.working_bank();
1239        self.last_blockhash = bank.last_blockhash();
1240        Ok(())
1241    }
1242
1243    pub fn warp_to_epoch(&mut self, warp_epoch: Epoch) -> Result<(), ProgramTestError> {
1244        let warp_slot = self
1245            .genesis_config
1246            .epoch_schedule
1247            .get_first_slot_in_epoch(warp_epoch);
1248        self.warp_to_slot(warp_slot)
1249    }
1250
1251    /// warp forward one more slot and force reward interval end
1252    pub fn warp_forward_force_reward_interval_end(&mut self) -> Result<(), ProgramTestError> {
1253        let mut bank_forks = self.bank_forks.write().unwrap();
1254        let bank = bank_forks.working_bank();
1255
1256        // Fill ticks until a new blockhash is recorded, otherwise retried transactions will have
1257        // the same signature
1258        bank.fill_bank_with_ticks_for_tests();
1259        let pre_warp_slot = bank.slot();
1260
1261        bank_forks.set_root(
1262            pre_warp_slot,
1263            None, // snapshot_controller
1264            Some(pre_warp_slot),
1265        );
1266
1267        // warp_bank is frozen so go forward to get unfrozen bank at warp_slot
1268        let warp_slot = pre_warp_slot + 1;
1269        let mut warp_bank = Bank::new_from_parent(bank, &Pubkey::default(), warp_slot);
1270
1271        warp_bank.force_reward_interval_end_for_tests();
1272        bank_forks.insert(warp_bank);
1273
1274        // Update block commitment cache, otherwise banks server will poll at
1275        // the wrong slot
1276        let mut w_block_commitment_cache = self.block_commitment_cache.write().unwrap();
1277        // HACK: The root set here should be `pre_warp_slot`, but since we're
1278        // in a testing environment, the root bank never updates after a warp.
1279        // The ticking thread only updates the working bank, and never the root
1280        // bank.
1281        w_block_commitment_cache.set_all_slots(warp_slot, warp_slot);
1282
1283        let bank = bank_forks.working_bank();
1284        self.last_blockhash = bank.last_blockhash();
1285        Ok(())
1286    }
1287
1288    /// Get a new latest blockhash, similar in spirit to RpcClient::get_latest_blockhash()
1289    pub async fn get_new_latest_blockhash(&mut self) -> io::Result<Hash> {
1290        let blockhash = self
1291            .banks_client
1292            .get_new_latest_blockhash(&self.last_blockhash)
1293            .await?;
1294        self.last_blockhash = blockhash;
1295        Ok(blockhash)
1296    }
1297
1298    /// record a hard fork slot in working bank; should be in the past
1299    pub fn register_hard_fork(&mut self, hard_fork_slot: Slot) {
1300        self.bank_forks
1301            .read()
1302            .unwrap()
1303            .working_bank()
1304            .register_hard_fork(hard_fork_slot)
1305    }
1306}