Moves second half of emit_sanitized_load_immediate(stack_slot_of_value_to_store, constant) into ANCHOR_TRANSLATE_MEMORY_ADDRESS.

Author: Lichtso

src/jit.rs

@@ -30,7 +30,7 @@ use crate::{
     memory_management::{
         allocate_pages, free_pages, get_system_page_size, protect_pages, round_to_page_size,
     },
-    memory_region::{AccessType, MemoryMapping},
+    memory_region::MemoryMapping,
     program::BuiltinFunction,
     vm::{get_runtime_environment_key, Config, ContextObject, EbpfVm},
     x86::*,
@@ -200,7 +200,7 @@ const ANCHOR_INTERNAL_FUNCTION_CALL_PROLOGUE: usize = 12;
 const ANCHOR_INTERNAL_FUNCTION_CALL_REG: usize = 13;
 const ANCHOR_CALL_REG_UNSUPPORTED_INSTRUCTION: usize = 14;
 const ANCHOR_TRANSLATE_MEMORY_ADDRESS: usize = 21;
-const ANCHOR_COUNT: usize = 30; // Update me when adding or removing anchors
+const ANCHOR_COUNT: usize = 34; // Update me when adding or removing anchors
 
 const REGISTER_MAP: [u8; 11] = [
     CALLER_SAVED_REGISTERS[0], // RAX
@@ -1155,11 +1155,10 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
                 self.emit_ins(X86Instruction::store(OperandSize::S64, reg, RSP, stack_slot_of_value_to_store));
             }
             Some(Value::Constant64(constant, user_provided)) => {
-                if user_provided && self.should_sanitize_constant(constant) {
-                    self.emit_sanitized_load_immediate(REGISTER_SCRATCH, constant);
-                } else {
-                    self.emit_ins(X86Instruction::load_immediate(REGISTER_SCRATCH, constant));
-                }
+                debug_assert!(user_provided);
+                // First half of emit_sanitized_load_immediate()
+                let lower_key = self.immediate_value_key as i32 as i64;
+                self.emit_ins(X86Instruction::load_immediate(REGISTER_SCRATCH, constant.wrapping_sub(lower_key)));
                 self.emit_ins(X86Instruction::store(OperandSize::S64, REGISTER_SCRATCH, RSP, stack_slot_of_value_to_store));
             }
             _ => {}
@@ -1185,8 +1184,12 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         }
 
         if self.config.enable_address_translation {
-            let access_type = if value.is_none() { AccessType::Load } else { AccessType::Store };
-            let anchor = ANCHOR_TRANSLATE_MEMORY_ADDRESS + len.trailing_zeros() as usize + 4 * (access_type as usize);
+            let access_type = match value {
+                Some(Value::Register(_reg)) => 4,
+                Some(Value::Constant64(_constant, _user_provided)) => 8,
+                _ => 0,
+            };
+            let anchor = ANCHOR_TRANSLATE_MEMORY_ADDRESS + access_type + len.trailing_zeros() as usize;
             self.emit_ins(X86Instruction::push_immediate(OperandSize::S64, self.pc as i32));
             self.emit_ins(X86Instruction::call_immediate(self.relative_to_anchor(anchor, 5)));
             if let Some(dst) = dst {
@@ -1597,21 +1600,16 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         // Translates a vm memory address to a host memory address
         let lower_key = self.immediate_value_key as i32 as i64;
         for (access_type, len) in &[
-            (AccessType::Load, 1i32),
-            (AccessType::Load, 2i32),
-            (AccessType::Load, 4i32),
-            (AccessType::Load, 8i32),
-            (AccessType::Store, 1i32),
-            (AccessType::Store, 2i32),
-            (AccessType::Store, 4i32),
-            (AccessType::Store, 8i32),
+            (0, 1i32), (0, 2i32), (0, 4i32), (0, 8i32),
+            (4, 1i32), (4, 2i32), (4, 4i32), (4, 8i32),
+            (8, 1i32), (8, 2i32), (8, 4i32), (8, 8i32),
         ] {
-            let target_offset = len.trailing_zeros() as usize + 4 * (*access_type as usize);
+            let target_offset = *access_type + len.trailing_zeros() as usize;
             self.set_anchor(ANCHOR_TRANSLATE_MEMORY_ADDRESS + target_offset);
             // Second half of emit_sanitized_load_immediate(REGISTER_SCRATCH, vm_addr)
             self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x81, 0, REGISTER_SCRATCH, lower_key, None));
             // call MemoryMapping::(load|store) storing the result in RuntimeEnvironmentSlot::ProgramResult
-            if *access_type == AccessType::Load {
+            if *access_type == 0 { // AccessType::Load
                 let load = match len {
                     1 => MemoryMapping::load::<u8> as *const u8 as i64,
                     2 => MemoryMapping::load::<u16> as *const u8 as i64,
@@ -1625,7 +1623,11 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
                     Argument { index: 1, value: Value::RegisterPlusConstant32(REGISTER_PTR_TO_VM, self.slot_in_vm(RuntimeEnvironmentSlot::MemoryMapping), false) },
                     Argument { index: 0, value: Value::RegisterPlusConstant32(REGISTER_PTR_TO_VM, self.slot_in_vm(RuntimeEnvironmentSlot::ProgramResult), false) },
                 ], None);
-            } else {
+            } else { // AccessType::Store
+                if *access_type == 8 {
+                    // Second half of emit_sanitized_load_immediate()
+                    self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x81, 0, RSP, lower_key, Some(X86IndirectAccess::OffsetIndexShift(-96, RSP, 0))));
+                }
                 let store = match len {
                     1 => MemoryMapping::store::<u8> as *const u8 as i64,
                     2 => MemoryMapping::store::<u16> as *const u8 as i64,
@@ -1648,7 +1650,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
             self.emit_ins(X86Instruction::xchg(OperandSize::S64, REGISTER_SCRATCH, RSP, Some(X86IndirectAccess::OffsetIndexShift(0, RSP, 0)))); // Swap return address and self.pc
             self.emit_ins(X86Instruction::conditional_jump_immediate(0x85, self.relative_to_anchor(ANCHOR_THROW_EXCEPTION, 6)));
 
-            if *access_type == AccessType::Load {
+            if *access_type == 0 { // AccessType::Load
                 // unwrap() the result into REGISTER_SCRATCH
                 self.emit_ins(X86Instruction::load(OperandSize::S64, REGISTER_PTR_TO_VM, REGISTER_SCRATCH, X86IndirectAccess::Offset(self.slot_in_vm(RuntimeEnvironmentSlot::ProgramResult) + std::mem::size_of::<u64>() as i32)));
             }