From b3d4cdfaaac3becb7c7015421303f9cf4c9f3202 Mon Sep 17 00:00:00 2001 From: Lucas Steuernagel Date: Fri, 18 Oct 2024 15:11:42 -0300 Subject: [PATCH] Update BuiltinProgram --- src/disassembler.rs | 2 +- src/elf.rs | 7 ++-- src/interpreter.rs | 2 +- src/jit.rs | 2 +- src/program.rs | 75 +++++++++++++++++++++++++++++++++--------- src/static_analysis.rs | 2 +- 6 files changed, 68 insertions(+), 22 deletions(-) diff --git a/src/disassembler.rs b/src/disassembler.rs index b500b0a56..3248ffc27 100644 --- a/src/disassembler.rs +++ b/src/disassembler.rs @@ -271,7 +271,7 @@ pub fn disassemble_instruction( function_name } else { name = "syscall"; - 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()) + loader.get_sparse_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 index 287f53e74..99f4912bf 100644 --- a/src/elf.rs +++ b/src/elf.rs @@ -315,7 +315,7 @@ impl Executable { self.get_config(), self.get_sbpf_version(), self.get_function_registry(), - self.loader.get_function_registry(), + self.loader.get_sparse_function_registry(), )?; Ok(()) } @@ -1071,7 +1071,10 @@ impl Executable { .entry(symbol.st_name) .or_insert_with(|| ebpf::hash_symbol_name(name)); if config.reject_broken_elfs - && loader.get_function_registry().lookup_by_key(hash).is_none() + && loader + .get_sparse_function_registry() + .lookup_by_key(hash) + .is_none() { return Err(ElfError::UnresolvedSymbol( String::from_utf8_lossy(name).to_string(), diff --git a/src/interpreter.rs b/src/interpreter.rs index 34ea6dd83..26e8c8009 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -536,7 +536,7 @@ impl<'a, 'b, C: ContextObject> Interpreter<'a, 'b, C> { }; if external { - if let Some((_function_name, function)) = self.executable.get_loader().get_function_registry().lookup_by_key(insn.imm as u32) { + if let Some((_function_name, function)) = self.executable.get_loader().get_sparse_function_registry().lookup_by_key(insn.imm as u32) { resolved = true; self.vm.due_insn_count = self.vm.previous_instruction_meter - self.vm.due_insn_count; diff --git a/src/jit.rs b/src/jit.rs index 1f48034e4..c961828f6 100644 --- a/src/jit.rs +++ b/src/jit.rs @@ -717,7 +717,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> { }; if external { - if let Some((_function_name, function)) = self.executable.get_loader().get_function_registry().lookup_by_key(insn.imm as u32) { + if let Some((_function_name, function)) = self.executable.get_loader().get_sparse_function_registry().lookup_by_key(insn.imm as u32) { self.emit_validate_and_profile_instruction_count(false, Some(0)); self.emit_ins(X86Instruction::load_immediate(OperandSize::S64, REGISTER_SCRATCH, function as usize as i64)); self.emit_ins(X86Instruction::call_immediate(self.relative_to_anchor(ANCHOR_EXTERNAL_FUNCTION_CALL, 5))); diff --git a/src/program.rs b/src/program.rs index 3c844317d..988b7be3d 100644 --- a/src/program.rs +++ b/src/program.rs @@ -153,7 +153,11 @@ impl FunctionRegistry { } else { ebpf::hash_symbol_name(&usize::from(value).to_le_bytes()) }; - if loader.get_function_registry().lookup_by_key(hash).is_some() { + if loader + .get_sparse_function_registry() + .lookup_by_key(hash) + .is_some() + { return Err(ElfError::SymbolHashCollision(hash)); } hash @@ -228,13 +232,17 @@ pub type BuiltinFunction = fn(*mut EbpfVm, u64, u64, u64, u64, u64); pub struct BuiltinProgram { /// Holds the Config if this is a loader program config: Option>, - /// Function pointers by symbol - functions: FunctionRegistry>, + /// Function pointers by symbol with sparse indexing + sparse_registry: FunctionRegistry>, + /// Function pointers by symbol with dense indexing + dense_registry: FunctionRegistry>, } impl PartialEq for BuiltinProgram { fn eq(&self, other: &Self) -> bool { - self.config.eq(&other.config) && self.functions.eq(&other.functions) + self.config.eq(&other.config) + && self.sparse_registry.eq(&other.sparse_registry) + && self.dense_registry.eq(&other.dense_registry) } } @@ -243,7 +251,8 @@ impl BuiltinProgram { pub fn new_loader(config: Config, functions: FunctionRegistry>) -> Self { Self { config: Some(Box::new(config)), - functions, + sparse_registry: functions, + dense_registry: FunctionRegistry::default(), } } @@ -251,7 +260,8 @@ impl BuiltinProgram { pub fn new_builtin(functions: FunctionRegistry>) -> Self { Self { config: None, - functions, + sparse_registry: functions, + dense_registry: FunctionRegistry::default(), } } @@ -259,7 +269,18 @@ impl BuiltinProgram { pub fn new_mock() -> Self { Self { config: Some(Box::default()), - functions: FunctionRegistry::default(), + sparse_registry: FunctionRegistry::default(), + dense_registry: FunctionRegistry::default(), + } + } + + /// Create a new loader with both dense and sparse function registrations + /// Use `BuiltinProgram::register_function` to register. + pub fn new_loader_with_dense_registration(config: Config) -> Self { + Self { + config: Some(Box::new(config)), + sparse_registry: FunctionRegistry::default(), + dense_registry: FunctionRegistry::default(), } } @@ -269,8 +290,8 @@ impl BuiltinProgram { } /// Get the function registry - pub fn get_function_registry(&self) -> &FunctionRegistry> { - &self.functions + pub fn get_sparse_function_registry(&self) -> &FunctionRegistry> { + &self.sparse_registry } /// Calculate memory size @@ -281,18 +302,40 @@ impl BuiltinProgram { } else { 0 }) - .saturating_add(self.functions.mem_size()) + .saturating_add(self.sparse_registry.mem_size()) + .saturating_add(self.dense_registry.mem_size()) + } + + /// Register a function both in the sparse and dense registries + pub fn register_function( + &mut self, + name: &str, + value: BuiltinFunction, + dense_key: u32, + ) -> Result<(), ElfError> { + self.sparse_registry.register_function_hashed(name, value)?; + self.dense_registry + .register_function(dense_key, name, value) } } impl std::fmt::Debug for BuiltinProgram { fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> { - writeln!(f, "{:?}", unsafe { - // `derive(Debug)` does not know that `C: ContextObject` does not need to implement `Debug` - std::mem::transmute::<&FunctionRegistry>, &FunctionRegistry>( - &self.functions, - ) - })?; + unsafe { + writeln!( + f, + "sparse: {:?}\n dense: {:?}", + // `derive(Debug)` does not know that `C: ContextObject` does not need to implement `Debug` + std::mem::transmute::< + &FunctionRegistry>, + &FunctionRegistry, + >(&self.sparse_registry,), + std::mem::transmute::< + &FunctionRegistry>, + &FunctionRegistry, + >(&self.dense_registry,) + )?; + } Ok(()) } } diff --git a/src/static_analysis.rs b/src/static_analysis.rs index aa5d0ae37..ab73c6cc8 100644 --- a/src/static_analysis.rs +++ b/src/static_analysis.rs @@ -236,7 +236,7 @@ impl<'a> Analysis<'a> { if let Some((function_name, _function)) = self .executable .get_loader() - .get_function_registry() + .get_sparse_function_registry() .lookup_by_key(insn.imm as u32) { if function_name == b"abort" {