solana_test_client/
program_test.rs

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