feat: add custom syscall registration

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<InvokeContext<'static, 'static>>,
+    ) -> 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> {

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]

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"))',
+]

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(())
+}

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<u64, Box<dyn std::error::Error>> {
+        assert_eq!(to_burn, CUS_TO_BURN);
+        invoke_context.consume_checked(to_burn)?;
+        Ok(0)
+    }
+);
+
+fn read_custom_syscall_program() -> Vec<u8> {
+    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());
+}