Skip to main content

rialo_s_program_runtime/
invoke_context.rs

1// Copyright (c) Subzero Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3// This file is either (a) original to Subzero Labs, Inc. or (b) derived from the Anza codebase and modified by Subzero Labs, Inc.
4
5use std::{
6    alloc::Layout,
7    cell::RefCell,
8    fmt::{self, Debug},
9    rc::Rc,
10};
11
12use rialo_hash::Hash;
13use rialo_s_account::{create_account_shared_data_for_test, AccountSharedData, StoredAccount};
14use rialo_s_clock::Slot;
15use rialo_s_compute_budget::compute_budget::ComputeBudget;
16use rialo_s_epoch_schedule::EpochSchedule;
17use rialo_s_feature_set::{
18    lift_cpi_caller_restriction, move_precompile_verification_to_svm,
19    remove_accounts_executable_flag_checks, FeatureSet,
20};
21use rialo_s_instruction::{error::InstructionError, AccountMeta};
22use rialo_s_log_collector::{ic_msg, LogCollector};
23use rialo_s_measure::measure::Measure;
24use rialo_s_precompiles::Precompile;
25use rialo_s_pubkey::Pubkey;
26use rialo_s_sdk_ids::{bpf_loader_deprecated, native_loader, sysvar};
27use rialo_s_stable_layout::stable_instruction::StableInstruction;
28use rialo_s_timings::{ExecuteDetailsTimings, ExecuteTimings};
29use rialo_s_transaction_context::{
30    IndexOfAccount, InstructionAccount, TransactionAccount, TransactionContext,
31};
32use rialo_s_type_overrides::sync::{atomic::Ordering, Arc};
33pub use rialo_stake_cache_interface::{
34    StakeCacheData, StakesHandle, ValidatorAccount, ValidatorAccountView,
35};
36use solana_sbpf::{
37    ebpf::MM_HEAP_START,
38    error::{EbpfError, ProgramResult},
39    memory_region::MemoryMapping,
40    program::{BuiltinFunction, SBPFVersion},
41    vm::{Config, ContextObject, EbpfVm},
42};
43
44use crate::{
45    loaded_programs::{
46        ProgramCacheEntry, ProgramCacheEntryType, ProgramCacheForTxBatch,
47        ProgramRuntimeEnvironments,
48    },
49    stable_log,
50    sysvar_cache::SysvarCache,
51};
52
53pub type BuiltinFunctionWithContext = BuiltinFunction<InvokeContext<'static>>;
54
55/// Adapter so we can unify the interfaces of built-in programs and syscalls
56#[macro_export]
57macro_rules! declare_process_instruction {
58    ($process_instruction:ident, $cu_to_consume:expr, |$invoke_context:ident| $inner:tt) => {
59        $crate::solana_sbpf::declare_builtin_function!(
60            $process_instruction,
61            fn rust(
62                invoke_context: &mut $crate::invoke_context::InvokeContext<'_>,
63                _arg0: u64,
64                _arg1: u64,
65                _arg2: u64,
66                _arg3: u64,
67                _arg4: u64,
68                _memory_mapping: &mut $crate::solana_sbpf::memory_region::MemoryMapping<'_>,
69            ) -> std::result::Result<u64, Box<dyn std::error::Error>> {
70                fn process_instruction_inner(
71                    $invoke_context: &mut $crate::invoke_context::InvokeContext<'_>,
72                ) -> std::result::Result<(), $crate::__private::InstructionError>
73                    $inner
74
75                let consumption_result = if $cu_to_consume > 0
76                {
77                    invoke_context.consume_checked($cu_to_consume)
78                } else {
79                    Ok(())
80                };
81                consumption_result
82                    .and_then(|_| {
83                        process_instruction_inner(invoke_context)
84                            .map(|_| 0)
85                            .map_err(|err| Box::new(err) as Box<dyn std::error::Error>)
86                    })
87                    .into()
88            }
89        );
90    };
91}
92
93impl ContextObject for InvokeContext<'_> {
94    fn trace(&mut self, state: [u64; 12]) {
95        self.syscall_context
96            .last_mut()
97            .unwrap()
98            .as_mut()
99            .unwrap()
100            .trace_log
101            .push(state);
102    }
103
104    fn consume(&mut self, amount: u64) {
105        // 1 to 1 instruction to compute unit mapping
106        // ignore overflow, Ebpf will bail if exceeded
107        let mut compute_meter = self.compute_meter.borrow_mut();
108        *compute_meter = compute_meter.saturating_sub(amount);
109    }
110
111    fn get_remaining(&self) -> u64 {
112        *self.compute_meter.borrow()
113    }
114}
115
116#[derive(Clone, PartialEq, Eq, Debug)]
117pub struct AllocErr;
118impl fmt::Display for AllocErr {
119    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
120        f.write_str("Error: Memory allocation failed")
121    }
122}
123
124pub struct BpfAllocator {
125    len: u64,
126    pos: u64,
127}
128
129impl BpfAllocator {
130    pub fn new(len: u64) -> Self {
131        Self { len, pos: 0 }
132    }
133
134    pub fn alloc(&mut self, layout: Layout) -> Result<u64, AllocErr> {
135        let bytes_to_align = (self.pos as *const u8).align_offset(layout.align()) as u64;
136        if self
137            .pos
138            .saturating_add(bytes_to_align)
139            .saturating_add(layout.size() as u64)
140            <= self.len
141        {
142            self.pos = self.pos.saturating_add(bytes_to_align);
143            let addr = MM_HEAP_START.saturating_add(self.pos);
144            self.pos = self.pos.saturating_add(layout.size() as u64);
145            Ok(addr)
146        } else {
147            Err(AllocErr)
148        }
149    }
150}
151
152pub struct EnvironmentConfig<'a> {
153    pub blockhash: Hash,
154    pub blockhash_kelvins_per_signature: u64,
155    epoch_total_stake: u64,
156    get_epoch_vote_account_stake_callback: &'a dyn Fn(&'a Pubkey) -> u64,
157    pub feature_set: Arc<FeatureSet>,
158    sysvar_cache: &'a SysvarCache,
159    pub random_seed: u64,
160    /// Handle to freeze the stake cache during instruction execution.
161    /// Builtin programs can call `freeze_stakes_and_get_validator_stakes()`
162    /// to append an O(1) clone of the pending stake cache data to the history.
163    /// The handle also contains the pending stake cache data for read-only access,
164    /// including the next epoch number.
165    stakes_handle: StakesHandle,
166}
167impl<'a> EnvironmentConfig<'a> {
168    #[allow(clippy::too_many_arguments)]
169    pub fn new(
170        blockhash: Hash,
171        blockhash_kelvins_per_signature: u64,
172        epoch_total_stake: u64,
173        get_epoch_vote_account_stake_callback: &'a dyn Fn(&'a Pubkey) -> u64,
174        feature_set: Arc<FeatureSet>,
175        sysvar_cache: &'a SysvarCache,
176        random_seed: u64,
177        stakes_handle: StakesHandle,
178    ) -> Self {
179        Self {
180            blockhash,
181            blockhash_kelvins_per_signature,
182            epoch_total_stake,
183            get_epoch_vote_account_stake_callback,
184            feature_set,
185            sysvar_cache,
186            random_seed,
187            stakes_handle,
188        }
189    }
190}
191
192pub struct SyscallContext {
193    pub allocator: BpfAllocator,
194    pub accounts_metadata: Vec<SerializedAccountMetadata>,
195    pub trace_log: Vec<[u64; 12]>,
196}
197
198#[derive(Debug, Clone)]
199pub struct SerializedAccountMetadata {
200    pub original_data_len: usize,
201    pub vm_data_addr: u64,
202    pub vm_key_addr: u64,
203    pub vm_kelvins_addr: u64,
204    pub vm_owner_addr: u64,
205}
206
207/// Main pipeline from runtime to program execution.
208pub struct InvokeContext<'a> {
209    /// Information about the currently executing transaction.
210    pub transaction_context: &'a mut TransactionContext,
211    /// The local program cache for the transaction batch.
212    pub program_cache_for_tx_batch: &'a mut ProgramCacheForTxBatch,
213    /// Runtime configurations used to provision the invocation environment.
214    pub environment_config: EnvironmentConfig<'a>,
215    /// The compute budget for the current invocation.
216    compute_budget: ComputeBudget,
217    /// Instruction compute meter, for tracking compute units consumed against
218    /// the designated compute budget during program execution.
219    compute_meter: RefCell<u64>,
220    log_collector: Option<Rc<RefCell<LogCollector>>>,
221    /// Latest measurement not yet accumulated in [ExecuteDetailsTimings::execute_us]
222    pub execute_time: Option<Measure>,
223    pub timings: ExecuteDetailsTimings,
224    pub syscall_context: Vec<Option<SyscallContext>>,
225    traces: Vec<Vec<[u64; 12]>>,
226}
227
228impl<'a> InvokeContext<'a> {
229    #[allow(clippy::too_many_arguments)]
230    pub fn new(
231        transaction_context: &'a mut TransactionContext,
232        program_cache_for_tx_batch: &'a mut ProgramCacheForTxBatch,
233        environment_config: EnvironmentConfig<'a>,
234        log_collector: Option<Rc<RefCell<LogCollector>>>,
235        compute_budget: ComputeBudget,
236    ) -> Self {
237        Self {
238            transaction_context,
239            program_cache_for_tx_batch,
240            environment_config,
241            log_collector,
242            compute_budget,
243            compute_meter: RefCell::new(compute_budget.compute_unit_limit),
244            execute_time: None,
245            timings: ExecuteDetailsTimings::default(),
246            syscall_context: Vec::new(),
247            traces: Vec::new(),
248        }
249    }
250
251    pub fn get_environments_for_slot(
252        &self,
253        effective_slot: Slot,
254    ) -> Result<&ProgramRuntimeEnvironments, InstructionError> {
255        let epoch_schedule = self.environment_config.sysvar_cache.get_epoch_schedule()?;
256        let epoch = epoch_schedule.get_epoch(effective_slot);
257        Ok(self
258            .program_cache_for_tx_batch
259            .get_environments_for_epoch(epoch))
260    }
261
262    /// Push a stack frame onto the invocation stack
263    pub fn push(&mut self) -> Result<(), InstructionError> {
264        let instruction_context = self
265            .transaction_context
266            .get_instruction_context_at_index_in_trace(
267                self.transaction_context.get_instruction_trace_length(),
268            )?;
269        let program_id = instruction_context
270            .get_last_program_key(self.transaction_context)
271            .map_err(|_| InstructionError::UnsupportedProgramId)?;
272        if self
273            .transaction_context
274            .get_instruction_context_stack_height()
275            != 0
276        {
277            let contains = (0..self
278                .transaction_context
279                .get_instruction_context_stack_height())
280                .any(|level| {
281                    self.transaction_context
282                        .get_instruction_context_at_nesting_level(level)
283                        .and_then(|instruction_context| {
284                            instruction_context
285                                .try_borrow_last_program_account(self.transaction_context)
286                        })
287                        .map(|program_account| program_account.get_key() == program_id)
288                        .unwrap_or(false)
289                });
290            let is_last = self
291                .transaction_context
292                .get_current_instruction_context()
293                .and_then(|instruction_context| {
294                    instruction_context.try_borrow_last_program_account(self.transaction_context)
295                })
296                .map(|program_account| program_account.get_key() == program_id)
297                .unwrap_or(false);
298            if contains && !is_last {
299                // Reentrancy not allowed unless caller is calling itself
300                return Err(InstructionError::ReentrancyNotAllowed);
301            }
302        }
303
304        self.syscall_context.push(None);
305        self.transaction_context.push()
306    }
307
308    /// Pop a stack frame from the invocation stack
309    fn pop(&mut self) -> Result<(), InstructionError> {
310        if let Some(Some(syscall_context)) = self.syscall_context.pop() {
311            self.traces.push(syscall_context.trace_log);
312        }
313        self.transaction_context.pop()
314    }
315
316    /// Current height of the invocation stack, top level instructions are height
317    /// `rialo_s_instruction::TRANSACTION_LEVEL_STACK_HEIGHT`
318    pub fn get_stack_height(&self) -> usize {
319        self.transaction_context
320            .get_instruction_context_stack_height()
321    }
322
323    /// Entrypoint for a cross-program invocation from a builtin program
324    pub fn native_invoke(
325        &mut self,
326        instruction: StableInstruction,
327        signers: &[Pubkey],
328    ) -> Result<(), InstructionError> {
329        let (instruction_accounts, program_indices) =
330            self.prepare_instruction(&instruction, signers)?;
331        let mut compute_units_consumed = 0;
332        self.process_instruction(
333            &instruction.data,
334            &instruction_accounts,
335            &program_indices,
336            &mut compute_units_consumed,
337            &mut ExecuteTimings::default(),
338        )?;
339        Ok(())
340    }
341
342    /// Helper to prepare for process_instruction()
343    #[allow(clippy::type_complexity)]
344    pub fn prepare_instruction(
345        &mut self,
346        instruction: &StableInstruction,
347        signers: &[Pubkey],
348    ) -> Result<(Vec<InstructionAccount>, Vec<IndexOfAccount>), InstructionError> {
349        self.prepare_instruction_inner(instruction.program_id, &instruction.accounts, signers)
350    }
351
352    /// Helper to prepare for process_instruction()
353    pub fn prepare_cpi_instruction(
354        &mut self,
355        program_id: Pubkey,
356        account_metas: &[AccountMeta],
357        signers: &[Pubkey],
358    ) -> Result<(Vec<InstructionAccount>, Vec<IndexOfAccount>), InstructionError> {
359        self.prepare_instruction_inner(program_id, account_metas, signers)
360    }
361
362    pub fn prepare_instruction_inner(
363        &mut self,
364        callee_program_id: Pubkey,
365        account_metas: &[AccountMeta],
366        signers: &[Pubkey],
367    ) -> Result<(Vec<InstructionAccount>, Vec<IndexOfAccount>), InstructionError> {
368        // Finds the index of each account in the instruction by its pubkey.
369        // Then normalizes / unifies the privileges of duplicate accounts.
370        // Note: This is an O(n^2) algorithm,
371        // but performed on a very small slice and requires no heap allocations.
372        let instruction_context = self.transaction_context.get_current_instruction_context()?;
373        let mut deduplicated_instruction_accounts: Vec<InstructionAccount> = Vec::new();
374        let mut duplicate_indicies = Vec::with_capacity(account_metas.len());
375        for (instruction_account_index, account_meta) in account_metas.iter().enumerate() {
376            let index_in_transaction = self
377                .transaction_context
378                .find_index_of_account(&account_meta.pubkey)
379                .ok_or_else(|| {
380                    ic_msg!(
381                        self,
382                        "Instruction references an unknown account {}",
383                        account_meta.pubkey,
384                    );
385                    InstructionError::MissingAccount
386                })?;
387            if let Some(duplicate_index) =
388                deduplicated_instruction_accounts
389                    .iter()
390                    .position(|instruction_account| {
391                        instruction_account.index_in_transaction == index_in_transaction
392                    })
393            {
394                duplicate_indicies.push(duplicate_index);
395                let instruction_account = deduplicated_instruction_accounts
396                    .get_mut(duplicate_index)
397                    .ok_or(InstructionError::NotEnoughAccountKeys)?;
398                instruction_account.is_signer |= account_meta.is_signer;
399                instruction_account.is_writable |= account_meta.is_writable;
400            } else {
401                let index_in_caller = instruction_context
402                    .find_index_of_instruction_account(
403                        self.transaction_context,
404                        &account_meta.pubkey,
405                    )
406                    .ok_or_else(|| {
407                        ic_msg!(
408                            self,
409                            "Instruction references an unknown account {}",
410                            account_meta.pubkey,
411                        );
412                        InstructionError::MissingAccount
413                    })?;
414                duplicate_indicies.push(deduplicated_instruction_accounts.len());
415                deduplicated_instruction_accounts.push(InstructionAccount {
416                    index_in_transaction,
417                    index_in_caller,
418                    index_in_callee: instruction_account_index as IndexOfAccount,
419                    is_signer: account_meta.is_signer,
420                    is_writable: account_meta.is_writable,
421                });
422            }
423        }
424        for instruction_account in deduplicated_instruction_accounts.iter() {
425            let borrowed_account = instruction_context.try_borrow_instruction_account(
426                self.transaction_context,
427                instruction_account.index_in_caller,
428            )?;
429
430            // Readonly in caller cannot become writable in callee
431            if instruction_account.is_writable && !borrowed_account.is_writable() {
432                ic_msg!(
433                    self,
434                    "{}'s writable privilege escalated",
435                    borrowed_account.get_key(),
436                );
437                return Err(InstructionError::PrivilegeEscalation);
438            }
439
440            // To be signed in the callee,
441            // it must be either signed in the caller or by the program
442            if instruction_account.is_signer
443                && !(borrowed_account.is_signer() || signers.contains(borrowed_account.get_key()))
444            {
445                ic_msg!(
446                    self,
447                    "{}'s signer privilege escalated",
448                    borrowed_account.get_key()
449                );
450                return Err(InstructionError::PrivilegeEscalation);
451            }
452        }
453        let instruction_accounts = duplicate_indicies
454            .into_iter()
455            .map(|duplicate_index| {
456                deduplicated_instruction_accounts
457                    .get(duplicate_index)
458                    .cloned()
459                    .ok_or(InstructionError::NotEnoughAccountKeys)
460            })
461            .collect::<Result<Vec<InstructionAccount>, InstructionError>>()?;
462
463        // Find and validate executables / program accounts
464        let program_account_index = if self
465            .get_feature_set()
466            .is_active(&lift_cpi_caller_restriction::id())
467        {
468            self.transaction_context
469                .find_index_of_program_account(&callee_program_id)
470                .ok_or_else(|| {
471                    ic_msg!(self, "Unknown program {}", callee_program_id);
472                    InstructionError::MissingAccount
473                })?
474        } else {
475            let program_account_index = instruction_context
476                .find_index_of_instruction_account(self.transaction_context, &callee_program_id)
477                .ok_or_else(|| {
478                    ic_msg!(self, "Unknown program {}", callee_program_id);
479                    InstructionError::MissingAccount
480                })?;
481            let borrowed_program_account = instruction_context
482                .try_borrow_instruction_account(self.transaction_context, program_account_index)?;
483            #[allow(deprecated)]
484            if !self
485                .get_feature_set()
486                .is_active(&remove_accounts_executable_flag_checks::id())
487                && !borrowed_program_account.is_executable()
488            {
489                ic_msg!(self, "Account {} is not executable", callee_program_id);
490                return Err(InstructionError::AccountNotExecutable);
491            }
492            borrowed_program_account.get_index_in_transaction()
493        };
494
495        Ok((instruction_accounts, vec![program_account_index]))
496    }
497
498    /// Processes an instruction and returns how many compute units were used
499    pub fn process_instruction(
500        &mut self,
501        instruction_data: &[u8],
502        instruction_accounts: &[InstructionAccount],
503        program_indices: &[IndexOfAccount],
504        compute_units_consumed: &mut u64,
505        timings: &mut ExecuteTimings,
506    ) -> Result<(), InstructionError> {
507        *compute_units_consumed = 0;
508        self.transaction_context
509            .get_next_instruction_context()?
510            .configure(program_indices, instruction_accounts, instruction_data);
511        self.push()?;
512        self.process_executable_chain(compute_units_consumed, timings)
513            // MUST pop if and only if `push` succeeded, independent of `result`.
514            // Thus, the `.and()` instead of an `.and_then()`.
515            .and(self.pop())
516    }
517
518    /// Processes a precompile instruction
519    pub fn process_precompile<'ix_data>(
520        &mut self,
521        precompile: &Precompile,
522        instruction_data: &[u8],
523        instruction_accounts: &[InstructionAccount],
524        program_indices: &[IndexOfAccount],
525        message_instruction_datas_iter: impl Iterator<Item = &'ix_data [u8]>,
526    ) -> Result<(), InstructionError> {
527        self.transaction_context
528            .get_next_instruction_context()?
529            .configure(program_indices, instruction_accounts, instruction_data);
530        self.push()?;
531
532        let feature_set = self.get_feature_set();
533        let move_precompile_verification_to_svm =
534            feature_set.is_active(&move_precompile_verification_to_svm::id());
535        if move_precompile_verification_to_svm {
536            let instruction_datas: Vec<_> = message_instruction_datas_iter.collect();
537            precompile
538                .verify(instruction_data, &instruction_datas, feature_set)
539                .map_err(InstructionError::from)
540                .and(self.pop())
541        } else {
542            self.pop()
543        }
544    }
545
546    /// Calls the instruction's program entrypoint method
547    fn process_executable_chain(
548        &mut self,
549        compute_units_consumed: &mut u64,
550        timings: &mut ExecuteTimings,
551    ) -> Result<(), InstructionError> {
552        let instruction_context = self.transaction_context.get_current_instruction_context()?;
553        let process_executable_chain_time = Measure::start("process_executable_chain_time");
554
555        let builtin_id = {
556            debug_assert!(instruction_context.get_number_of_program_accounts() <= 1);
557            let borrowed_root_account = instruction_context
558                .try_borrow_program_account(self.transaction_context, 0)
559                .map_err(|_| InstructionError::UnsupportedProgramId)?;
560            let owner_id = borrowed_root_account.get_owner();
561            if native_loader::check_id(owner_id) {
562                *borrowed_root_account.get_key()
563            } else {
564                *owner_id
565            }
566        };
567
568        // The Murmur3 hash value (used by RBPF) of the string "entrypoint"
569        const ENTRYPOINT_KEY: u32 = 0x71E3CF81;
570        let entry = self
571            .program_cache_for_tx_batch
572            .find(&builtin_id)
573            .ok_or(InstructionError::UnsupportedProgramId)?;
574        let function = match &entry.program {
575            ProgramCacheEntryType::Builtin(program) => program
576                .get_function_registry()
577                .lookup_by_key(ENTRYPOINT_KEY)
578                .map(|(_name, function)| function),
579            _ => None,
580        }
581        .ok_or(InstructionError::UnsupportedProgramId)?;
582        entry.ix_usage_counter.fetch_add(1, Ordering::Relaxed);
583
584        let program_id = *instruction_context.get_last_program_key(self.transaction_context)?;
585        self.transaction_context
586            .set_return_data(program_id, Vec::new())?;
587        let logger = self.get_log_collector();
588        stable_log::program_invoke(&logger, &program_id, self.get_stack_height());
589        let pre_remaining_units = self.get_remaining();
590        // In program-runtime v2 we will create this VM instance only once per transaction.
591        // `program_runtime_environment_v2.get_config()` will be used instead of `mock_config`.
592        // For now, only built-ins are invoked from here, so the VM and its Config are irrelevant.
593        let mock_config = Config::default();
594        let empty_memory_mapping =
595            MemoryMapping::new(Vec::new(), &mock_config, SBPFVersion::V0).unwrap();
596        let mut vm = EbpfVm::new(
597            self.program_cache_for_tx_batch
598                .environments
599                .program_runtime_v2
600                .clone(),
601            SBPFVersion::V0,
602            // Removes lifetime tracking
603            unsafe { std::mem::transmute::<&mut InvokeContext<'_>, &mut InvokeContext<'_>>(self) },
604            empty_memory_mapping,
605            0,
606        );
607        vm.invoke_function(function);
608        let result = match vm.program_result {
609            ProgramResult::Ok(_) => {
610                stable_log::program_success(&logger, &program_id);
611                Ok(())
612            }
613            ProgramResult::Err(ref err) => {
614                if let EbpfError::SyscallError(syscall_error) = err {
615                    if let Some(instruction_err) = syscall_error.downcast_ref::<InstructionError>()
616                    {
617                        stable_log::program_failure(&logger, &program_id, instruction_err);
618                        Err(instruction_err.clone())
619                    } else {
620                        stable_log::program_failure(&logger, &program_id, syscall_error);
621                        Err(InstructionError::ProgramFailedToComplete)
622                    }
623                } else {
624                    stable_log::program_failure(&logger, &program_id, err);
625                    Err(InstructionError::ProgramFailedToComplete)
626                }
627            }
628        };
629        let post_remaining_units = self.get_remaining();
630        *compute_units_consumed = pre_remaining_units.saturating_sub(post_remaining_units);
631
632        if builtin_id == program_id && result.is_ok() && *compute_units_consumed == 0 {
633            return Err(InstructionError::BuiltinProgramsMustConsumeComputeUnits);
634        }
635
636        timings
637            .execute_accessories
638            .process_instructions
639            .process_executable_chain_us += process_executable_chain_time.end_as_us();
640        result
641    }
642
643    /// Get this invocation's LogCollector
644    pub fn get_log_collector(&self) -> Option<Rc<RefCell<LogCollector>>> {
645        self.log_collector.clone()
646    }
647
648    /// Consume compute units
649    pub fn consume_checked(&self, amount: u64) -> Result<(), Box<dyn std::error::Error>> {
650        let mut compute_meter = self.compute_meter.borrow_mut();
651        let exceeded = *compute_meter < amount;
652        *compute_meter = compute_meter.saturating_sub(amount);
653        if exceeded {
654            return Err(Box::new(InstructionError::ComputationalBudgetExceeded));
655        }
656        Ok(())
657    }
658
659    /// Set compute units
660    ///
661    /// Only use for tests and benchmarks
662    pub fn mock_set_remaining(&self, remaining: u64) {
663        *self.compute_meter.borrow_mut() = remaining;
664    }
665
666    /// Get this invocation's compute budget
667    pub fn get_compute_budget(&self) -> &ComputeBudget {
668        &self.compute_budget
669    }
670
671    /// Get the current feature set.
672    pub fn get_feature_set(&self) -> &FeatureSet {
673        &self.environment_config.feature_set
674    }
675
676    /// Set feature set.
677    ///
678    /// Only use for tests and benchmarks.
679    pub fn mock_set_feature_set(&mut self, feature_set: Arc<FeatureSet>) {
680        self.environment_config.feature_set = feature_set;
681    }
682
683    /// Get cached sysvars
684    pub fn get_sysvar_cache(&self) -> &SysvarCache {
685        self.environment_config.sysvar_cache
686    }
687
688    /// Get cached epoch total stake.
689    pub fn get_epoch_total_stake(&self) -> u64 {
690        self.environment_config.epoch_total_stake
691    }
692
693    /// Get cached stake for the epoch vote account.
694    pub fn get_epoch_vote_account_stake(&self, pubkey: &'a Pubkey) -> u64 {
695        (self
696            .environment_config
697            .get_epoch_vote_account_stake_callback)(pubkey)
698    }
699
700    /// Get the CURRENT epoch (the epoch that is actively running).
701    ///
702    /// This is the epoch to use when checking stake activation/deactivation status,
703    /// unbonding periods, lockup periods, etc.
704    ///
705    /// # Implementation Note
706    /// The stake cache stores the NEXT epoch (the epoch being prepared for when
707    /// stake changes become effective). We subtract 1 to get the current epoch.
708    /// `saturating_sub` handles the edge case of epoch 0.
709    pub fn get_current_epoch(&self) -> u64 {
710        self.environment_config
711            .stakes_handle
712            .pending
713            .0
714            .read()
715            .expect("Failed to acquire lock")
716            .epoch
717            .saturating_sub(1)
718    }
719
720    /// Get the NEXT epoch (the epoch being prepared in the stake cache).
721    ///
722    /// This is the epoch that accumulated stake changes will become effective in
723    /// AFTER the next FreezeStakes call. Most instruction logic should use
724    /// `get_current_epoch()` instead; this method is primarily for internal use
725    /// and the FreezeStakes instruction.
726    pub fn get_next_epoch(&self) -> u64 {
727        self.environment_config
728            .stakes_handle
729            .pending
730            .0
731            .read()
732            .expect("Failed to acquire lock")
733            .epoch
734    }
735
736    /// Set the stake cache epoch for testing.
737    ///
738    /// Only use for tests and benchmarks.
739    pub fn mock_set_stake_cache_epoch(&mut self, epoch: u64) {
740        self.environment_config
741            .stakes_handle
742            .pending
743            .0
744            .write()
745            .expect("Failed to acquire lock")
746            .epoch = epoch;
747    }
748
749    /// Freeze the pending stake cache data and return the validator stakes
750    /// as a sorted `Vec<(Pubkey, u64)>`.
751    ///
752    /// This performs an O(1) clone of the stake cache data (via ImHashMap structural sharing)
753    /// and pushes it to the stake history. This is typically called by the ValidatorRegistry
754    /// program's FreezeStakes instruction to capture the validator set at a specific point.
755    /// It then extracts validator pubkeys and their stakes from the frozen snapshot.
756    pub fn freeze_stakes_and_get_validator_stakes(&self) -> Vec<(Pubkey, u64)> {
757        self.environment_config
758            .stakes_handle
759            .freeze_and_get_validator_stakes()
760    }
761
762    // Should alignment be enforced during user pointer translation
763    pub fn get_check_aligned(&self) -> bool {
764        self.transaction_context
765            .get_current_instruction_context()
766            .and_then(|instruction_context| {
767                let program_account =
768                    instruction_context.try_borrow_last_program_account(self.transaction_context);
769                debug_assert!(program_account.is_ok());
770                program_account
771            })
772            .map(|program_account| *program_account.get_owner() != bpf_loader_deprecated::id())
773            .unwrap_or(true)
774    }
775
776    // Set this instruction syscall context
777    pub fn set_syscall_context(
778        &mut self,
779        syscall_context: SyscallContext,
780    ) -> Result<(), InstructionError> {
781        *self
782            .syscall_context
783            .last_mut()
784            .ok_or(InstructionError::CallDepth)? = Some(syscall_context);
785        Ok(())
786    }
787
788    // Get this instruction's SyscallContext
789    pub fn get_syscall_context(&self) -> Result<&SyscallContext, InstructionError> {
790        self.syscall_context
791            .last()
792            .and_then(std::option::Option::as_ref)
793            .ok_or(InstructionError::CallDepth)
794    }
795
796    // Get this instruction's SyscallContext
797    pub fn get_syscall_context_mut(&mut self) -> Result<&mut SyscallContext, InstructionError> {
798        self.syscall_context
799            .last_mut()
800            .and_then(|syscall_context| syscall_context.as_mut())
801            .ok_or(InstructionError::CallDepth)
802    }
803
804    /// Return a references to traces
805    pub fn get_traces(&self) -> &Vec<Vec<[u64; 12]>> {
806        &self.traces
807    }
808}
809
810#[macro_export]
811macro_rules! with_mock_invoke_context {
812    (
813        $invoke_context:ident,
814        $transaction_context:ident,
815        $transaction_accounts:expr $(,)?
816    ) => {
817        use rialo_s_compute_budget::compute_budget::ComputeBudget;
818        use rialo_s_feature_set::FeatureSet;
819        use rialo_s_log_collector::LogCollector;
820        use rialo_s_type_overrides::sync::Arc;
821        use $crate::{
822            __private::{Hash, ReadableAccount, Rent, TransactionContext},
823            invoke_context::{EnvironmentConfig, InvokeContext},
824            loaded_programs::ProgramCacheForTxBatch,
825            sysvar_cache::SysvarCache,
826        };
827        let compute_budget = ComputeBudget::default();
828        let mut $transaction_context = TransactionContext::new(
829            $transaction_accounts,
830            Rent::default(),
831            compute_budget.max_instruction_stack_depth,
832            compute_budget.max_instruction_trace_length,
833        );
834        let mut sysvar_cache = SysvarCache::default();
835        sysvar_cache.fill_missing_entries(|pubkey, callback| {
836            for index in 0..$transaction_context.get_number_of_accounts() {
837                if $transaction_context
838                    .get_key_of_account_at_index(index)
839                    .unwrap()
840                    == pubkey
841                {
842                    callback(
843                        $transaction_context
844                            .get_account_at_index(index)
845                            .unwrap()
846                            .borrow()
847                            .data(),
848                    );
849                }
850            }
851        });
852        // Create a stakes handle with default epoch (0) - tests that need specific epochs
853        // should set up the stake cache data appropriately
854        let mock_stakes_handle = $crate::invoke_context::StakesHandle::default();
855        let environment_config = EnvironmentConfig::new(
856            Hash::default(),
857            0,
858            0,
859            &|_| 0,
860            Arc::new(FeatureSet::all_enabled()),
861            &sysvar_cache,
862            0,
863            mock_stakes_handle,
864        );
865        let mut program_cache_for_tx_batch = ProgramCacheForTxBatch::default();
866        let mut $invoke_context = InvokeContext::new(
867            &mut $transaction_context,
868            &mut program_cache_for_tx_batch,
869            environment_config,
870            Some(LogCollector::new_ref()),
871            compute_budget,
872        );
873    };
874}
875
876pub fn mock_process_instruction<
877    T: Into<StoredAccount>,
878    F: FnMut(&mut InvokeContext<'_>),
879    G: FnMut(&mut InvokeContext<'_>),
880>(
881    loader_id: &Pubkey,
882    mut program_indices: Vec<IndexOfAccount>,
883    instruction_data: &[u8],
884    transaction_accounts: Vec<(Pubkey, T)>,
885    instruction_account_metas: Vec<AccountMeta>,
886    expected_result: Result<(), InstructionError>,
887    builtin_function: BuiltinFunctionWithContext,
888    mut pre_adjustments: F,
889    mut post_adjustments: G,
890) -> Vec<StoredAccount> {
891    // Convert to StoredAccount upfront since we need to mutate the vector
892    let mut transaction_accounts: Vec<TransactionAccount> = transaction_accounts
893        .into_iter()
894        .map(|(key, account)| (key, account.into()))
895        .collect();
896    let mut instruction_accounts: Vec<InstructionAccount> =
897        Vec::with_capacity(instruction_account_metas.len());
898    for (instruction_account_index, account_meta) in instruction_account_metas.iter().enumerate() {
899        let index_in_transaction = transaction_accounts
900            .iter()
901            .position(|(key, _account)| *key == account_meta.pubkey)
902            .unwrap_or(transaction_accounts.len())
903            as IndexOfAccount;
904        let index_in_callee = instruction_accounts
905            .get(0..instruction_account_index)
906            .unwrap()
907            .iter()
908            .position(|instruction_account| {
909                instruction_account.index_in_transaction == index_in_transaction
910            })
911            .unwrap_or(instruction_account_index) as IndexOfAccount;
912        instruction_accounts.push(InstructionAccount {
913            index_in_transaction,
914            index_in_caller: index_in_transaction,
915            index_in_callee,
916            is_signer: account_meta.is_signer,
917            is_writable: account_meta.is_writable,
918        });
919    }
920    if program_indices.is_empty() {
921        program_indices.insert(0, transaction_accounts.len() as IndexOfAccount);
922        let processor_account =
923            StoredAccount::Data(AccountSharedData::new(0, 0, &native_loader::id()));
924        transaction_accounts.push((*loader_id, processor_account));
925    }
926    let pop_epoch_schedule_account = if !transaction_accounts
927        .iter()
928        .any(|(key, _)| *key == sysvar::epoch_schedule::id())
929    {
930        transaction_accounts.push((
931            sysvar::epoch_schedule::id(),
932            create_account_shared_data_for_test(&EpochSchedule::default()),
933        ));
934        true
935    } else {
936        false
937    };
938    with_mock_invoke_context!(invoke_context, transaction_context, transaction_accounts);
939    let mut program_cache_for_tx_batch = ProgramCacheForTxBatch::default();
940    program_cache_for_tx_batch.replenish(
941        *loader_id,
942        Arc::new(ProgramCacheEntry::new_builtin(0, 0, builtin_function)),
943    );
944    invoke_context.program_cache_for_tx_batch = &mut program_cache_for_tx_batch;
945    pre_adjustments(&mut invoke_context);
946    let result = invoke_context.process_instruction(
947        instruction_data,
948        &instruction_accounts,
949        &program_indices,
950        &mut 0,
951        &mut ExecuteTimings::default(),
952    );
953    assert_eq!(result, expected_result);
954    post_adjustments(&mut invoke_context);
955    let mut transaction_accounts = transaction_context.deconstruct_without_keys().unwrap();
956    if pop_epoch_schedule_account {
957        transaction_accounts.pop();
958    }
959    transaction_accounts.pop();
960    transaction_accounts
961}
962
963#[cfg(test)]
964mod tests {
965    use rialo_s_compute_budget::compute_budget_limits;
966    use rialo_s_instruction::Instruction;
967    use rialo_s_rent::Rent;
968    use serde::{Deserialize, Serialize};
969
970    use super::*;
971
972    #[derive(Debug, Serialize, Deserialize)]
973    enum MockInstruction {
974        NoopSuccess,
975        NoopFail,
976        ModifyOwned,
977        ModifyNotOwned,
978        ModifyReadonly,
979        UnbalancedPush,
980        UnbalancedPop,
981        ConsumeComputeUnits {
982            compute_units_to_consume: u64,
983            desired_result: Result<(), InstructionError>,
984        },
985        Resize {
986            new_len: u64,
987        },
988    }
989
990    const MOCK_BUILTIN_COMPUTE_UNIT_COST: u64 = 1;
991
992    declare_process_instruction!(
993        MockBuiltin,
994        MOCK_BUILTIN_COMPUTE_UNIT_COST,
995        |invoke_context| {
996            let transaction_context = &invoke_context.transaction_context;
997            let instruction_context = transaction_context.get_current_instruction_context()?;
998            let instruction_data = instruction_context.get_instruction_data();
999            let program_id = instruction_context.get_last_program_key(transaction_context)?;
1000            let instruction_accounts = (0..4)
1001                .map(|instruction_account_index| InstructionAccount {
1002                    index_in_transaction: instruction_account_index,
1003                    index_in_caller: instruction_account_index,
1004                    index_in_callee: instruction_account_index,
1005                    is_signer: false,
1006                    is_writable: false,
1007                })
1008                .collect::<Vec<_>>();
1009            assert_eq!(
1010                program_id,
1011                instruction_context
1012                    .try_borrow_instruction_account(transaction_context, 0)?
1013                    .get_owner()
1014            );
1015            assert_ne!(
1016                instruction_context
1017                    .try_borrow_instruction_account(transaction_context, 1)?
1018                    .get_owner(),
1019                instruction_context
1020                    .try_borrow_instruction_account(transaction_context, 0)?
1021                    .get_key()
1022            );
1023
1024            if let Ok(instruction) = bincode::deserialize(instruction_data) {
1025                match instruction {
1026                    MockInstruction::NoopSuccess => (),
1027                    MockInstruction::NoopFail => return Err(InstructionError::GenericError),
1028                    MockInstruction::ModifyOwned => instruction_context
1029                        .try_borrow_instruction_account(transaction_context, 0)?
1030                        .set_data_from_slice(&[1])?,
1031                    MockInstruction::ModifyNotOwned => instruction_context
1032                        .try_borrow_instruction_account(transaction_context, 1)?
1033                        .set_data_from_slice(&[1])?,
1034                    MockInstruction::ModifyReadonly => instruction_context
1035                        .try_borrow_instruction_account(transaction_context, 2)?
1036                        .set_data_from_slice(&[1])?,
1037                    MockInstruction::UnbalancedPush => {
1038                        instruction_context
1039                            .try_borrow_instruction_account(transaction_context, 0)?
1040                            .checked_add_kelvins(1)?;
1041                        let program_id = *transaction_context.get_key_of_account_at_index(3)?;
1042                        let metas = vec![
1043                            AccountMeta::new_readonly(
1044                                *transaction_context.get_key_of_account_at_index(0)?,
1045                                false,
1046                            ),
1047                            AccountMeta::new_readonly(
1048                                *transaction_context.get_key_of_account_at_index(1)?,
1049                                false,
1050                            ),
1051                        ];
1052                        let inner_instruction = Instruction::new_with_bincode(
1053                            program_id,
1054                            &MockInstruction::NoopSuccess,
1055                            metas,
1056                        );
1057                        invoke_context
1058                            .transaction_context
1059                            .get_next_instruction_context()
1060                            .unwrap()
1061                            .configure(&[3], &instruction_accounts, &[]);
1062                        let result = invoke_context.push();
1063                        assert_eq!(result, Err(InstructionError::UnbalancedInstruction));
1064                        result?;
1065                        invoke_context
1066                            .native_invoke(inner_instruction.into(), &[])
1067                            .and(invoke_context.pop())?;
1068                    }
1069                    MockInstruction::UnbalancedPop => instruction_context
1070                        .try_borrow_instruction_account(transaction_context, 0)?
1071                        .checked_add_kelvins(1)?,
1072                    MockInstruction::ConsumeComputeUnits {
1073                        compute_units_to_consume,
1074                        desired_result,
1075                    } => {
1076                        invoke_context
1077                            .consume_checked(compute_units_to_consume)
1078                            .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
1079                        return desired_result;
1080                    }
1081                    MockInstruction::Resize { new_len } => instruction_context
1082                        .try_borrow_instruction_account(transaction_context, 0)?
1083                        .set_data(vec![0; new_len as usize])?,
1084                }
1085            } else {
1086                return Err(InstructionError::InvalidInstructionData);
1087            }
1088            Ok(())
1089        }
1090    );
1091
1092    #[test]
1093    fn test_instruction_stack_height() {
1094        let one_more_than_max_depth = ComputeBudget::default()
1095            .max_instruction_stack_depth
1096            .saturating_add(1);
1097        let mut invoke_stack = vec![];
1098        let mut transaction_accounts = vec![];
1099        let mut instruction_accounts = vec![];
1100        for index in 0..one_more_than_max_depth {
1101            invoke_stack.push(rialo_s_pubkey::new_rand());
1102            transaction_accounts.push((
1103                rialo_s_pubkey::new_rand(),
1104                AccountSharedData::new(index as u64, 1, invoke_stack.get(index).unwrap()),
1105            ));
1106            instruction_accounts.push(InstructionAccount {
1107                index_in_transaction: index as IndexOfAccount,
1108                index_in_caller: index as IndexOfAccount,
1109                index_in_callee: instruction_accounts.len() as IndexOfAccount,
1110                is_signer: false,
1111                is_writable: true,
1112            });
1113        }
1114        for (index, program_id) in invoke_stack.iter().enumerate() {
1115            transaction_accounts.push((
1116                *program_id,
1117                AccountSharedData::new(1, 1, &rialo_s_pubkey::Pubkey::default()),
1118            ));
1119            instruction_accounts.push(InstructionAccount {
1120                index_in_transaction: index as IndexOfAccount,
1121                index_in_caller: index as IndexOfAccount,
1122                index_in_callee: index as IndexOfAccount,
1123                is_signer: false,
1124                is_writable: false,
1125            });
1126        }
1127        let transaction_accounts: Vec<(Pubkey, StoredAccount)> = transaction_accounts
1128            .into_iter()
1129            .map(|(k, v)| (k, v.into()))
1130            .collect();
1131        with_mock_invoke_context!(invoke_context, transaction_context, transaction_accounts);
1132
1133        // Check call depth increases and has a limit
1134        let mut depth_reached = 0;
1135        for _ in 0..invoke_stack.len() {
1136            invoke_context
1137                .transaction_context
1138                .get_next_instruction_context()
1139                .unwrap()
1140                .configure(
1141                    &[one_more_than_max_depth.saturating_add(depth_reached) as IndexOfAccount],
1142                    &instruction_accounts,
1143                    &[],
1144                );
1145            if Err(InstructionError::CallDepth) == invoke_context.push() {
1146                break;
1147            }
1148            depth_reached = depth_reached.saturating_add(1);
1149        }
1150        assert_ne!(depth_reached, 0);
1151        assert!(depth_reached < one_more_than_max_depth);
1152    }
1153
1154    #[test]
1155    fn test_max_instruction_trace_length() {
1156        const MAX_INSTRUCTIONS: usize = 8;
1157        let mut transaction_context = TransactionContext::new::<StoredAccount>(
1158            Vec::new(),
1159            Rent::default(),
1160            1,
1161            MAX_INSTRUCTIONS,
1162        );
1163        for _ in 0..MAX_INSTRUCTIONS {
1164            transaction_context.push().unwrap();
1165            transaction_context.pop().unwrap();
1166        }
1167        assert_eq!(
1168            transaction_context.push(),
1169            Err(InstructionError::MaxInstructionTraceLengthExceeded)
1170        );
1171    }
1172
1173    #[test]
1174    fn test_process_instruction() {
1175        let callee_program_id = rialo_s_pubkey::new_rand();
1176        let owned_account = AccountSharedData::new(42, 1, &callee_program_id);
1177        let not_owned_account = AccountSharedData::new(84, 1, &rialo_s_pubkey::new_rand());
1178        let readonly_account = AccountSharedData::new(168, 1, &rialo_s_pubkey::new_rand());
1179        let loader_account = AccountSharedData::new(0, 1, &native_loader::id());
1180        let program_account = AccountSharedData::new(1, 1, &native_loader::id());
1181        let transaction_accounts = vec![
1182            (rialo_s_pubkey::new_rand(), owned_account),
1183            (rialo_s_pubkey::new_rand(), not_owned_account),
1184            (rialo_s_pubkey::new_rand(), readonly_account),
1185            (callee_program_id, program_account),
1186            (rialo_s_pubkey::new_rand(), loader_account),
1187        ];
1188        let metas = vec![
1189            AccountMeta::new(transaction_accounts.first().unwrap().0, false),
1190            AccountMeta::new(transaction_accounts.get(1).unwrap().0, false),
1191            AccountMeta::new_readonly(transaction_accounts.get(2).unwrap().0, false),
1192        ];
1193        let instruction_accounts = (0..4)
1194            .map(|instruction_account_index| InstructionAccount {
1195                index_in_transaction: instruction_account_index,
1196                index_in_caller: instruction_account_index,
1197                index_in_callee: instruction_account_index,
1198                is_signer: false,
1199                is_writable: instruction_account_index < 2,
1200            })
1201            .collect::<Vec<_>>();
1202        let transaction_accounts: Vec<(Pubkey, StoredAccount)> = transaction_accounts
1203            .into_iter()
1204            .map(|(k, v)| (k, v.into()))
1205            .collect();
1206        with_mock_invoke_context!(invoke_context, transaction_context, transaction_accounts);
1207        let mut program_cache_for_tx_batch = ProgramCacheForTxBatch::default();
1208        program_cache_for_tx_batch.replenish(
1209            callee_program_id,
1210            Arc::new(ProgramCacheEntry::new_builtin(0, 1, MockBuiltin::vm)),
1211        );
1212        invoke_context.program_cache_for_tx_batch = &mut program_cache_for_tx_batch;
1213
1214        // Account modification tests
1215        let cases = vec![
1216            (MockInstruction::NoopSuccess, Ok(())),
1217            (
1218                MockInstruction::NoopFail,
1219                Err(InstructionError::GenericError),
1220            ),
1221            (MockInstruction::ModifyOwned, Ok(())),
1222            (
1223                MockInstruction::ModifyNotOwned,
1224                Err(InstructionError::ExternalAccountDataModified),
1225            ),
1226            (
1227                MockInstruction::ModifyReadonly,
1228                Err(InstructionError::ReadonlyDataModified),
1229            ),
1230            (
1231                MockInstruction::UnbalancedPush,
1232                Err(InstructionError::UnbalancedInstruction),
1233            ),
1234            (
1235                MockInstruction::UnbalancedPop,
1236                Err(InstructionError::UnbalancedInstruction),
1237            ),
1238        ];
1239        for case in cases {
1240            invoke_context
1241                .transaction_context
1242                .get_next_instruction_context()
1243                .unwrap()
1244                .configure(&[4], &instruction_accounts, &[]);
1245            invoke_context.push().unwrap();
1246            let inner_instruction =
1247                Instruction::new_with_bincode(callee_program_id, &case.0, metas.clone());
1248            let result = invoke_context
1249                .native_invoke(inner_instruction.into(), &[])
1250                .and(invoke_context.pop());
1251            assert_eq!(result, case.1);
1252        }
1253
1254        // Compute unit consumption tests
1255        let compute_units_to_consume = 10;
1256        let expected_results = vec![Ok(()), Err(InstructionError::GenericError)];
1257        for expected_result in expected_results {
1258            invoke_context
1259                .transaction_context
1260                .get_next_instruction_context()
1261                .unwrap()
1262                .configure(&[4], &instruction_accounts, &[]);
1263            invoke_context.push().unwrap();
1264            let inner_instruction = Instruction::new_with_bincode(
1265                callee_program_id,
1266                &MockInstruction::ConsumeComputeUnits {
1267                    compute_units_to_consume,
1268                    desired_result: expected_result.clone(),
1269                },
1270                metas.clone(),
1271            );
1272            let inner_instruction = StableInstruction::from(inner_instruction);
1273            let (inner_instruction_accounts, program_indices) = invoke_context
1274                .prepare_instruction(&inner_instruction, &[])
1275                .unwrap();
1276
1277            let mut compute_units_consumed = 0;
1278            let result = invoke_context.process_instruction(
1279                &inner_instruction.data,
1280                &inner_instruction_accounts,
1281                &program_indices,
1282                &mut compute_units_consumed,
1283                &mut ExecuteTimings::default(),
1284            );
1285
1286            // Because the instruction had compute cost > 0, then regardless of the execution result,
1287            // the number of compute units consumed should be a non-default which is something greater
1288            // than zero.
1289            assert!(compute_units_consumed > 0);
1290            assert_eq!(
1291                compute_units_consumed,
1292                compute_units_to_consume.saturating_add(MOCK_BUILTIN_COMPUTE_UNIT_COST),
1293            );
1294            assert_eq!(result, expected_result);
1295
1296            invoke_context.pop().unwrap();
1297        }
1298    }
1299
1300    #[test]
1301    fn test_invoke_context_compute_budget() {
1302        let transaction_accounts = vec![(rialo_s_pubkey::new_rand(), AccountSharedData::default())];
1303        let transaction_accounts: Vec<(Pubkey, StoredAccount)> = transaction_accounts
1304            .into_iter()
1305            .map(|(k, v)| (k, v.into()))
1306            .collect();
1307
1308        with_mock_invoke_context!(invoke_context, transaction_context, transaction_accounts);
1309        invoke_context.compute_budget = ComputeBudget::new(
1310            compute_budget_limits::DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT as u64,
1311        );
1312
1313        invoke_context
1314            .transaction_context
1315            .get_next_instruction_context()
1316            .unwrap()
1317            .configure(&[0], &[], &[]);
1318        invoke_context.push().unwrap();
1319        assert_eq!(
1320            *invoke_context.get_compute_budget(),
1321            ComputeBudget::new(
1322                compute_budget_limits::DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT as u64
1323            )
1324        );
1325        invoke_context.pop().unwrap();
1326    }
1327
1328    #[test]
1329    fn test_process_instruction_accounts_resize_delta() {
1330        let program_key = Pubkey::new_unique();
1331        let user_account_data_len = 123u64;
1332        let user_account =
1333            AccountSharedData::new(100, user_account_data_len as usize, &program_key);
1334        let dummy_account = AccountSharedData::new(10, 0, &program_key);
1335        let program_account = AccountSharedData::new(500, 500, &native_loader::id());
1336        let transaction_accounts = vec![
1337            (Pubkey::new_unique(), user_account),
1338            (Pubkey::new_unique(), dummy_account),
1339            (program_key, program_account),
1340        ];
1341        let instruction_accounts = [
1342            InstructionAccount {
1343                index_in_transaction: 0,
1344                index_in_caller: 0,
1345                index_in_callee: 0,
1346                is_signer: false,
1347                is_writable: true,
1348            },
1349            InstructionAccount {
1350                index_in_transaction: 1,
1351                index_in_caller: 1,
1352                index_in_callee: 1,
1353                is_signer: false,
1354                is_writable: false,
1355            },
1356        ];
1357        let transaction_accounts: Vec<(Pubkey, StoredAccount)> = transaction_accounts
1358            .into_iter()
1359            .map(|(k, v)| (k, v.into()))
1360            .collect();
1361        with_mock_invoke_context!(invoke_context, transaction_context, transaction_accounts);
1362        let mut program_cache_for_tx_batch = ProgramCacheForTxBatch::default();
1363        program_cache_for_tx_batch.replenish(
1364            program_key,
1365            Arc::new(ProgramCacheEntry::new_builtin(0, 0, MockBuiltin::vm)),
1366        );
1367        invoke_context.program_cache_for_tx_batch = &mut program_cache_for_tx_batch;
1368
1369        // Test: Resize the account to *the same size*, so not consuming any additional size; this must succeed
1370        {
1371            let resize_delta: i64 = 0;
1372            let new_len = (user_account_data_len as i64).saturating_add(resize_delta) as u64;
1373            let instruction_data =
1374                bincode::serialize(&MockInstruction::Resize { new_len }).unwrap();
1375
1376            let result = invoke_context.process_instruction(
1377                &instruction_data,
1378                &instruction_accounts,
1379                &[2],
1380                &mut 0,
1381                &mut ExecuteTimings::default(),
1382            );
1383
1384            assert!(result.is_ok());
1385            assert_eq!(
1386                invoke_context
1387                    .transaction_context
1388                    .accounts_resize_delta()
1389                    .unwrap(),
1390                resize_delta
1391            );
1392        }
1393
1394        // Test: Resize the account larger; this must succeed
1395        {
1396            let resize_delta: i64 = 1;
1397            let new_len = (user_account_data_len as i64).saturating_add(resize_delta) as u64;
1398            let instruction_data =
1399                bincode::serialize(&MockInstruction::Resize { new_len }).unwrap();
1400
1401            let result = invoke_context.process_instruction(
1402                &instruction_data,
1403                &instruction_accounts,
1404                &[2],
1405                &mut 0,
1406                &mut ExecuteTimings::default(),
1407            );
1408
1409            assert!(result.is_ok());
1410            assert_eq!(
1411                invoke_context
1412                    .transaction_context
1413                    .accounts_resize_delta()
1414                    .unwrap(),
1415                resize_delta
1416            );
1417        }
1418
1419        // Test: Resize the account smaller; this must succeed
1420        {
1421            let resize_delta: i64 = -1;
1422            let new_len = (user_account_data_len as i64).saturating_add(resize_delta) as u64;
1423            let instruction_data =
1424                bincode::serialize(&MockInstruction::Resize { new_len }).unwrap();
1425
1426            let result = invoke_context.process_instruction(
1427                &instruction_data,
1428                &instruction_accounts,
1429                &[2],
1430                &mut 0,
1431                &mut ExecuteTimings::default(),
1432            );
1433
1434            assert!(result.is_ok());
1435            assert_eq!(
1436                invoke_context
1437                    .transaction_context
1438                    .accounts_resize_delta()
1439                    .unwrap(),
1440                resize_delta
1441            );
1442        }
1443    }
1444}