Skip to content

Commit

Permalink
Refactor - vm interface of BuiltinFunctions (#535)
Browse files Browse the repository at this point in the history 
* 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!().
  • Loading branch information
@Lichtso
Lichtso authored Oct 1, 2023
1 parent b1b50f7 commit ae16614
Show file tree
Hide file tree
Showing 5 changed files with 91 additions and 92 deletions.
4 changes: 0 additions & 4 deletions fuzz/fuzz_targets/common.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
Expand All @@ -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,
})
Expand All @@ -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 {
Expand All @@ -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()
Expand Down
25 changes: 8 additions & 17 deletions src/interpreter.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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],
Expand All @@ -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,
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
}
Expand All @@ -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();
}
}
}

Expand All @@ -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]);
Expand All @@ -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);
}
Expand Down
59 changes: 19 additions & 40 deletions src/jit.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand All @@ -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
}
}

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand Down
64 changes: 44 additions & 20 deletions src/program.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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},
};
Expand Down Expand Up @@ -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)]
Expand Down Expand Up @@ -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(())
}
Expand All @@ -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();
}
}
}
};
Expand Down
Loading

0 comments on commit ae16614

Please sign in to comment.