diff --git a/crates/litesvm/src/lib.rs b/crates/litesvm/src/lib.rs index dbb3cb7f5..689dc7b36 100644 --- a/crates/litesvm/src/lib.rs +++ b/crates/litesvm/src/lib.rs @@ -343,6 +343,7 @@ use { solana_program_runtime::{ invoke_context::{BuiltinFunctionWithContext, EnvironmentConfig, InvokeContext}, loaded_programs::{LoadProgramMetrics, ProgramCacheEntry}, + solana_sbpf::program::BuiltinFunction, }, solana_rent::Rent, solana_sdk_ids::{ @@ -1622,6 +1623,45 @@ impl LiteSVM { ) { self.invocation_inspect_callback = Arc::new(callback); } + + /// Registers a custom syscall in both program runtime environments (v1 and v2). + /// + /// **Must be called after `with_builtins()`** (which recreates the environments + /// from scratch) and **before `with_default_programs()`** (which clones the + /// environment Arcs into program cache entries, preventing further mutation). + /// + /// Panics if the runtime environments cannot be mutated or if registration + /// fails. This is intentional — a misconfigured syscall should fail loudly + /// rather than silently. + pub fn with_custom_syscall( + mut self, + name: &str, + syscall: BuiltinFunction>, + ) -> Self { + let (Some(program_runtime_v1), Some(program_runtime_v2)) = ( + Arc::get_mut(&mut self.accounts.environments.program_runtime_v1), + Arc::get_mut(&mut self.accounts.environments.program_runtime_v2), + ) else { + panic!("with_custom_syscall: can't mutate program runtimes"); + }; + + // Once unregister_function is available, users could replace existing built-in + // syscalls. + + // TODO: uncomment once https://github.com/anza-xyz/sbpf/pull/153 is available. + // let _ = program_runtime_v1.unregister_function(name); + program_runtime_v1 + .register_function(name, syscall) + .unwrap_or_else(|e| panic!("failed to register syscall '{name}' in runtime_v1: {e}")); + + // TODO: uncomment once https://github.com/anza-xyz/sbpf/pull/153 is available. + // let _ = program_runtime_v2.unregister_function(name); + program_runtime_v2 + .register_function(name, syscall) + .unwrap_or_else(|e| panic!("failed to register syscall '{name}' in runtime_v2: {e}")); + + self + } } struct CheckAndProcessTransactionSuccessCore<'ix_data> { diff --git a/crates/litesvm/test_programs/Cargo.lock b/crates/litesvm/test_programs/Cargo.lock index bff0a9930..1e8ab292b 100644 --- a/crates/litesvm/test_programs/Cargo.lock +++ b/crates/litesvm/test_programs/Cargo.lock @@ -750,6 +750,16 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "test-program-custom-syscall" +version = "0.1.0" +dependencies = [ + "solana-account-info", + "solana-program-entrypoint", + "solana-program-error", + "solana-pubkey 3.0.0", +] + [[package]] name = "tinyvec" version = "1.10.0" diff --git a/crates/litesvm/test_programs/Cargo.toml b/crates/litesvm/test_programs/Cargo.toml index 6c690b7ba..188481e6b 100644 --- a/crates/litesvm/test_programs/Cargo.toml +++ b/crates/litesvm/test_programs/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["clock-example", "counter", "failure"] +members = ["clock-example", "counter", "failure", "custom-syscall"] resolver = "2" [workspace.dependencies] diff --git a/crates/litesvm/test_programs/custom-syscall/Cargo.toml b/crates/litesvm/test_programs/custom-syscall/Cargo.toml new file mode 100644 index 000000000..3396c1397 --- /dev/null +++ b/crates/litesvm/test_programs/custom-syscall/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "test-program-custom-syscall" +version = "0.1.0" +edition = "2021" + +[dependencies] +solana-account-info = { workspace = true } +solana-program-entrypoint = { workspace = true } +solana-program-error = { workspace = true } +solana-pubkey = { workspace = true } + +[lib] +crate-type = ["cdylib", "lib"] + +[lints.rust.unexpected_cfgs] +level = "warn" +check-cfg = [ + 'cfg(feature, values("custom-heap", "custom-panic"))', + 'cfg(target_os, values("solana"))', +] diff --git a/crates/litesvm/test_programs/custom-syscall/src/lib.rs b/crates/litesvm/test_programs/custom-syscall/src/lib.rs new file mode 100644 index 000000000..a4d8cf9c0 --- /dev/null +++ b/crates/litesvm/test_programs/custom-syscall/src/lib.rs @@ -0,0 +1,29 @@ +#![cfg(target_os = "solana")] + +use {solana_account_info::AccountInfo, solana_program_error::ProgramError, solana_pubkey::Pubkey}; + +// Declare the custom syscall that we expect to be registered. +// This matches the `sol_burn_cus` syscall from the test. +extern "C" { + fn sol_burn_cus(to_burn: u64) -> u64; +} + +solana_program_entrypoint::entrypoint!(process_instruction); + +fn process_instruction( + _program_id: &Pubkey, + _accounts: &[AccountInfo], + input: &[u8], +) -> Result<(), ProgramError> { + let to_burn = input + .get(0..8) + .and_then(|bytes| bytes.try_into().map(u64::from_le_bytes).ok()) + .ok_or(ProgramError::InvalidInstructionData)?; + + // Call the custom syscall to burn CUs. + unsafe { + sol_burn_cus(to_burn); + } + + Ok(()) +} diff --git a/crates/litesvm/tests/custom_syscall.rs b/crates/litesvm/tests/custom_syscall.rs new file mode 100644 index 000000000..d47276b1e --- /dev/null +++ b/crates/litesvm/tests/custom_syscall.rs @@ -0,0 +1,77 @@ +use { + agave_feature_set::FeatureSet, + litesvm::LiteSVM, + solana_address::address, + solana_keypair::Keypair, + solana_message::{Instruction, Message}, + solana_native_token::LAMPORTS_PER_SOL, + solana_program_runtime::{ + invoke_context::InvokeContext, + solana_sbpf::{declare_builtin_function, memory_region::MemoryMapping}, + }, + solana_signer::Signer, + solana_transaction::Transaction, + std::path::PathBuf, +}; + +const CUS_TO_BURN: u64 = 1234; + +declare_builtin_function!( + /// A custom syscall to burn CUs. + SyscallBurnCus, + fn rust( + invoke_context: &mut InvokeContext, + to_burn: u64, + _arg2: u64, + _arg3: u64, + _arg4: u64, + _arg5: u64, + _memory_mapping: &mut MemoryMapping, + ) -> Result> { + assert_eq!(to_burn, CUS_TO_BURN); + invoke_context.consume_checked(to_burn)?; + Ok(0) + } +); + +fn read_custom_syscall_program() -> Vec { + let mut so_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + so_path.push("test_programs/target/deploy/test_program_custom_syscall.so"); + std::fs::read(so_path).unwrap() +} + +fn litesvm_ctor() -> LiteSVM { + LiteSVM::default() + .with_feature_set(FeatureSet::all_enabled()) + .with_builtins() + .with_custom_syscall("sol_burn_cus", SyscallBurnCus::vm) + .with_lamports(1_000_000u64.wrapping_mul(LAMPORTS_PER_SOL)) + .with_sysvars() + .with_default_programs() + .with_sigverify(true) + .with_blockhash_check(true) +} + +#[test] +pub fn test_custom_syscall() { + let mut svm = litesvm_ctor(); + let payer_kp = Keypair::new(); + let payer_pk = payer_kp.pubkey(); + let program_id = address!("GtdambwDgHWrDJdVPBkEHGhCwokqgAoch162teUjJse2"); + svm.add_program(program_id, &read_custom_syscall_program()) + .unwrap(); + svm.airdrop(&payer_pk, 1000000000).unwrap(); + let blockhash = svm.latest_blockhash(); + let msg = Message::new_with_blockhash( + &[Instruction { + program_id, + accounts: vec![], + data: CUS_TO_BURN.to_le_bytes().to_vec(), + }], + Some(&payer_pk), + &blockhash, + ); + let tx = Transaction::new(&[payer_kp], msg, blockhash); + let res = svm.send_transaction(tx); + assert!(res.is_ok(), "custom syscall tx failed: {:?}", res.err()); +}