Refactor - vm interface of BuiltinFunctions (#535)

* Replaces all context parameters of BuiltinFunction with a reference to the vm.

* Deduplicates instruction metering of interpreter and JIT by moving it into declare_builtin_function.

* Moves get_runtime_environment_key() from JIT into VM.

* Passes the encrypted VM pointer to BuiltinFunctions.

* Puts due_insn_count into the VM.

* Moves the calculation of due_insn_count from the VM into the JIT epilogue.

* Makes $ContextObject generic in declare_builtin_function!().

Author: Lichtso

fuzz/fuzz_targets/common.rs

@@ -12,7 +12,6 @@ pub struct ConfigTemplate {
     enable_stack_frame_gaps: bool,
     enable_symbol_and_section_labels: bool,
     sanitize_user_provided_values: bool,
-    encrypt_runtime_environment: bool,
     reject_callx_r10: bool,
     optimize_rodata: bool,
 }
@@ -27,7 +26,6 @@ impl<'a> Arbitrary<'a> for ConfigTemplate {
             enable_stack_frame_gaps: bools & (1 << 0) != 0,
             enable_symbol_and_section_labels: bools & (1 << 1) != 0,
             sanitize_user_provided_values: bools & (1 << 3) != 0,
-            encrypt_runtime_environment: bools & (1 << 4) != 0,
             reject_callx_r10: bools & (1 << 6) != 0,
             optimize_rodata: bools & (1 << 9) != 0,
         })
@@ -51,7 +49,6 @@ impl From<ConfigTemplate> for Config {
                 enable_stack_frame_gaps,
                 enable_symbol_and_section_labels,
                 sanitize_user_provided_values,
-                encrypt_runtime_environment,
                 reject_callx_r10,
                 optimize_rodata,
             } => Config {
@@ -61,7 +58,6 @@ impl From<ConfigTemplate> for Config {
                 enable_symbol_and_section_labels,
                 noop_instruction_rate,
                 sanitize_user_provided_values,
-                encrypt_runtime_environment,
                 reject_callx_r10,
                 optimize_rodata,
                 ..Default::default()

src/interpreter.rs

@@ -16,7 +16,7 @@ use crate::{
     ebpf::{self, STACK_PTR_REG},
     elf::Executable,
     error::EbpfError,
-    vm::{Config, ContextObject, EbpfVm, ProgramResult},
+    vm::{get_runtime_environment_key, Config, ContextObject, EbpfVm, ProgramResult},
 };
 
 /// Virtual memory operation helper.
@@ -91,7 +91,6 @@ pub struct Interpreter<'a, 'b, C: ContextObject> {
     pub(crate) executable: &'a Executable<C>,
     pub(crate) program: &'a [u8],
     pub(crate) program_vm_addr: u64,
-    pub(crate) due_insn_count: u64,
 
     /// General purpose registers and pc
     pub reg: [u64; 12],
@@ -115,7 +114,6 @@ impl<'a, 'b, C: ContextObject> Interpreter<'a, 'b, C> {
             executable,
             program,
             program_vm_addr,
-            due_insn_count: 0,
             reg: registers,
             #[cfg(feature = "debugger")]
             debug_state: DebugState::Continue,
@@ -161,7 +159,7 @@ impl<'a, 'b, C: ContextObject> Interpreter<'a, 'b, C> {
     pub fn step(&mut self) -> bool {
         let config = &self.executable.get_config();
 
-        self.due_insn_count += 1;
+        self.vm.due_insn_count += 1;
         let mut next_pc = self.reg[11] + 1;
         if next_pc as usize * ebpf::INSN_SIZE > self.program.len() {
             throw_error!(self, EbpfError::ExecutionOverrun);
@@ -456,7 +454,7 @@ impl<'a, 'b, C: ContextObject> Interpreter<'a, 'b, C> {
                 }
                 check_pc!(self, next_pc, (target_pc - self.program_vm_addr) / ebpf::INSN_SIZE as u64);
                 if self.executable.get_sbpf_version().static_syscalls() && self.executable.get_function_registry().lookup_by_key(next_pc as u32).is_none() {
-                    self.due_insn_count += 1;
+                    self.vm.due_insn_count += 1;
                     self.reg[11] = next_pc;
                     throw_error!(self, EbpfError::UnsupportedInstruction);
                 }
@@ -476,27 +474,20 @@ impl<'a, 'b, C: ContextObject> Interpreter<'a, 'b, C> {
                     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 {
-                            self.vm.context_object_pointer.consume(self.due_insn_count);
-                        }
-                        self.due_insn_count = 0;
+                        self.vm.due_insn_count = self.vm.previous_instruction_meter - self.vm.due_insn_count;
                         function(
-                            self.vm.context_object_pointer,
+                            unsafe { (self.vm as *mut _ as *mut u64).offset(get_runtime_environment_key() as isize) as *mut _ },
                             self.reg[1],
                             self.reg[2],
                             self.reg[3],
                             self.reg[4],
                             self.reg[5],
-                            &mut self.vm.memory_mapping,
-                            &mut self.vm.program_result,
                         );
+                        self.vm.due_insn_count = 0;
                         self.reg[0] = match &self.vm.program_result {
                             ProgramResult::Ok(value) => *value,
                             ProgramResult::Err(_err) => return false,
                         };
-                        if config.enable_instruction_meter {
-                            self.vm.previous_instruction_meter = self.vm.context_object_pointer.get_remaining();
-                        }
                     }
                 }
 
@@ -519,7 +510,7 @@ impl<'a, 'b, C: ContextObject> Interpreter<'a, 'b, C> {
 
             ebpf::EXIT       => {
                 if self.vm.call_depth == 0 {
-                    if config.enable_instruction_meter && self.due_insn_count > self.vm.previous_instruction_meter {
+                    if config.enable_instruction_meter && self.vm.due_insn_count > self.vm.previous_instruction_meter {
                         throw_error!(self, EbpfError::ExceededMaxInstructions);
                     }
                     self.vm.program_result = ProgramResult::Ok(self.reg[0]);
@@ -542,7 +533,7 @@ impl<'a, 'b, C: ContextObject> Interpreter<'a, 'b, C> {
             _ => throw_error!(self, EbpfError::UnsupportedInstruction),
         }
 
-        if config.enable_instruction_meter && self.due_insn_count >= self.vm.previous_instruction_meter {
+        if config.enable_instruction_meter && self.vm.due_insn_count >= self.vm.previous_instruction_meter {
             self.reg[11] += 1;
             throw_error!(self, EbpfError::ExceededMaxInstructions);
         }

src/jit.rs

@@ -21,20 +21,14 @@ use crate::{
         allocate_pages, free_pages, get_system_page_size, protect_pages, round_to_page_size,
     },
     memory_region::{AccessType, MemoryMapping},
-    vm::{Config, ContextObject, EbpfVm, ProgramResult},
+    vm::{get_runtime_environment_key, Config, ContextObject, EbpfVm, ProgramResult},
     x86::*,
 };
 
-/// Shift the RUNTIME_ENVIRONMENT_KEY by this many bits to the LSB
-///
-/// 3 bits for 8 Byte alignment, and 1 bit to have encoding space for the RuntimeEnvironment.
-const PROGRAM_ENVIRONMENT_KEY_SHIFT: u32 = 4;
 const MAX_EMPTY_PROGRAM_MACHINE_CODE_LENGTH: usize = 4096;
 const MAX_MACHINE_CODE_LENGTH_PER_INSTRUCTION: usize = 110;
 const MACHINE_CODE_PER_INSTRUCTION_METER_CHECKPOINT: usize = 13;
 
-static RUNTIME_ENVIRONMENT_KEY: std::sync::OnceLock<i32> = std::sync::OnceLock::<i32>::new();
-
 pub struct JitProgram {
     /// OS page size in bytes and the alignment of the sections
     page_size: usize,
@@ -100,10 +94,8 @@ impl JitProgram {
         _config: &Config,
         vm: &mut EbpfVm<C>,
         registers: [u64; 12],
-    ) -> i64 {
+    ) {
         unsafe {
-            let mut instruction_meter =
-                (vm.previous_instruction_meter as i64).wrapping_add(registers[11] as i64);
             std::arch::asm!(
                 // RBP and RBX must be saved and restored manually in the current version of rustc and llvm.
                 "push rbx",
@@ -124,19 +116,17 @@ impl JitProgram {
                 "mov rbp, [r11 + 0x50]",
                 "mov r11, [r11 + 0x58]",
                 "call r10",
-                "mov rax, rbx",
                 "pop rbp",
                 "pop rbx",
                 host_stack_pointer = in(reg) &mut vm.host_stack_pointer,
-                inlateout("rax") instruction_meter,
-                inlateout("rdi") (vm as *mut _ as *mut u64).offset(*RUNTIME_ENVIRONMENT_KEY.get().unwrap() as isize) => _,
+                inlateout("rdi") (vm as *mut _ as *mut u64).offset(get_runtime_environment_key() as isize) => _,
+                inlateout("rax") (vm.previous_instruction_meter as i64).wrapping_add(registers[11] as i64) => _,
                 inlateout("r10") self.pc_section[registers[11] as usize] => _,
                 inlateout("r11") &registers => _,
                 lateout("rsi") _, lateout("rdx") _, lateout("rcx") _, lateout("r8") _,
                 lateout("r9") _, lateout("r12") _, lateout("r13") _, lateout("r14") _, lateout("r15") _,
                 // lateout("rbp") _, lateout("rbx") _,
             );
-            instruction_meter
         }
     }
 
@@ -255,11 +245,12 @@ enum RuntimeEnvironmentSlot {
     StackPointer = 2,
     ContextObjectPointer = 3,
     PreviousInstructionMeter = 4,
-    StopwatchNumerator = 5,
-    StopwatchDenominator = 6,
-    Registers = 7,
-    ProgramResult = 19,
-    MemoryMapping = 27,
+    DueInsnCount = 5,
+    StopwatchNumerator = 6,
+    StopwatchDenominator = 7,
+    Registers = 8,
+    ProgramResult = 20,
+    MemoryMapping = 28,
 }
 
 /* Explaination of the Instruction Meter
@@ -357,13 +348,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
             code_length_estimate += pc / config.instruction_meter_checkpoint_distance * MACHINE_CODE_PER_INSTRUCTION_METER_CHECKPOINT;
         }
 
-        let runtime_environment_key = *RUNTIME_ENVIRONMENT_KEY.get_or_init(|| {
-            if config.encrypt_runtime_environment {
-                rand::thread_rng().gen::<i32>() >> PROGRAM_ENVIRONMENT_KEY_SHIFT
-            } else {
-                0
-            }
-        });
+        let runtime_environment_key = get_runtime_environment_key();
         let mut diversification_rng = SmallRng::from_rng(rand::thread_rng()).map_err(|_| EbpfError::JitNotCompiled)?;
 
         Ok(Self {
@@ -1326,6 +1311,10 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         if self.config.enable_instruction_meter {
             self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x81, 5, REGISTER_INSTRUCTION_METER, 1, None)); // REGISTER_INSTRUCTION_METER -= 1;
             self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x29, REGISTER_SCRATCH, REGISTER_INSTRUCTION_METER, 0, None)); // REGISTER_INSTRUCTION_METER -= pc;
+            // *DueInsnCount = *PreviousInstructionMeter - REGISTER_INSTRUCTION_METER;
+            self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x2B, REGISTER_INSTRUCTION_METER, REGISTER_PTR_TO_VM, 0, Some(X86IndirectAccess::Offset(self.slot_in_vm(RuntimeEnvironmentSlot::PreviousInstructionMeter))))); // REGISTER_INSTRUCTION_METER -= *PreviousInstructionMeter;
+            self.emit_ins(X86Instruction::alu(OperandSize::S64, 0xf7, 3, REGISTER_INSTRUCTION_METER, 0, None)); // REGISTER_INSTRUCTION_METER = -REGISTER_INSTRUCTION_METER;
+            self.emit_ins(X86Instruction::store(OperandSize::S64, REGISTER_INSTRUCTION_METER, REGISTER_PTR_TO_VM, X86IndirectAccess::Offset(self.slot_in_vm(RuntimeEnvironmentSlot::DueInsnCount)))); // *DueInsnCount = REGISTER_INSTRUCTION_METER;
         }
         // Print stop watch value
         fn stopwatch_result(numerator: u64, denominator: u64) {
@@ -1398,29 +1387,18 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         self.set_anchor(ANCHOR_EXTERNAL_FUNCTION_CALL);
         self.emit_ins(X86Instruction::push_immediate(OperandSize::S64, -1)); // Used as PC value in error case, acts as stack padding otherwise
         if self.config.enable_instruction_meter {
-            // REGISTER_INSTRUCTION_METER = *PreviousInstructionMeter - REGISTER_INSTRUCTION_METER;
-            self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x2B, REGISTER_INSTRUCTION_METER, REGISTER_PTR_TO_VM, 0, Some(X86IndirectAccess::Offset(self.slot_in_vm(RuntimeEnvironmentSlot::PreviousInstructionMeter))))); // REGISTER_INSTRUCTION_METER -= *PreviousInstructionMeter;
-            self.emit_ins(X86Instruction::alu(OperandSize::S64, 0xf7, 3, REGISTER_INSTRUCTION_METER, 0, None)); // REGISTER_INSTRUCTION_METER = -REGISTER_INSTRUCTION_METER;
-            self.emit_rust_call(Value::Constant64(C::consume as *const u8 as i64, false), &[
-                Argument { index: 1, value: Value::Register(REGISTER_INSTRUCTION_METER) },
-                Argument { index: 0, value: Value::RegisterIndirect(REGISTER_PTR_TO_VM, self.slot_in_vm(RuntimeEnvironmentSlot::ContextObjectPointer), false) },
-            ], None);
+            self.emit_ins(X86Instruction::store(OperandSize::S64, REGISTER_INSTRUCTION_METER, REGISTER_PTR_TO_VM, X86IndirectAccess::Offset(self.slot_in_vm(RuntimeEnvironmentSlot::DueInsnCount)))); // *DueInsnCount = REGISTER_INSTRUCTION_METER;
         }
         self.emit_rust_call(Value::Register(REGISTER_SCRATCH), &[
-            Argument { index: 7, value: Value::RegisterPlusConstant32(REGISTER_PTR_TO_VM, self.slot_in_vm(RuntimeEnvironmentSlot::ProgramResult), false) },
-            Argument { index: 6, value: Value::RegisterPlusConstant32(REGISTER_PTR_TO_VM, self.slot_in_vm(RuntimeEnvironmentSlot::MemoryMapping), false) },
             Argument { index: 5, value: Value::Register(ARGUMENT_REGISTERS[5]) },
             Argument { index: 4, value: Value::Register(ARGUMENT_REGISTERS[4]) },
             Argument { index: 3, value: Value::Register(ARGUMENT_REGISTERS[3]) },
             Argument { index: 2, value: Value::Register(ARGUMENT_REGISTERS[2]) },
             Argument { index: 1, value: Value::Register(ARGUMENT_REGISTERS[1]) },
-            Argument { index: 0, value: Value::RegisterIndirect(REGISTER_PTR_TO_VM, self.slot_in_vm(RuntimeEnvironmentSlot::ContextObjectPointer), false) },
+            Argument { index: 0, value: Value::Register(REGISTER_PTR_TO_VM) },
         ], None);
         if self.config.enable_instruction_meter {
-            self.emit_rust_call(Value::Constant64(C::get_remaining as *const u8 as i64, false), &[
-                Argument { index: 0, value: Value::RegisterIndirect(REGISTER_PTR_TO_VM, self.slot_in_vm(RuntimeEnvironmentSlot::ContextObjectPointer), false) },
-            ], Some(REGISTER_INSTRUCTION_METER));
-            self.emit_ins(X86Instruction::store(OperandSize::S64, REGISTER_INSTRUCTION_METER, REGISTER_PTR_TO_VM, X86IndirectAccess::Offset(self.slot_in_vm(RuntimeEnvironmentSlot::PreviousInstructionMeter)))); // *PreviousInstructionMeter = REGISTER_INSTRUCTION_METER;
+            self.emit_ins(X86Instruction::load(OperandSize::S64, REGISTER_PTR_TO_VM, REGISTER_INSTRUCTION_METER, X86IndirectAccess::Offset(self.slot_in_vm(RuntimeEnvironmentSlot::PreviousInstructionMeter)))); // REGISTER_INSTRUCTION_METER = *PreviousInstructionMeter;
         }
 
         // Test if result indicates that an error occured
@@ -1650,6 +1628,7 @@ mod tests {
         check_slot!(env, stack_pointer, StackPointer);
         check_slot!(env, context_object_pointer, ContextObjectPointer);
         check_slot!(env, previous_instruction_meter, PreviousInstructionMeter);
+        check_slot!(env, due_insn_count, DueInsnCount);
         check_slot!(env, stopwatch_numerator, StopwatchNumerator);
         check_slot!(env, stopwatch_denominator, StopwatchDenominator);
         check_slot!(env, registers, Registers);

src/program.rs

@@ -3,8 +3,7 @@ use {
     crate::{
         ebpf,
         elf::ElfError,
-        memory_region::MemoryMapping,
-        vm::{Config, ContextObject, ProgramResult},
+        vm::{Config, ContextObject, EbpfVm},
     },
     std::collections::{btree_map::Entry, BTreeMap},
 };
@@ -210,8 +209,7 @@ impl<T: Copy + PartialEq> FunctionRegistry<T> {
 }
 
 /// Syscall function without context
-pub type BuiltinFunction<C> =
-    fn(&mut C, u64, u64, u64, u64, u64, &mut MemoryMapping, &mut ProgramResult);
+pub type BuiltinFunction<C> = fn(*mut EbpfVm<C>, u64, u64, u64, u64, u64);
 
 /// Represents the interface to a fixed functionality program
 #[derive(Eq)]
@@ -279,10 +277,9 @@ impl<C: ContextObject> std::fmt::Debug for BuiltinProgram<C> {
     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<BuiltinFunction<C>>,
-                &FunctionRegistry<BuiltinFunction<*const ()>>,
-            >(&self.functions)
+            std::mem::transmute::<&FunctionRegistry<BuiltinFunction<C>>, &FunctionRegistry<usize>>(
+                &self.functions,
+            )
         })?;
         Ok(())
     }
@@ -291,28 +288,55 @@ impl<C: ContextObject> std::fmt::Debug for BuiltinProgram<C> {
 /// Generates an adapter for a BuiltinFunction between the Rust and the VM interface
 #[macro_export]
 macro_rules! declare_builtin_function {
-    ($(#[$attr:meta])* $name:ident, $rust:item) => {
+    ($(#[$attr:meta])* $name:ident, fn rust(
+        $vm:ident : &mut $ContextObject:ty,
+        $arg_a:ident : u64,
+        $arg_b:ident : u64,
+        $arg_c:ident : u64,
+        $arg_d:ident : u64,
+        $arg_e:ident : u64,
+        $memory_mapping:ident : &mut $MemoryMapping:ty,
+    ) -> Result<u64, EbpfError> $rust:tt) => {
         $(#[$attr])*
         pub struct $name {}
         impl $name {
             /// Rust interface
-            $rust
+            pub fn rust(
+                $vm: &mut $ContextObject,
+                $arg_a: u64,
+                $arg_b: u64,
+                $arg_c: u64,
+                $arg_d: u64,
+                $arg_e: u64,
+                $memory_mapping: &mut $MemoryMapping,
+            ) -> Result<u64, EbpfError> {
+                $rust
+            }
             /// VM interface
             #[allow(clippy::too_many_arguments)]
             pub fn vm(
-                context_object: &mut TestContextObject,
-                arg_a: u64,
-                arg_b: u64,
-                arg_c: u64,
-                arg_d: u64,
-                arg_e: u64,
-                memory_mapping: &mut $crate::memory_region::MemoryMapping,
-                program_result: &mut $crate::vm::ProgramResult,
+                $vm: *mut $crate::vm::EbpfVm<$ContextObject>,
+                $arg_a: u64,
+                $arg_b: u64,
+                $arg_c: u64,
+                $arg_d: u64,
+                $arg_e: u64,
             ) {
+                use $crate::vm::ContextObject;
+                let vm = unsafe {
+                    &mut *(($vm as *mut u64).offset(-($crate::vm::get_runtime_environment_key() as isize)) as *mut $crate::vm::EbpfVm<$ContextObject>)
+                };
+                let config = vm.loader.get_config();
+                if config.enable_instruction_meter {
+                    vm.context_object_pointer.consume(vm.previous_instruction_meter - vm.due_insn_count);
+                }
                 let converted_result: $crate::vm::ProgramResult = Self::rust(
-                    context_object, arg_a, arg_b, arg_c, arg_d, arg_e, memory_mapping,
+                    vm.context_object_pointer, $arg_a, $arg_b, $arg_c, $arg_d, $arg_e, &mut vm.memory_mapping,
                 ).into();
-                *program_result = converted_result;
+                vm.program_result = converted_result;
+                if config.enable_instruction_meter {
+                    vm.previous_instruction_meter = vm.context_object_pointer.get_remaining();
+                }
             }
         }
     };