Use FunctionRegistry in BuiltinProgram too.

Lichtso · Lichtso · commit f6fa3a3f7b4d · 2023-07-11T17:00:39.000+01:00
diff --git a/benches/elf_loader.rs b/benches/elf_loader.rs
@@ -11,20 +11,23 @@ extern crate test;
 extern crate test_utils;
 
 use solana_rbpf::{
-    elf::Executable,
+    elf::{Executable, FunctionRegistry},
     syscalls::bpf_syscall_u64,
     verifier::TautologyVerifier,
-    vm::{BuiltinProgram, Config, TestContextObject},
+    vm::{BuiltinFunction, BuiltinProgram, Config, TestContextObject},
 };
 use std::{fs::File, io::Read, sync::Arc};
 use test::Bencher;
 
 fn loader() -> Arc<BuiltinProgram<TestContextObject>> {
-    let mut loader = BuiltinProgram::new_loader(Config::default());
-    loader
-        .register_function(b"log_64", bpf_syscall_u64)
+    let mut function_registry = FunctionRegistry::<BuiltinFunction<TestContextObject>>::default();
+    function_registry
+        .register_function_hashed(b"log_64", bpf_syscall_u64)
         .unwrap();
-    Arc::new(loader)
+    Arc::new(BuiltinProgram::new_loader(
+        Config::default(),
+        function_registry,
+    ))
 }
 
 #[bench]
diff --git a/benches/jit_compile.rs b/benches/jit_compile.rs
@@ -12,7 +12,7 @@ extern crate test;
 use solana_rbpf::{
     elf::Executable,
     verifier::{RequisiteVerifier, TautologyVerifier},
-    vm::{BuiltinProgram, Config, TestContextObject},
+    vm::{BuiltinProgram, TestContextObject},
 };
 use std::{fs::File, io::Read, sync::Arc};
 use test::Bencher;
@@ -25,7 +25,7 @@ fn bench_init_vm(bencher: &mut Bencher) {
     file.read_to_end(&mut elf).unwrap();
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_elf(
         &elf,
-        Arc::new(BuiltinProgram::new_loader(Config::default())),
+        Arc::new(BuiltinProgram::new_mock()),
     )
     .unwrap();
     let verified_executable =
@@ -52,7 +52,7 @@ fn bench_jit_compile(bencher: &mut Bencher) {
     file.read_to_end(&mut elf).unwrap();
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_elf(
         &elf,
-        Arc::new(BuiltinProgram::new_loader(Config::default())),
+        Arc::new(BuiltinProgram::new_mock()),
     )
     .unwrap();
     let mut verified_executable =
diff --git a/benches/vm_execution.rs b/benches/vm_execution.rs
@@ -11,7 +11,7 @@ extern crate test;
 
 use solana_rbpf::{
     ebpf,
-    elf::Executable,
+    elf::{Executable, FunctionRegistry},
     memory_region::MemoryRegion,
     verifier::{RequisiteVerifier, TautologyVerifier},
     vm::{BuiltinProgram, Config, TestContextObject},
@@ -27,7 +27,7 @@ fn bench_init_interpreter_execution(bencher: &mut Bencher) {
     file.read_to_end(&mut elf).unwrap();
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_elf(
         &elf,
-        Arc::new(BuiltinProgram::new_loader(Config::default())),
+        Arc::new(BuiltinProgram::new_mock()),
     )
     .unwrap();
     let verified_executable =
@@ -56,7 +56,7 @@ fn bench_init_jit_execution(bencher: &mut Bencher) {
     file.read_to_end(&mut elf).unwrap();
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_elf(
         &elf,
-        Arc::new(BuiltinProgram::new_loader(Config::default())),
+        Arc::new(BuiltinProgram::new_mock()),
     )
     .unwrap();
     let mut verified_executable =
@@ -88,7 +88,10 @@ fn bench_jit_vs_interpreter(
 ) {
     let executable = solana_rbpf::assembler::assemble::<TestContextObject>(
         assembly,
-        Arc::new(BuiltinProgram::new_loader(config)),
+        Arc::new(BuiltinProgram::new_loader(
+            config,
+            FunctionRegistry::default(),
+        )),
     )
     .unwrap();
     let mut verified_executable =
diff --git a/cli/src/main.rs b/cli/src/main.rs
@@ -3,7 +3,7 @@ use solana_rbpf::{
     aligned_memory::AlignedMemory,
     assembler::assemble,
     ebpf,
-    elf::Executable,
+    elf::{Executable, FunctionRegistry},
     memory_region::{MemoryMapping, MemoryRegion},
     static_analysis::Analysis,
     verifier::{RequisiteVerifier, TautologyVerifier},
@@ -92,11 +92,15 @@ fn main() {
         )
         .get_matches();
 
-    let loader = Arc::new(BuiltinProgram::new_loader(Config {
-        enable_instruction_tracing: matches.is_present("trace") || matches.is_present("profile"),
-        enable_symbol_and_section_labels: true,
-        ..Config::default()
-    }));
+    let loader = Arc::new(BuiltinProgram::new_loader(
+        Config {
+            enable_instruction_tracing: matches.is_present("trace")
+                || matches.is_present("profile"),
+            enable_symbol_and_section_labels: true,
+            ..Config::default()
+        },
+        FunctionRegistry::default(),
+    ));
     let executable = match matches.value_of("assembler") {
         Some(asm_file_name) => {
             let mut file = File::open(Path::new(asm_file_name)).unwrap();
diff --git a/examples/disassemble.rs b/examples/disassemble.rs
@@ -31,7 +31,7 @@ fn main() {
         0x00, 0x00, 0x00, 0x00, 0xb7, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x95, 0x00, 0x00,
         0x00, 0x00, 0x00, 0x00, 0x00,
     ];
-    let loader = Arc::new(BuiltinProgram::default());
+    let loader = Arc::new(BuiltinProgram::new_mock());
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(
         program,
         loader,
diff --git a/examples/to_json.rs b/examples/to_json.rs
@@ -30,7 +30,7 @@ use std::sync::Arc;
 fn to_json(program: &[u8]) -> String {
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(
         program,
-        Arc::new(BuiltinProgram::default()),
+        Arc::new(BuiltinProgram::new_mock()),
         SBPFVersion::V2,
         FunctionRegistry::default(),
     )
diff --git a/fuzz/fuzz_targets/dumb.rs b/fuzz/fuzz_targets/dumb.rs
@@ -35,7 +35,7 @@ fuzz_target!(|data: DumbFuzzData| {
     let mut mem = data.mem;
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(
         &prog,
-        std::sync::Arc::new(BuiltinProgram::new_loader(config)),
+        std::sync::Arc::new(BuiltinProgram::new_loader(config, FunctionRegistry::default())),
         SBPFVersion::V2,
         function_registry,
     )
diff --git a/fuzz/fuzz_targets/smart.rs b/fuzz/fuzz_targets/smart.rs
@@ -39,7 +39,7 @@ fuzz_target!(|data: FuzzData| {
     let mut mem = data.mem;
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(
         prog.into_bytes(),
-        std::sync::Arc::new(BuiltinProgram::new_loader(config)),
+        std::sync::Arc::new(BuiltinProgram::new_loader(config, FunctionRegistry::default())),
         SBPFVersion::V2,
         function_registry,
     )
diff --git a/fuzz/fuzz_targets/smart_jit_diff.rs b/fuzz/fuzz_targets/smart_jit_diff.rs
@@ -47,7 +47,7 @@ fuzz_target!(|data: FuzzData| {
     let mut jit_mem = data.mem;
     let mut executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(
         prog.into_bytes(),
-        std::sync::Arc::new(BuiltinProgram::new_loader(config)),
+        std::sync::Arc::new(BuiltinProgram::new_loader(config, FunctionRegistry::default())),
         SBPFVersion::V2,
         function_registry,
     )
diff --git a/fuzz/fuzz_targets/smarter_jit_diff.rs b/fuzz/fuzz_targets/smarter_jit_diff.rs
@@ -46,7 +46,7 @@ fuzz_target!(|data: FuzzData| {
     let mut jit_mem = data.mem;
     let mut executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(
         prog.into_bytes(),
-        std::sync::Arc::new(BuiltinProgram::new_loader(config)),
+        std::sync::Arc::new(BuiltinProgram::new_loader(config, FunctionRegistry::default())),
         SBPFVersion::V2,
         function_registry,
     )
diff --git a/src/assembler.rs b/src/assembler.rs
@@ -191,7 +191,7 @@ fn insn(opc: u8, dst: i64, src: i64, off: i64, imm: i64) -> Result<Insn, String>
 ///     be16 r0
 ///     neg64 r2
 ///     exit",
-///     std::sync::Arc::new(BuiltinProgram::new_loader(Config::default())),
+///     std::sync::Arc::new(BuiltinProgram::new_mock()),
 /// ).unwrap();
 /// let program = executable.get_text_bytes().1;
 /// println!("{:?}", program);
diff --git a/src/disassembler.rs b/src/disassembler.rs
@@ -264,7 +264,7 @@ pub fn disassemble_instruction<C: ContextObject>(
                 function_name
             } else {
                 name = "syscall";
-                loader.lookup_function(insn.imm as u32).map(|(function_name, _)| String::from_utf8_lossy(function_name).to_string()).unwrap_or_else(|| "[invalid]".to_string())
+                loader.get_function_registry().lookup_by_key(insn.imm as u32).map(|(function_name, _)| String::from_utf8_lossy(function_name).to_string()).unwrap_or_else(|| "[invalid]".to_string())
             };
             desc = format!("{name} {function_name}");
         },
diff --git a/src/elf.rs b/src/elf.rs
@@ -279,11 +279,7 @@ impl<T: Copy + PartialEq> FunctionRegistry<T> {
     }
 
     /// Register a symbol with an implicit key
-    pub fn register_function_hashed<C: ContextObject>(
-        &mut self,
-        name: &[u8],
-        value: T,
-    ) -> Result<u32, ElfError> {
+    pub fn register_function_hashed(&mut self, name: &[u8], value: T) -> Result<u32, ElfError> {
         let key = ebpf::hash_symbol_name(name);
         self.register_function(key, name.to_vec(), value)?;
         Ok(key)
@@ -310,7 +306,7 @@ impl<T: Copy + PartialEq> FunctionRegistry<T> {
                 ebpf::hash_symbol_name(&key)
             };
             if config.external_internal_function_hash_collision
-                && loader.lookup_function(hash).is_some()
+                && loader.get_function_registry().lookup_by_key(hash).is_some()
             {
                 return Err(ElfError::SymbolHashCollision(hash));
             }
@@ -1278,7 +1274,9 @@ impl<V: Verifier, C: ContextObject> Executable<V, C> {
                         let hash = *syscall_cache
                             .entry(symbol.st_name())
                             .or_insert_with(|| ebpf::hash_symbol_name(name));
-                        if config.reject_broken_elfs && loader.lookup_function(hash).is_none() {
+                        if config.reject_broken_elfs
+                            && loader.get_function_registry().lookup_by_key(hash).is_none()
+                        {
                             return Err(ElfError::UnresolvedSymbol(
                                 String::from_utf8_lossy(name).to_string(),
                                 r_offset
@@ -1371,22 +1369,26 @@ mod test {
         },
         fuzz::fuzz,
         syscalls,
-        vm::{ProgramResult, TestContextObject},
+        vm::{BuiltinFunction, ProgramResult, TestContextObject},
     };
     use rand::{distributions::Uniform, Rng};
     use std::{fs::File, io::Read};
     use test_utils::assert_error;
     type ElfExecutable = Executable<TautologyVerifier, TestContextObject>;
 
     fn loader() -> Arc<BuiltinProgram<TestContextObject>> {
-        let mut loader = BuiltinProgram::new_loader(Config::default());
-        loader
-            .register_function(b"log", syscalls::bpf_syscall_string)
+        let mut function_registry =
+            FunctionRegistry::<BuiltinFunction<TestContextObject>>::default();
+        function_registry
+            .register_function_hashed(b"log", syscalls::bpf_syscall_string)
             .unwrap();
-        loader
-            .register_function(b"log_64", syscalls::bpf_syscall_u64)
+        function_registry
+            .register_function_hashed(b"log_64", syscalls::bpf_syscall_u64)
             .unwrap();
-        Arc::new(loader)
+        Arc::new(BuiltinProgram::new_loader(
+            Config::default(),
+            function_registry,
+        ))
     }
 
     #[test]
diff --git a/src/interpreter.rs b/src/interpreter.rs
@@ -443,7 +443,7 @@ impl<'a, 'b, V: Verifier, C: ContextObject> Interpreter<'a, 'b, V, C> {
                 };
 
                 if external {
-                    if let Some((_function_name, function)) = self.executable.get_loader().lookup_function(insn.imm as u32) {
+                    if let Some((_function_name, function)) = self.executable.get_loader().get_function_registry().lookup_by_key(insn.imm as u32) {
                         resolved = true;
 
                         if config.enable_instruction_meter {
diff --git a/src/jit.rs b/src/jit.rs
@@ -599,7 +599,7 @@ impl<'a, V: Verifier, C: ContextObject> JitCompiler<'a, V, C> {
                     };
 
                     if external {
-                        if let Some((_function_name, function)) = self.executable.get_loader().lookup_function(insn.imm as u32) {
+                        if let Some((_function_name, function)) = self.executable.get_loader().get_function_registry().lookup_by_key(insn.imm as u32) {
                             self.emit_validate_and_profile_instruction_count(true, Some(0));
                             self.emit_ins(X86Instruction::load_immediate(OperandSize::S64, R11, function as usize as i64));
                             self.emit_ins(X86Instruction::call_immediate(self.relative_to_anchor(ANCHOR_EXTERNAL_FUNCTION_CALL, 5)));
@@ -1555,7 +1555,7 @@ mod tests {
         elf::{FunctionRegistry, SBPFVersion},
         syscalls,
         verifier::TautologyVerifier,
-        vm::{BuiltinProgram, TestContextObject},
+        vm::{BuiltinFunction, BuiltinProgram, TestContextObject},
     };
     use byteorder::{ByteOrder, LittleEndian};
     use std::sync::Arc;
@@ -1598,13 +1598,18 @@ mod tests {
     fn create_mockup_executable(
         program: &[u8],
     ) -> Executable<TautologyVerifier, TestContextObject> {
-        let mut loader = BuiltinProgram::new_loader(Config {
-            noop_instruction_rate: 0,
-            ..Config::default()
-        });
-        loader
-            .register_function(b"gather_bytes", syscalls::bpf_gather_bytes)
+        let mut function_registry =
+            FunctionRegistry::<BuiltinFunction<TestContextObject>>::default();
+        function_registry
+            .register_function_hashed(b"gather_bytes", syscalls::bpf_gather_bytes)
             .unwrap();
+        let loader = BuiltinProgram::new_loader(
+            Config {
+                noop_instruction_rate: 0,
+                ..Config::default()
+            },
+            function_registry,
+        );
         let mut function_registry = FunctionRegistry::default();
         function_registry
             .register_function(8, b"function_foo".to_vec(), 8)
diff --git a/src/static_analysis.rs b/src/static_analysis.rs
@@ -231,7 +231,8 @@ impl<'a> Analysis<'a> {
                     if let Some((function_name, _function)) = self
                         .executable
                         .get_loader()
-                        .lookup_function(insn.imm as u32)
+                        .get_function_registry()
+                        .lookup_by_key(insn.imm as u32)
                     {
                         if function_name == b"abort" {
                             self.cfg_nodes
diff --git a/src/vm.rs b/src/vm.rs
diff --git a/tests/assembler.rs b/tests/assembler.rs
diff --git a/tests/disassembler.rs b/tests/disassembler.rs
diff --git a/tests/execution.rs b/tests/execution.rs
diff --git a/tests/verifier.rs b/tests/verifier.rs

Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ use std::sync::Arc;`
`30`	`30`	`fn to_json(program: &[u8]) -> String {`
`31`	`31`	`let executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(`
`32`	`32`	`program,`
`33`		`- Arc::new(BuiltinProgram::default()),`
	`33`	`+ Arc::new(BuiltinProgram::new_mock()),`
`34`	`34`	`SBPFVersion::V2,`
`35`	`35`	`FunctionRegistry::default(),`
`36`	`36`	`)`
Original file line number	Diff line number	Diff line change
`@@ -35,7 +35,7 @@ fuzz_target!(\|data: DumbFuzzData\| {`
`35`	`35`	`let mut mem = data.mem;`
`36`	`36`	`let executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(`
`37`	`37`	`&prog,`
`38`		`- std::sync::Arc::new(BuiltinProgram::new_loader(config)),`
	`38`	`+ std::sync::Arc::new(BuiltinProgram::new_loader(config, FunctionRegistry::default())),`
`39`	`39`	`SBPFVersion::V2,`
`40`	`40`	`function_registry,`
`41`	`41`	`)`
Original file line number	Diff line number	Diff line change
`@@ -39,7 +39,7 @@ fuzz_target!(\|data: FuzzData\| {`
`39`	`39`	`let mut mem = data.mem;`
`40`	`40`	`let executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(`
`41`	`41`	`prog.into_bytes(),`
`42`		`- std::sync::Arc::new(BuiltinProgram::new_loader(config)),`
	`42`	`+ std::sync::Arc::new(BuiltinProgram::new_loader(config, FunctionRegistry::default())),`
`43`	`43`	`SBPFVersion::V2,`
`44`	`44`	`function_registry,`
`45`	`45`	`)`
Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ fuzz_target!(\|data: FuzzData\| {`
`47`	`47`	`let mut jit_mem = data.mem;`
`48`	`48`	`let mut executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(`
`49`	`49`	`prog.into_bytes(),`
`50`		`- std::sync::Arc::new(BuiltinProgram::new_loader(config)),`
	`50`	`+ std::sync::Arc::new(BuiltinProgram::new_loader(config, FunctionRegistry::default())),`
`51`	`51`	`SBPFVersion::V2,`
`52`	`52`	`function_registry,`
`53`	`53`	`)`
Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,7 @@ fuzz_target!(\|data: FuzzData\| {`
`46`	`46`	`let mut jit_mem = data.mem;`
`47`	`47`	`let mut executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(`
`48`	`48`	`prog.into_bytes(),`
`49`		`- std::sync::Arc::new(BuiltinProgram::new_loader(config)),`
	`49`	`+ std::sync::Arc::new(BuiltinProgram::new_loader(config, FunctionRegistry::default())),`
`50`	`50`	`SBPFVersion::V2,`
`51`	`51`	`function_registry,`
`52`	`52`	`)`