Simplify syscall register and bind (#308)

cli/src/main.rs

@@ -215,7 +215,7 @@ fn main() {
     }
 
     for (hash, name) in executable.get_syscall_symbols() {
-        vm.bind_syscall_context_object(Box::new(MockSyscall { name: name.clone() }), Some(*hash))
+        vm.bind_syscall_context_objects(Box::new(MockSyscall { name: name.clone() }), Some(*hash))
             .unwrap();
     }
     let result = if matches.value_of("use").unwrap() == "interpreter" {

src/elf.rs

@@ -1030,7 +1030,7 @@ mod test {
         ebpf,
         elf::scroll::Pwrite,
         fuzz::fuzz,
-        syscalls::{BpfSyscallString, BpfSyscallU64},
+        syscalls::{BpfSyscallContext, BpfSyscallString, BpfSyscallU64},
         user_error::UserError,
         vm::{SyscallObject, TestInstructionMeter},
     };
@@ -1041,10 +1041,18 @@ mod test {
     fn syscall_registry() -> SyscallRegistry {
         let mut syscall_registry = SyscallRegistry::default();
         syscall_registry
-            .register_syscall_by_name(b"log", BpfSyscallString::call)
+            .register_syscall_by_name(
+                b"log",
+                BpfSyscallString::init::<BpfSyscallContext, UserError>,
+                BpfSyscallString::call,
+            )
             .unwrap();
         syscall_registry
-            .register_syscall_by_name(b"log_64", BpfSyscallU64::call)
+            .register_syscall_by_name(
+                b"log_64",
+                BpfSyscallU64::init::<BpfSyscallContext, UserError>,
+                BpfSyscallU64::call,
+            )
             .unwrap();
         syscall_registry
     }
@@ -1800,6 +1808,6 @@ mod test {
             Executable::jit_compile(&mut executable).unwrap();
         }
 
-        assert_eq!(18616, executable.mem_size());
+        assert_eq!(18640, executable.mem_size());
     }
 }

src/jit.rs

@@ -1846,6 +1846,7 @@ mod tests {
         syscall_registry
             .register_syscall_by_hash(
                 0xFFFFFFFF,
+                syscalls::BpfGatherBytes::init::<syscalls::BpfSyscallContext, UserError>,
                 syscalls::BpfGatherBytes::call,
             )
             .unwrap();

src/syscalls.rs

@@ -29,6 +29,9 @@ use crate::{
 };
 use std::{slice::from_raw_parts, str::from_utf8, u64};
 
+/// Test syscall context
+pub type BpfSyscallContext = u64;
+
 /// Return type of syscalls
 pub type Result = std::result::Result<u64, EbpfError<UserError>>;
 
@@ -80,6 +83,12 @@ pub const BPF_TRACE_PRINTK_IDX: u32 = 6;
 /// This would equally print the three numbers in `/sys/kernel/debug/tracing` file each time the
 /// program is run.
 pub struct BpfTracePrintf {}
+impl BpfTracePrintf {
+    /// new
+    pub fn init<C, E>(_unused: C) -> Box<dyn SyscallObject<UserError>> {
+        Box::new(Self {})
+    }
+}
 impl SyscallObject<UserError> for BpfTracePrintf {
     fn call(
         &mut self,
@@ -128,6 +137,12 @@ impl SyscallObject<UserError> for BpfTracePrintf {
 /// assert_eq!(result.unwrap(), 0x1122334455);
 /// ```
 pub struct BpfGatherBytes {}
+impl BpfGatherBytes {
+    /// new
+    pub fn init<C, E>(_unused: C) -> Box<dyn SyscallObject<UserError>> {
+        Box::new(Self {})
+    }
+}
 impl SyscallObject<UserError> for BpfGatherBytes {
     fn call(
         &mut self,
@@ -173,6 +188,12 @@ impl SyscallObject<UserError> for BpfGatherBytes {
 /// assert_eq!(val, &[0x00, 0x00, 0x00, 0x00, 0x00, 0x11, 0x22, 0x33]);
 /// ```
 pub struct BpfMemFrob {}
+impl BpfMemFrob {
+    /// new
+    pub fn init<C, E>(_unused: C) -> Box<dyn SyscallObject<UserError>> {
+        Box::new(Self {})
+    }
+}
 impl SyscallObject<UserError> for BpfMemFrob {
     fn call(
         &mut self,
@@ -221,6 +242,12 @@ impl SyscallObject<UserError> for BpfMemFrob {
 /// assert!(result.unwrap() != 0);
 /// ```
 pub struct BpfStrCmp {}
+impl BpfStrCmp {
+    /// new
+    pub fn init<C, E>(_unused: C) -> Box<dyn SyscallObject<UserError>> {
+        Box::new(Self {})
+    }
+}
 impl SyscallObject<UserError> for BpfStrCmp {
     fn call(
         &mut self,
@@ -261,6 +288,12 @@ impl SyscallObject<UserError> for BpfStrCmp {
 
 /// Prints a NULL-terminated UTF-8 string.
 pub struct BpfSyscallString {}
+impl BpfSyscallString {
+    /// new
+    pub fn init<C, E>(_unused: C) -> Box<dyn SyscallObject<UserError>> {
+        Box::new(Self {})
+    }
+}
 impl SyscallObject<UserError> for BpfSyscallString {
     fn call(
         &mut self,
@@ -291,6 +324,12 @@ impl SyscallObject<UserError> for BpfSyscallString {
 
 /// Prints the five arguments formated as u64 in decimal.
 pub struct BpfSyscallU64 {}
+impl BpfSyscallU64 {
+    /// new
+    pub fn init<C, E>(_unused: C) -> Box<dyn SyscallObject<UserError>> {
+        Box::new(Self {})
+    }
+}
 impl SyscallObject<UserError> for BpfSyscallU64 {
     fn call(
         &mut self,
@@ -313,7 +352,13 @@ impl SyscallObject<UserError> for BpfSyscallU64 {
 /// Example of a syscall with internal state.
 pub struct SyscallWithContext {
     /// Mutable state
-    pub context: u64,
+    pub context: BpfSyscallContext,
+}
+impl SyscallWithContext {
+    /// new
+    pub fn init<C, E>(context: BpfSyscallContext) -> Box<dyn SyscallObject<UserError>> {
+        Box::new(Self { context })
+    }
 }
 impl SyscallObject<UserError> for SyscallWithContext {
     fn call(

src/vm.rs

@@ -61,6 +61,9 @@ macro_rules! question_mark {
     }};
 }
 
+/// Syscall initialization function
+pub type SyscallInit<'a, C, E> = fn(C) -> Box<(dyn SyscallObject<E> + 'a)>;
+
 /// Syscall function without context
 pub type SyscallFunction<E, O> =
     fn(O, u64, u64, u64, u64, u64, &MemoryMapping, &mut ProgramResult<E>);
@@ -84,6 +87,8 @@ pub trait SyscallObject<E: UserDefinedError> {
 /// Syscall function and binding slot for a context object
 #[derive(Debug, PartialEq)]
 pub struct Syscall {
+    /// Syscall init
+    pub init: u64,
     /// Call the syscall function
     pub function: u64,
     /// Slot of context object
@@ -123,18 +128,21 @@ pub struct SyscallRegistry {
 
 impl SyscallRegistry {
     /// Register a syscall function by its symbol hash
-    pub fn register_syscall_by_hash<E: UserDefinedError, O: SyscallObject<E>>(
+    pub fn register_syscall_by_hash<'a, C, E: UserDefinedError, O: SyscallObject<E>>(
         &mut self,
         hash: u32,
+        init: SyscallInit<'a, C, E>,
         function: SyscallFunction<E, &mut O>,
     ) -> Result<(), EbpfError<E>> {
+        let init = init as *const u8 as u64;
         let function = function as *const u8 as u64;
         let context_object_slot = self.entries.len();
         if self
             .entries
             .insert(
                 hash,
                 Syscall {
+                    init,
                     function,
                     context_object_slot,
                 },
@@ -152,12 +160,13 @@ impl SyscallRegistry {
     }
 
     /// Register a syscall function by its symbol name
-    pub fn register_syscall_by_name<E: UserDefinedError, O: SyscallObject<E>>(
+    pub fn register_syscall_by_name<'a, C, E: UserDefinedError, O: SyscallObject<E>>(
         &mut self,
         name: &[u8],
+        init: SyscallInit<'a, C, E>,
         function: SyscallFunction<E, &mut O>,
     ) -> Result<(), EbpfError<E>> {
-        self.register_syscall_by_hash(ebpf::hash_symbol_name(name), function)
+        self.register_syscall_by_hash::<C, E, O>(ebpf::hash_symbol_name(name), init, function)
     }
 
     /// Get a symbol's function pointer and context object slot
@@ -568,7 +577,7 @@ impl<'a, E: UserDefinedError, I: InstructionMeter> EbpfVm<'a, E, I> {
         &self.tracer
     }
 
-    /// Bind a context object instance to a previously registered syscall
+    /// Bind a context objects instance to a previously registered syscalls
     ///
     /// # Examples
     ///
@@ -594,45 +603,55 @@ impl<'a, E: UserDefinedError, I: InstructionMeter> EbpfVm<'a, E, I> {
     /// // On running the program this syscall will print the content of registers r3, r4 and r5 to
     /// // standard output.
     /// let mut syscall_registry = SyscallRegistry::default();
-    /// syscall_registry.register_syscall_by_hash(6, BpfTracePrintf::call).unwrap();
+    /// syscall_registry.register_syscall_by_hash(6, BpfTracePrintf::init::<u64, UserError>, BpfTracePrintf::call).unwrap();
     /// // Instantiate an Executable and VM
     /// let config = Config::default();
     /// let mut bpf_functions = std::collections::BTreeMap::new();
     /// register_bpf_function(&config, &mut bpf_functions, &syscall_registry, 0, "entrypoint").unwrap();
     /// let mut executable = Executable::<UserError, TestInstructionMeter>::from_text_bytes(prog, None, config, syscall_registry, bpf_functions).unwrap();
     /// let mut vm = EbpfVm::<UserError, TestInstructionMeter>::new(&executable, &mut [], Vec::new()).unwrap();
     /// // Bind a context object instance to the previously registered syscall
-    /// vm.bind_syscall_context_object(Box::new(BpfTracePrintf {}), None);
+    /// vm.bind_syscall_context_objects(0, None);
     /// ```
-    pub fn bind_syscall_context_object(
+    pub fn bind_syscall_context_objects<C: Clone>(
         &mut self,
-        syscall_context_object: Box<dyn SyscallObject<E> + 'a>,
+        syscall_context: C,
         hash: Option<u32>,
     ) -> Result<(), EbpfError<E>> {
-        let fat_ptr: DynTraitFatPointer = unsafe { std::mem::transmute(&*syscall_context_object) };
         let syscall_registry = self.executable.get_syscall_registry();
-        let slot = match hash {
-            Some(hash) => {
-                syscall_registry
-                    .lookup_syscall(hash)
-                    .ok_or(EbpfError::SyscallNotRegistered(hash as usize))?
-                    .context_object_slot
+
+        for (_hash, syscall) in syscall_registry.entries.iter() {
+            let syscall_object_init_fn: SyscallInit<C, E> =
+                unsafe { std::mem::transmute(syscall.init) };
+            let syscall_context_object: Box<dyn SyscallObject<E> + 'a> =
+                syscall_object_init_fn(syscall_context.clone());
+            let fat_ptr: DynTraitFatPointer =
+                unsafe { std::mem::transmute(&*syscall_context_object) };
+
+            let slot = match hash {
+                Some(hash) => {
+                    syscall_registry
+                        .lookup_syscall(hash)
+                        .ok_or(EbpfError::SyscallNotRegistered(hash as usize))?
+                        .context_object_slot
+                }
+                None => syscall_registry
+                    .lookup_context_object_slot(fat_ptr.vtable.methods[0] as u64)
+                    .ok_or(EbpfError::SyscallNotRegistered(
+                        fat_ptr.vtable.methods[0] as usize,
+                    ))?,
+            };
+            if !self.syscall_context_objects[SYSCALL_CONTEXT_OBJECTS_OFFSET + slot].is_null() {
+                return Err(EbpfError::SyscallAlreadyBound(slot));
+            } else {
+                self.syscall_context_objects[SYSCALL_CONTEXT_OBJECTS_OFFSET + slot] = fat_ptr.data;
+                // Keep the dyn trait objects so that they can be dropped properly later
+                self.syscall_context_object_pool
+                    .push(syscall_context_object);
             }
-            None => syscall_registry
-                .lookup_context_object_slot(fat_ptr.vtable.methods[0] as u64)
-                .ok_or(EbpfError::SyscallNotRegistered(
-                    fat_ptr.vtable.methods[0] as usize,
-                ))?,
-        };
-        if !self.syscall_context_objects[SYSCALL_CONTEXT_OBJECTS_OFFSET + slot].is_null() {
-            Err(EbpfError::SyscallAlreadyBound(slot))
-        } else {
-            self.syscall_context_objects[SYSCALL_CONTEXT_OBJECTS_OFFSET + slot] = fat_ptr.data;
-            // Keep the dyn trait objects so that they can be dropped properly later
-            self.syscall_context_object_pool
-                .push(syscall_context_object);
-            Ok(())
         }
+
+        Ok(())
     }
 
     /// Lookup a syscall context object by its function pointer. Used for testing and validation.

tests/misc.rs

@@ -24,7 +24,7 @@ extern crate test_utils;
 use solana_rbpf::{
     elf::Executable,
     fuzz::fuzz,
-    syscalls::{BpfSyscallString, BpfSyscallU64},
+    syscalls::{BpfSyscallContext, BpfSyscallString, BpfSyscallU64},
     user_error::UserError,
     verifier::check,
     vm::{Config, EbpfVm, SyscallObject, SyscallRegistry, TestInstructionMeter},
@@ -114,10 +114,18 @@ fn test_fuzz_execute() {
         |bytes: &mut [u8]| {
             let mut syscall_registry = SyscallRegistry::default();
             syscall_registry
-                .register_syscall_by_name::<UserError, _>(b"log", BpfSyscallString::call)
+                .register_syscall_by_name(
+                    b"log",
+                    BpfSyscallString::init::<BpfSyscallContext, UserError>,
+                    BpfSyscallString::call,
+                )
                 .unwrap();
             syscall_registry
-                .register_syscall_by_name::<UserError, _>(b"log_64", BpfSyscallU64::call)
+                .register_syscall_by_name(
+                    b"log_64",
+                    BpfSyscallU64::init::<BpfSyscallContext, UserError>,
+                    BpfSyscallU64::call,
+                )
                 .unwrap();
             if let Ok(executable) = Executable::<UserError, TestInstructionMeter>::from_elf(
                 bytes,
@@ -131,10 +139,8 @@ fn test_fuzz_execute() {
                     Vec::new(),
                 )
                 .unwrap();
-                vm.bind_syscall_context_object(Box::new(BpfSyscallString {}), None)
-                    .unwrap();
-                vm.bind_syscall_context_object(Box::new(BpfSyscallU64 {}), None)
-                    .unwrap();
+                vm.bind_syscall_context_objects(0, None).unwrap();
+                vm.bind_syscall_context_objects(0, None).unwrap();
                 let _ = vm.execute_program_interpreted(&mut TestInstructionMeter {
                     remaining: 1_000_000,
                 });