Skip to content

Commit

Permalink
jit/x86: emit endbr64 on anchors
Browse files Browse the repository at this point in the history
  • Loading branch information
Richard Patel committed Mar 21, 2022
1 parent f4093e7 commit e36a585
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 19 deletions.
58 changes: 39 additions & 19 deletions src/jit.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -367,8 +367,9 @@ fn emit_call<E: UserDefinedError>(jit: &mut JitCompiler, target_pc: usize) -> Re
}

#[inline]
fn set_anchor(jit: &mut JitCompiler, target: usize) {
fn set_anchor<E: UserDefinedError>(jit: &mut JitCompiler, target: usize) -> Result<(), EbpfError<E>> {
jit.handler_anchors.insert(target, jit.offset_in_text_section);
emit_indirect_branch_target(jit)
}

/// Indices of slots inside the struct at inital RSP
Expand Down Expand Up @@ -866,6 +867,11 @@ fn emit_set_exception_kind<E: UserDefinedError>(jit: &mut JitCompiler, err: Ebpf
X86Instruction::store_immediate(OperandSize::S64, R10, X86IndirectAccess::Offset(8), err_kind as i64).emit(jit)
}

#[inline]
fn emit_indirect_branch_target<E: UserDefinedError>(jit: &mut JitCompiler) -> Result<(), EbpfError<E>> {
X86Instruction::end_branch().emit(jit)
}

#[derive(Debug)]
struct Jump {
location: usize,
Expand Down Expand Up @@ -980,6 +986,20 @@ impl JitCompiler {
})
}

#[cfg(test)]
pub(crate) fn new_mock() -> Self {
let res = Self::new::<UserError>(&[], &Config {
noop_instruction_ratio: 0.0,
..Config::default()
});
res.expect("failed to create mock JitCompiler")
}

#[cfg(test)]
pub(crate) fn get_text_result(&self) -> &[u8] {
self.result.text_section
}

fn compile<E: UserDefinedError, I: InstructionMeter>(&mut self,
executable: &Pin<Box<Executable<E, I>>>) -> Result<(), EbpfError<E>> {
let (program_vm_addr, program) = executable.get_text_bytes();
Expand Down Expand Up @@ -1335,7 +1355,7 @@ impl JitCompiler {
fn generate_helper_routines<E: UserDefinedError, I: InstructionMeter>(&mut self) -> Result<(), EbpfError<E>> {
// Routine for instruction tracing
if self.config.enable_instruction_tracing {
set_anchor(self, TARGET_PC_TRACE);
set_anchor(self, TARGET_PC_TRACE)?;
// Save registers on stack
X86Instruction::push(R11, None).emit(self)?;
for reg in REGISTER_MAP.iter().rev() {
Expand All @@ -1356,7 +1376,7 @@ impl JitCompiler {
}

// Routine for syscall
set_anchor(self, TARGET_PC_SYSCALL);
set_anchor(self, TARGET_PC_SYSCALL)?;
X86Instruction::push(R11, None).emit(self)?; // Padding for stack alignment
if self.config.enable_instruction_meter {
// RDI = *PrevInsnMeter - RDI;
Expand Down Expand Up @@ -1390,7 +1410,7 @@ impl JitCompiler {
X86Instruction::return_near().emit(self)?;

// Routine for prologue of emit_bpf_call()
set_anchor(self, TARGET_PC_BPF_CALL_PROLOGUE);
set_anchor(self, TARGET_PC_BPF_CALL_PROLOGUE)?;
emit_alu(self, OperandSize::S64, 0x81, 5, RSP, 8 * (SCRATCH_REGS + 1) as i64, None)?; // alloca
X86Instruction::store(OperandSize::S64, R11, RSP, X86IndirectAccess::OffsetIndexShift(0, RSP, 0)).emit(self)?; // Save original R11
X86Instruction::load(OperandSize::S64, RSP, R11, X86IndirectAccess::OffsetIndexShift(8 * (SCRATCH_REGS + 1) as i32, RSP, 0)).emit(self)?; // Load return address
Expand All @@ -1409,7 +1429,7 @@ impl JitCompiler {
X86Instruction::return_near().emit(self)?;

// Routine for emit_bpf_call(Value::Register())
set_anchor(self, TARGET_PC_BPF_CALL_REG);
set_anchor(self, TARGET_PC_BPF_CALL_REG)?;
// Force alignment of RAX
emit_alu(self, OperandSize::S64, 0x81, 4, REGISTER_MAP[0], !(INSN_SIZE as i64 - 1), None)?; // RAX &= !(INSN_SIZE - 1);
// Upper bound check
Expand Down Expand Up @@ -1442,10 +1462,10 @@ impl JitCompiler {
X86Instruction::return_near().emit(self)?;

// Translates a host pc back to a BPF pc by linear search of the pc_section table
set_anchor(self, TARGET_PC_TRANSLATE_PC);
set_anchor(self, TARGET_PC_TRANSLATE_PC)?;
X86Instruction::push(REGISTER_MAP[0], None).emit(self)?; // Save REGISTER_MAP[0]
X86Instruction::load_immediate(OperandSize::S64, REGISTER_MAP[0], self.result.pc_section.as_ptr() as i64 - 8).emit(self)?; // Loop index and pointer to look up
set_anchor(self, TARGET_PC_TRANSLATE_PC_LOOP); // Loop label
set_anchor(self, TARGET_PC_TRANSLATE_PC_LOOP)?; // Loop label
emit_alu(self, OperandSize::S64, 0x81, 0, REGISTER_MAP[0], 8, None)?; // Increase index
X86Instruction::cmp(OperandSize::S64, R11, REGISTER_MAP[0], Some(X86IndirectAccess::Offset(8))).emit(self)?; // Look up and compare against value at next index
emit_jcc(self, 0x86, TARGET_PC_TRANSLATE_PC_LOOP)?; // Continue while *REGISTER_MAP[0] <= R11
Expand All @@ -1469,7 +1489,7 @@ impl JitCompiler {
] {
let target_offset = len.trailing_zeros() as usize + 4 * (*access_type as usize);

set_anchor(self, TARGET_PC_TRANSLATE_MEMORY_ADDRESS + target_offset);
set_anchor(self, TARGET_PC_TRANSLATE_MEMORY_ADDRESS + target_offset)?;
X86Instruction::push(R11, None).emit(self)?;
X86Instruction::push(RAX, None).emit(self)?;
X86Instruction::push(RCX, None).emit(self)?;
Expand Down Expand Up @@ -1519,7 +1539,7 @@ impl JitCompiler {
emit_alu(self, OperandSize::S64, 0x81, 0, RSP, 8, None)?;
X86Instruction::return_near().emit(self)?;

set_anchor(self, TARGET_PC_MEMORY_ACCESS_VIOLATION + target_offset);
set_anchor(self, TARGET_PC_MEMORY_ACCESS_VIOLATION + target_offset)?;
emit_alu(self, OperandSize::S64, 0x31, R11, R11, 0, None)?; // R11 = 0;
X86Instruction::load(OperandSize::S64, RSP, R11, X86IndirectAccess::OffsetIndexShift(stack_offset, R11, 0)).emit(self)?;
emit_rust_call(self, Value::Constant64(MemoryMapping::generate_access_violation::<UserError> as *const u8 as i64, false), &[
Expand All @@ -1539,45 +1559,45 @@ impl JitCompiler {

fn generate_exception_handlers<E: UserDefinedError>(&mut self) -> Result<(), EbpfError<E>> {
// Handler for EbpfError::ExceededMaxInstructions
set_anchor(self, TARGET_PC_CALL_EXCEEDED_MAX_INSTRUCTIONS);
set_anchor(self, TARGET_PC_CALL_EXCEEDED_MAX_INSTRUCTIONS)?;
emit_set_exception_kind::<E>(self, EbpfError::ExceededMaxInstructions(0, 0))?;
X86Instruction::mov(OperandSize::S64, ARGUMENT_REGISTERS[0], R11).emit(self)?; // R11 = instruction_meter;
emit_profile_instruction_count_finalize(self, true)?;
emit_jmp(self, TARGET_PC_EPILOGUE)?;

// Handler for EbpfError::CallDepthExceeded
set_anchor(self, TARGET_PC_CALL_DEPTH_EXCEEDED);
set_anchor(self, TARGET_PC_CALL_DEPTH_EXCEEDED)?;
emit_set_exception_kind::<E>(self, EbpfError::CallDepthExceeded(0, 0))?;
X86Instruction::store_immediate(OperandSize::S64, R10, X86IndirectAccess::Offset(24), self.config.max_call_depth as i64).emit(self)?; // depth = jit.config.max_call_depth;
emit_jmp(self, TARGET_PC_EXCEPTION_AT)?;

// Handler for EbpfError::CallOutsideTextSegment
set_anchor(self, TARGET_PC_CALL_OUTSIDE_TEXT_SEGMENT);
set_anchor(self, TARGET_PC_CALL_OUTSIDE_TEXT_SEGMENT)?;
emit_set_exception_kind::<E>(self, EbpfError::CallOutsideTextSegment(0, 0))?;
X86Instruction::store(OperandSize::S64, REGISTER_MAP[0], R10, X86IndirectAccess::Offset(24)).emit(self)?; // target_address = RAX;
emit_jmp(self, TARGET_PC_EXCEPTION_AT)?;

// Handler for EbpfError::DivideByZero
set_anchor(self, TARGET_PC_DIV_BY_ZERO);
set_anchor(self, TARGET_PC_DIV_BY_ZERO)?;
emit_set_exception_kind::<E>(self, EbpfError::DivideByZero(0))?;
emit_jmp(self, TARGET_PC_EXCEPTION_AT)?;

// Handler for EbpfError::UnsupportedInstruction
set_anchor(self, TARGET_PC_CALLX_UNSUPPORTED_INSTRUCTION);
set_anchor(self, TARGET_PC_CALLX_UNSUPPORTED_INSTRUCTION)?;
// Load BPF target pc from stack (which was saved in TARGET_PC_BPF_CALL_REG)
X86Instruction::load(OperandSize::S64, RSP, R11, X86IndirectAccess::OffsetIndexShift(-16, RSP, 0)).emit(self)?; // R11 = RSP[-16];
// emit_jmp(self, TARGET_PC_CALL_UNSUPPORTED_INSTRUCTION)?; // Fall-through

// Handler for EbpfError::UnsupportedInstruction
set_anchor(self, TARGET_PC_CALL_UNSUPPORTED_INSTRUCTION);
set_anchor(self, TARGET_PC_CALL_UNSUPPORTED_INSTRUCTION)?;
if self.config.enable_instruction_tracing {
emit_call(self, TARGET_PC_TRACE)?;
}
emit_set_exception_kind::<E>(self, EbpfError::UnsupportedInstruction(0))?;
// emit_jmp(self, TARGET_PC_EXCEPTION_AT)?; // Fall-through

// Handler for exceptions which report their pc
set_anchor(self, TARGET_PC_EXCEPTION_AT);
set_anchor(self, TARGET_PC_EXCEPTION_AT)?;
// Validate that we did not reach the instruction meter limit before the exception occured
if self.config.enable_instruction_meter {
emit_validate_instruction_count(self, false, None)?;
Expand All @@ -1586,7 +1606,7 @@ impl JitCompiler {
emit_jmp(self, TARGET_PC_EPILOGUE)?;

// Handler for syscall exceptions
set_anchor(self, TARGET_PC_RUST_EXCEPTION);
set_anchor(self, TARGET_PC_RUST_EXCEPTION)?;
emit_profile_instruction_count_finalize(self, false)?;
emit_jmp(self, TARGET_PC_EPILOGUE)
}
Expand Down Expand Up @@ -1639,7 +1659,7 @@ impl JitCompiler {

fn generate_epilogue<E: UserDefinedError>(&mut self) -> Result<(), EbpfError<E>> {
// Quit gracefully
set_anchor(self, TARGET_PC_EXIT);
set_anchor(self, TARGET_PC_EXIT)?;
emit_validate_instruction_count(self, false, None)?;
emit_profile_instruction_count_finalize(self, false)?;

Expand All @@ -1649,7 +1669,7 @@ impl JitCompiler {
X86Instruction::store(OperandSize::S64, REGISTER_MAP[0], R10, X86IndirectAccess::Offset(0)).emit(self)?; // result.is_error = false;

// Epilogue
set_anchor(self, TARGET_PC_EPILOGUE);
set_anchor(self, TARGET_PC_EPILOGUE)?;

// Print stop watch value
fn stopwatch_result(numerator: u64, denominator: u64) {
Expand Down
47 changes: 47 additions & 0 deletions src/x86.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ pub enum FenceType {
#[derive(PartialEq, Eq, Copy, Clone)]
pub struct X86Instruction {
pub size: OperandSize,
pub repeat_string: Option<bool>,
pub opcode_escape_sequence: u8,
pub opcode: u8,
pub modrm: bool,
Expand All @@ -82,6 +83,7 @@ impl Default for X86Instruction {
fn default() -> Self {
Self {
size: OperandSize::S64,
repeat_string: None,
opcode_escape_sequence: 0,
opcode: 0,
modrm: true,
Expand All @@ -96,6 +98,11 @@ impl Default for X86Instruction {

impl X86Instruction {
pub fn emit<E: UserDefinedError>(&self, jit: &mut JitCompiler) -> Result<(), EbpfError<E>> {
match self.repeat_string {
Some(false) => emit(jit, 0xF2u8)?, // repnz
Some(true) => emit(jit, 0xF3u8)?, // repz
_ => (),
}
let mut rex = X86Rex {
w: self.size == OperandSize::S64,
r: self.first_operand & 0b1000 != 0,
Expand Down Expand Up @@ -581,4 +588,44 @@ impl X86Instruction {
..Self::default()
}
}

/// endbr64
///
/// aka Intel CET "Terminate Indirect Branch".
/// Marks a valid indirect branch target.
#[allow(dead_code)]
pub fn end_branch() -> Self {
Self {
// repz
repeat_string: Some(true),
// nop
opcode_escape_sequence: 1,
opcode: 0x1e,
size: OperandSize::S0,
// edx
modrm: true,
first_operand: 7, // reserved
second_operand: 2, // edx
..Self::default()
}
}
}

#[cfg(test)]
#[cfg(not(windows))]
mod tests {
use super::*;
use crate::user_error::UserError;

#[test]
fn test_end_branch() {
let mut compiler = JitCompiler::new_mock();
X86Instruction::end_branch()
.emit::<UserError>(&mut compiler)
.unwrap();
assert_eq!(
&compiler.get_text_result()[..4],
&[0xF3u8, 0x0Fu8, 0x1Eu8, 0xFAu8]
);
}
}

0 comments on commit e36a585

Please sign in to comment.