Fix callx error discrepancy between interpreter and jit

src/interpreter.rs

@@ -519,8 +519,6 @@ impl<'a, 'b, C: ContextObject> Interpreter<'a, 'b, C> {
                 }
                 check_pc!(self, next_pc, target_pc.wrapping_sub(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.vm.due_insn_count += 1;
-                    self.reg[11] = next_pc;
                     throw_error!(self, EbpfError::UnsupportedInstruction);
                 }
             },

src/jit.rs

@@ -198,6 +198,7 @@ const ANCHOR_CALL_UNSUPPORTED_INSTRUCTION: usize = 10;
 const ANCHOR_EXTERNAL_FUNCTION_CALL: usize = 11;
 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
 
@@ -978,9 +979,20 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
     }
 
     #[inline]
-    fn emit_undo_profile_instruction_count(&mut self, target_pc: usize) {
+    fn emit_undo_profile_instruction_count(&mut self, target_pc: Value) {
         if self.config.enable_instruction_meter {
-            self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x81, 0, REGISTER_INSTRUCTION_METER, self.pc as i64 + 1 - target_pc as i64, None)); // instruction_meter += (self.pc + 1) - target_pc;
+            match target_pc {
+                Value::Constant64(target_pc, _) => {
+                    self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x81, 0, REGISTER_INSTRUCTION_METER, self.pc as i64 + 1 - target_pc, None)); // instruction_meter += (self.pc + 1) - target_pc;
+                }
+                Value::Register(reg) => {
+                    self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x29, reg, REGISTER_INSTRUCTION_METER, 0, None)); // instruction_meter -= guest_target_pc
+                    self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x81, 0, REGISTER_INSTRUCTION_METER, 1, None)); // instruction_meter += 1
+                    self.emit_ins(X86Instruction::load(OperandSize::S64, RSP, REGISTER_SCRATCH, X86IndirectAccess::OffsetIndexShift(-16, RSP, 0)));
+                    self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x01, REGISTER_SCRATCH, REGISTER_INSTRUCTION_METER, 0, None)); // instruction_meter += guest_current_pc
+                }
+                _ => {},
+            }
         }
     }
 
@@ -1124,7 +1136,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
             }
         }
 
-        self.emit_undo_profile_instruction_count(0);
+        self.emit_undo_profile_instruction_count(Value::Constant64(0, false));
 
         // Restore the previous frame pointer
         self.emit_ins(X86Instruction::pop(REGISTER_MAP[FRAME_PTR_REG]));
@@ -1138,7 +1150,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         self.emit_validate_and_profile_instruction_count(false, Some(0));
         self.emit_ins(X86Instruction::load_immediate(OperandSize::S64, REGISTER_SCRATCH, function as usize as i64));
         self.emit_ins(X86Instruction::call_immediate(self.relative_to_anchor(ANCHOR_EXTERNAL_FUNCTION_CALL, 5)));
-        self.emit_undo_profile_instruction_count(0);
+        self.emit_undo_profile_instruction_count(Value::Constant64(0, false));
     }
 
     #[inline]
@@ -1223,7 +1235,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         self.emit_ins(X86Instruction::load_immediate(OperandSize::S64, REGISTER_SCRATCH, target_pc as i64));
         let jump_offset = self.relative_to_target_pc(target_pc, 6);
         self.emit_ins(X86Instruction::conditional_jump_immediate(op, jump_offset));
-        self.emit_undo_profile_instruction_count(target_pc);
+        self.emit_undo_profile_instruction_count(Value::Constant64(target_pc as i64, false));
     }
 
     #[inline]
@@ -1244,7 +1256,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         self.emit_ins(X86Instruction::load_immediate(OperandSize::S64, REGISTER_SCRATCH, target_pc as i64));
         let jump_offset = self.relative_to_target_pc(target_pc, 6);
         self.emit_ins(X86Instruction::conditional_jump_immediate(op, jump_offset));
-        self.emit_undo_profile_instruction_count(target_pc);
+        self.emit_undo_profile_instruction_count(Value::Constant64(target_pc as i64, false));
     }
 
     fn emit_shift(&mut self, size: OperandSize, opcode_extension: u8, source: u8, destination: u8, immediate: Option<i64>) {
@@ -1549,6 +1561,10 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         // Outputs: Guest target pc in REGISTER_SCRATCH, Host target address in RIP
         self.set_anchor(ANCHOR_INTERNAL_FUNCTION_CALL_REG);
         self.emit_ins(X86Instruction::push(REGISTER_MAP[0], None));
+        // REGISTER_SCRATCH contains the current program counter, and we must store it for proper
+        // error handling. We can discard the value if callx succeeds, so we are not incrementing
+        // the stack pointer (RSP).
+        self.emit_ins(X86Instruction::store(OperandSize::S64, REGISTER_SCRATCH, RSP, X86IndirectAccess::OffsetIndexShift(-8, RSP, 0)));
         self.emit_ins(X86Instruction::mov(OperandSize::S64, REGISTER_MAP[FRAME_PTR_REG], REGISTER_MAP[0]));
         // Calculate offset relative to program_vm_addr
         self.emit_ins(X86Instruction::load_immediate(OperandSize::S64, REGISTER_MAP[FRAME_PTR_REG], self.program_vm_addr as i64));
@@ -1582,6 +1598,16 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         self.emit_ins(X86Instruction::xchg(OperandSize::S64, REGISTER_MAP[0], RSP, Some(X86IndirectAccess::OffsetIndexShift(0, RSP, 0)))); // Swap REGISTER_MAP[0] and host_target_address
         self.emit_ins(X86Instruction::return_near()); // Tail call to host_target_address
 
+        // If callx lands in an invalid address, we must undo the changes in the instruction meter
+        // so that we can correctly calculate the number of executed instructions for error handling.
+        self.set_anchor(ANCHOR_CALL_REG_UNSUPPORTED_INSTRUCTION);
+        self.emit_ins(X86Instruction::mov(OperandSize::S64, REGISTER_SCRATCH, REGISTER_MAP[0]));
+        // Retrieve the current program from the stack. `return_near` popped an element from the stack,
+        // so the offset is 16. Check `ANCHOR_INTERNAL_FUNCTION_CALL_REG` for more details.
+        self.emit_ins(X86Instruction::load(OperandSize::S64, RSP, REGISTER_SCRATCH, X86IndirectAccess::OffsetIndexShift(-16, RSP, 0)));
+        self.emit_undo_profile_instruction_count(Value::Register(REGISTER_MAP[0]));
+        self.emit_ins(X86Instruction::jump_immediate(self.relative_to_anchor(ANCHOR_CALL_UNSUPPORTED_INSTRUCTION, 5)));
+
         // Translates a vm memory address to a host memory address
         for (access_type, len) in &[
             (AccessType::Load, 1i32),
@@ -1678,8 +1704,8 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
                 - mem::size_of::<i32>() as i32; // Jump from end of instruction
             unsafe { ptr::write_unaligned(jump.location as *mut i32, offset_value); }
         }
-        // There is no `VerifierError::JumpToMiddleOfLDDW` for `call imm` so patch it here
-        let call_unsupported_instruction = self.anchors[ANCHOR_CALL_UNSUPPORTED_INSTRUCTION] as usize;
+        // Patch addresses to which `callx` may raise an unsupported instruction error
+        let call_unsupported_instruction = self.anchors[ANCHOR_CALL_REG_UNSUPPORTED_INSTRUCTION] as usize;
         if self.executable.get_sbpf_version().static_syscalls() {
             let mut prev_pc = 0;
             for current_pc in self.executable.get_function_registry().keys() {

tests/execution.rs

@@ -41,14 +41,31 @@ macro_rules! test_interpreter_and_jit {
             .unwrap();
     };
     ($executable:expr, $mem:tt, $context_object:expr, $expected_result:expr $(,)?) => {
-        test_interpreter_and_jit!(true, $executable, $mem, $context_object, $expected_result)
+        test_interpreter_and_jit!(
+            false,
+            true,
+            $executable,
+            $mem,
+            $context_object,
+            $expected_result
+        )
     };
     ($verify:literal, $executable:expr, $mem:tt, $context_object:expr, $expected_result:expr $(,)?) => {
+        test_interpreter_and_jit!(
+            false,
+            $verify,
+            $executable,
+            $mem,
+            $context_object,
+            $expected_result
+        )
+    };
+    ($override_budget:literal, $verify:literal, $executable:expr, $mem:tt, $context_object:expr, $expected_result:expr $(,)?) => {
         let expected_instruction_count = $context_object.get_remaining();
         #[allow(unused_mut)]
         let mut context_object = $context_object;
         let expected_result = format!("{:?}", $expected_result);
-        if !expected_result.contains("ExceededMaxInstructions") {
+        if !$override_budget && !expected_result.contains("ExceededMaxInstructions") {
             context_object.remaining = INSTRUCTION_METER_BUDGET;
         }
         if $verify {
@@ -2347,18 +2364,74 @@ fn test_callx() {
 
 #[test]
 fn test_err_callx_unregistered() {
+    let config = Config {
+        enabled_sbpf_versions: SBPFVersion::V1..=SBPFVersion::V1,
+        ..Config::default()
+    };
+
+    // Callx jumps to `mov64 r0, 0x2A`
     test_interpreter_and_jit_asm!(
         "
         mov64 r0, 0x0
-        or64 r8, 0x20
+        lddw r8, 0x100000028
         callx r8
         exit
         mov64 r0, 0x2A
         exit",
+        config,
         [],
-        TestContextObject::new(4),
+        TestContextObject::new(6),
+        ProgramResult::Ok(42),
+    );
+
+    let config = Config {
+        enabled_sbpf_versions: SBPFVersion::V2..=SBPFVersion::V2,
+        ..Config::default()
+    };
+
+    // Callx jumps to `mov64 r0, 0x2A`
+    test_interpreter_and_jit_asm!(
+        "
+        mov64 r0, 0x0
+        or64 r8, 0x28
+        callx r8
+        exit
+        mov64 r0, 0x2A
+        exit",
+        config,
+        [],
+        TestContextObject::new(3),
         ProgramResult::Err(EbpfError::UnsupportedInstruction),
     );
+
+    let versions = [SBPFVersion::V1, SBPFVersion::V2];
+    let expected_errors = [
+        EbpfError::CallOutsideTextSegment,
+        EbpfError::UnsupportedInstruction,
+    ];
+
+    // We execute three instructions when callx errors out.
+    for (version, error) in versions.iter().zip(expected_errors) {
+        let config = Config {
+            enabled_sbpf_versions: *version..=*version,
+            ..Config::default()
+        };
+
+        // Callx jumps to a location outside text segment
+        test_interpreter_and_jit_asm!(
+            "
+            mov64 r0, 0x0
+            or64 r8, 0x20
+            callx r8
+            exit
+            mov64 r0, 0x2A
+            exit",
+            config,
+            [],
+            TestContextObject::new(3),
+            ProgramResult::Err(error),
+        );
+    }
 }
 
 #[test]
@@ -3446,6 +3519,56 @@ fn test_invalid_exit_or_return() {
     }
 }
 
+#[test]
+fn callx_unsupported_instruction_and_exceeded_max_instructions() {
+    let program = "
+        sub32 r7, r1
+        sub64 r5, 8
+        sub64 r7, 0
+        callx r5
+        callx r5
+        return
+        ";
+    test_interpreter_and_jit_asm!(
+        program,
+        [],
+        TestContextObject::new(4),
+        ProgramResult::Err(EbpfError::UnsupportedInstruction),
+    );
+
+    let loader = Arc::new(BuiltinProgram::new_loader(
+        Config::default(),
+        FunctionRegistry::default(),
+    ));
+
+    let mut executable = assemble(program, loader).unwrap();
+    test_interpreter_and_jit!(
+        true,
+        false,
+        executable,
+        [],
+        TestContextObject::new(4),
+        ProgramResult::Err(EbpfError::UnsupportedInstruction)
+    );
+}
+
+#[test]
+fn test_capped_after_callx() {
+    test_interpreter_and_jit_asm!(
+        "
+        mov64 r0, 0x0
+        or64 r8, 0x20
+        callx r8
+        exit
+        function_foo:
+        mov64 r0, 0x2A
+        exit",
+        [],
+        TestContextObject::new(3),
+        ProgramResult::Err(EbpfError::ExceededMaxInstructions),
+    );
+}
+
 // SBPFv1 only [DEPRECATED]
 
 #[test]