Skip to content

Commit

Permalink
Stores u32 instead of u64 elements in pc_section.
Browse files Browse the repository at this point in the history
  • Loading branch information
@Lichtso
Lichtso committed Dec 31, 2024
1 parent 1bff78a commit 305a3d7
Showing 1 changed file with 20 additions and 11 deletions.
31 changes: 20 additions & 11 deletions src/jit.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,22 +48,22 @@ const MAX_START_PADDING_LENGTH: usize = 256;
pub struct JitProgram {
/// OS page size in bytes and the alignment of the sections
page_size: usize,
/// A `*const u8` pointer into the text_section for each BPF instruction
pc_section: &'static mut [usize],
/// Byte offset in the text_section for each BPF instruction
pc_section: &'static mut [u32],
/// The x86 machinecode
text_section: &'static mut [u8],
}

impl JitProgram {
fn new(pc: usize, code_size: usize) -> Result<Self, EbpfError> {
let page_size = get_system_page_size();
let pc_loc_table_size = round_to_page_size(pc * 8, page_size);
let pc_loc_table_size = round_to_page_size(pc * std::mem::size_of::<u32>(), page_size);
let over_allocated_code_size = round_to_page_size(code_size, page_size);
unsafe {
let raw = allocate_pages(pc_loc_table_size + over_allocated_code_size)?;
Ok(Self {
page_size,
pc_section: std::slice::from_raw_parts_mut(raw.cast::<usize>(), pc),
pc_section: std::slice::from_raw_parts_mut(raw.cast::<u32>(), pc),
text_section: std::slice::from_raw_parts_mut(
raw.add(pc_loc_table_size),
over_allocated_code_size,
Expand All @@ -77,7 +77,10 @@ impl JitProgram {
return Ok(());
}
let raw = self.pc_section.as_ptr() as *mut u8;
let pc_loc_table_size = round_to_page_size(self.pc_section.len() * 8, self.page_size);
let pc_loc_table_size = round_to_page_size(
self.pc_section.len() * std::mem::size_of::<u32>(),
self.page_size,
);
let over_allocated_code_size = round_to_page_size(self.text_section.len(), self.page_size);
let code_size = round_to_page_size(text_section_usage, self.page_size);
unsafe {
Expand Down Expand Up @@ -153,15 +156,21 @@ impl JitProgram {
}

pub fn mem_size(&self) -> usize {
let pc_loc_table_size = round_to_page_size(self.pc_section.len() * 8, self.page_size);
let pc_loc_table_size = round_to_page_size(
self.pc_section.len() * std::mem::size_of::<u32>(),
self.page_size,
);
let code_size = round_to_page_size(self.text_section.len(), self.page_size);
pc_loc_table_size + code_size
}
}

impl Drop for JitProgram {
fn drop(&mut self) {
let pc_loc_table_size = round_to_page_size(self.pc_section.len() * 8, self.page_size);
let pc_loc_table_size = round_to_page_size(
self.pc_section.len() * std::mem::size_of::<u32>(),
self.page_size,
);
let code_size = round_to_page_size(self.text_section.len(), self.page_size);
if pc_loc_table_size + code_size > 0 {
unsafe {
Expand Down Expand Up @@ -409,7 +418,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
return Err(EbpfError::ExhaustedTextSegment(self.pc));
}
let mut insn = ebpf::get_insn_unchecked(self.program, self.pc);
self.result.pc_section[self.pc] = self.offset_in_text_section;
self.result.pc_section[self.pc] = self.offset_in_text_section as u32;

// Regular instruction meter checkpoints to prevent long linear runs from exceeding their budget
if self.last_instruction_meter_validation_pc + self.config.instruction_meter_checkpoint_distance <= self.pc {
Expand All @@ -430,7 +439,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
ebpf::LD_DW_IMM if !self.executable.get_sbpf_version().disable_lddw() => {
self.emit_validate_and_profile_instruction_count(Some(self.pc + 2));
self.pc += 1;
self.result.pc_section[self.pc] = unsafe { self.anchors[ANCHOR_CALL_UNSUPPORTED_INSTRUCTION].offset_from(self.result.text_section.as_ptr()) as _ };
self.result.pc_section[self.pc] = unsafe { self.anchors[ANCHOR_CALL_UNSUPPORTED_INSTRUCTION].offset_from(self.result.text_section.as_ptr()) as u32 };
ebpf::augment_lddw_unchecked(self.program, &mut insn);
if self.should_sanitize_constant(insn.imm) {
self.emit_sanitized_load_immediate(dst, insn.imm);
Expand Down Expand Up @@ -1567,7 +1576,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x01, REGISTER_SCRATCH, REGISTER_INSTRUCTION_METER, None)); // instruction_meter += guest_target_pc;
// Load host target_address offset from self.result.pc_section
self.emit_ins(X86Instruction::load_immediate(REGISTER_MAP[0], self.result.pc_section.as_ptr() as i64)); // host_target_address = self.result.pc_section;
self.emit_ins(X86Instruction::load(OperandSize::S64, REGISTER_MAP[0], REGISTER_MAP[0], X86IndirectAccess::OffsetIndexShift(0, REGISTER_SCRATCH, shift_amount as u8))); // host_target_address = self.result.pc_section[guest_target_pc];
self.emit_ins(X86Instruction::load(OperandSize::S32, REGISTER_MAP[0], REGISTER_MAP[0], X86IndirectAccess::OffsetIndexShift(0, REGISTER_SCRATCH, 2))); // host_target_address = self.result.pc_section[guest_target_pc];
// Offset host target_address by self.result.text_section
self.emit_ins(X86Instruction::alu_immediate(OperandSize::S64, 0xc1, 1, REGISTER_MAP[0], 32, None)); // rotate_right(32)
self.emit_ins(X86Instruction::alu_immediate(OperandSize::S64, 0x81, 0, REGISTER_MAP[0], (self.result.text_section.as_ptr() as u64 >> 32) as u32 as i64, None)); // wrapping_add(self.result.text_section >> 32)
Expand Down Expand Up @@ -1688,7 +1697,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
unsafe { ptr::write_unaligned(jump.location as *mut i32, offset_value); }
}
// Patch addresses to which `callx` may raise an unsupported instruction error
let call_unsupported_instruction = unsafe { self.anchors[ANCHOR_CALL_REG_UNSUPPORTED_INSTRUCTION].offset_from(self.result.text_section.as_ptr()) as _ };
let call_unsupported_instruction = unsafe { self.anchors[ANCHOR_CALL_REG_UNSUPPORTED_INSTRUCTION].offset_from(self.result.text_section.as_ptr()) as u32 };
if self.executable.get_sbpf_version().static_syscalls() {
let mut prev_pc = 0;
for current_pc in self.executable.get_function_registry().keys() {
Expand Down

0 comments on commit 305a3d7

Please sign in to comment.