Support registering custom syscalls (#126)

Author: buffalojoec

Cargo.lock

@@ -29,6 +29,7 @@ audit:
 # Build test programs
 build-test-programs:
 	@cargo build-sbf --manifest-path test-programs/cpi-target/Cargo.toml
+	@cargo build-sbf --manifest-path test-programs/custom-syscall/Cargo.toml
 	@cargo build-sbf --manifest-path test-programs/primary/Cargo.toml
 
 # Pre-publish checks
@@ -84,6 +85,7 @@ check-features:
 
 # Run tests
 test:
+	@$(MAKE) build-test-programs
 	@cargo test --all-features
 
 # Run all checks in sequence

Makefile

@@ -503,6 +503,7 @@ impl Default for Mollusk {
              solana_runtime::message_processor=debug,\
              solana_runtime::system_instruction_processor=trace",
         );
+        let compute_budget = ComputeBudget::default();
         #[cfg(feature = "fuzz")]
         let feature_set = {
             // Omit "test features" (they have the same u64 ID).
@@ -515,12 +516,13 @@ impl Default for Mollusk {
         };
         #[cfg(not(feature = "fuzz"))]
         let feature_set = FeatureSet::all_enabled();
+        let program_cache = ProgramCache::new(&feature_set, &compute_budget);
         Self {
             config: Config::default(),
-            compute_budget: ComputeBudget::default(),
+            compute_budget,
             feature_set,
             logger: None,
-            program_cache: ProgramCache::default(),
+            program_cache,
             sysvars: Sysvars::default(),
             #[cfg(feature = "fuzz-fd")]
             slot: 0,
@@ -573,13 +575,7 @@ impl Mollusk {
         elf: &[u8],
         loader_key: &Pubkey,
     ) {
-        self.program_cache.add_program(
-            program_id,
-            loader_key,
-            elf,
-            &self.compute_budget,
-            &self.feature_set,
-        );
+        self.program_cache.add_program(program_id, loader_key, elf);
     }
 
     /// Warp the test environment to a slot by updating sysvars.

harness/src/lib.rs

@@ -8,8 +8,12 @@ use {
     solana_loader_v3_interface::state::UpgradeableLoaderState,
     solana_loader_v4_interface::state::{LoaderV4State, LoaderV4Status},
     solana_program_runtime::{
-        invoke_context::BuiltinFunctionWithContext,
+        invoke_context::{BuiltinFunctionWithContext, InvokeContext},
         loaded_programs::{LoadProgramMetrics, ProgramCacheEntry, ProgramCacheForTxBatch},
+        solana_sbpf::{
+            elf::ElfError,
+            program::{BuiltinFunction, BuiltinProgram},
+        },
     },
     solana_pubkey::Pubkey,
     solana_rent::Rent,
@@ -57,13 +61,23 @@ pub struct ProgramCache {
     //
     // K: program ID, V: loader key
     entries_cache: Rc<RefCell<HashMap<Pubkey, Pubkey>>>,
+    // The function registry (syscalls) to use for verifying and loading
+    // program ELFs.
+    program_runtime_environment: BuiltinProgram<InvokeContext<'static>>,
 }
 
-impl Default for ProgramCache {
-    fn default() -> Self {
+impl ProgramCache {
+    pub fn new(feature_set: &FeatureSet, compute_budget: &ComputeBudget) -> Self {
         let me = Self {
             cache: Rc::new(RefCell::new(ProgramCacheForTxBatch::default())),
             entries_cache: Rc::new(RefCell::new(HashMap::new())),
+            program_runtime_environment: create_program_runtime_environment_v1(
+                feature_set,
+                compute_budget,
+                /* reject_deployment_of_broken_elfs */ false,
+                /* debugging_features */ false,
+            )
+            .unwrap(),
         };
         BUILTINS.iter().for_each(|builtin| {
             let program_id = builtin.program_id;
@@ -72,9 +86,7 @@ impl Default for ProgramCache {
         });
         me
     }
-}
 
-impl ProgramCache {
     pub(crate) fn cache(&self) -> RefMut<ProgramCacheForTxBatch> {
         self.cache.borrow_mut()
     }
@@ -94,18 +106,24 @@ impl ProgramCache {
     }
 
     /// Add a program to the cache.
-    pub fn add_program(
-        &mut self,
-        program_id: &Pubkey,
-        loader_key: &Pubkey,
-        elf: &[u8],
-        compute_budget: &ComputeBudget,
-        feature_set: &FeatureSet,
-    ) {
-        let environment = Arc::new(
-            create_program_runtime_environment_v1(feature_set, compute_budget, false, false)
-                .unwrap(),
-        );
+    pub fn add_program(&mut self, program_id: &Pubkey, loader_key: &Pubkey, elf: &[u8]) {
+        // This might look rough, but it's actually functionally the same as
+        // calling `create_program_runtime_environment_v1` on every addition.
+        let environment = {
+            let config = self.program_runtime_environment.get_config().clone();
+            let mut loader = BuiltinProgram::new_loader(config);
+
+            for (_key, (name, value)) in self
+                .program_runtime_environment
+                .get_function_registry()
+                .iter()
+            {
+                let name = std::str::from_utf8(name).unwrap();
+                loader.register_function(name, value).unwrap();
+            }
+
+            Arc::new(loader)
+        };
         self.replenish(
             *program_id,
             Arc::new(
@@ -149,6 +167,19 @@ impl ProgramCache {
             })
             .collect()
     }
+
+    /// Register a new function (syscall) with the program runtime environment.
+    ///
+    /// **Important**: You should register all custom syscalls BEFORE adding any
+    /// programs to the store that will use them.
+    pub fn register_function(
+        &mut self,
+        name: &str,
+        value: BuiltinFunction<InvokeContext<'static>>,
+    ) -> Result<(), ElfError> {
+        self.program_runtime_environment
+            .register_function(name, value)
+    }
 }
 
 pub struct Builtin {

harness/src/program.rs

@@ -0,0 +1,67 @@
+use {
+    mollusk_svm::{result::Check, Mollusk},
+    solana_instruction::Instruction,
+    solana_program_runtime::{
+        invoke_context::InvokeContext,
+        solana_sbpf::{declare_builtin_function, memory_region::MemoryMapping},
+    },
+    solana_pubkey::Pubkey,
+};
+
+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>> {
+        invoke_context.consume_checked(to_burn)?;
+        Ok(0)
+    }
+);
+
+fn instruction_burn_cus(program_id: &Pubkey, to_burn: u64) -> Instruction {
+    Instruction::new_with_bytes(*program_id, &to_burn.to_le_bytes(), vec![])
+}
+
+#[test]
+fn test_custom_syscall() {
+    std::env::set_var("SBF_OUT_DIR", "../target/deploy");
+
+    let program_id = Pubkey::new_unique();
+
+    let mollusk = {
+        let mut mollusk = Mollusk::default();
+        mollusk
+            .program_cache
+            .register_function("sol_burn_cus", SyscallBurnCus::vm)
+            .unwrap();
+        mollusk.add_program(
+            &program_id,
+            "custom_syscall_program",
+            &mollusk_svm::program::loader_keys::LOADER_V3,
+        );
+        mollusk
+    };
+
+    let base_cus = mollusk
+        .process_and_validate_instruction(
+            &instruction_burn_cus(&program_id, 0),
+            &[],
+            &[Check::success()],
+        )
+        .compute_units_consumed;
+
+    for to_burn in [100, 1_000, 10_000] {
+        mollusk.process_and_validate_instruction(
+            &instruction_burn_cus(&program_id, to_burn),
+            &[],
+            &[Check::success(), Check::compute_units(base_cus + to_burn)],
+        );
+    }
+}

harness/tests/custom_syscall.rs

@@ -0,0 +1,16 @@
+[package]
+name = "custom-syscall-program"
+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]
+workspace = true

test-programs/custom-syscall/Cargo.toml

@@ -0,0 +1,27 @@
+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(())
+}