src/jit.rs: Fix offset when skipping insns for handling divisions by 0

When dividing by the content of a register, we have a specific case to
handle divisions by zero. When the denominator is not 0, we skip the
related instructions by emitted a near jump with the relevant offset.

However, this offset was incorrect. We accounted for two instructions: a
2-byte long XOR of the destination register with itself to set it to 0,
and a 5-byte long jump. The former may in fact emit 3 bytes under
certain condition, depending on the destination in use. For example,
when dividing from r7, we emit 3 bytes, and the fixed offset of 7 bytes
we used is incorrect, and triggers a segfault.

To fix this, let's add a check to determine the offset value we should
use. We also add related tests to make sure we don't regress on this in
the future.

Fixes: ce7d416d9a98 ("src/jit.rs: Update behaviour for divisions by zero (JIT)")
Fixes: https://github.com/qmonnet/rbpf/issues/88
Reported-by: Officeyutong
Signed-off-by: Quentin Monnet <quentin@isovalent.com>

src/jit.rs

@@ -148,6 +148,10 @@ impl JitCompiler {
         }
     }
 
+    fn basix_rex_would_set_bits(&self, w: u8, src: u8, dst: u8) -> bool {
+        w != 0 || (src & 0b1000) != 0 || (dst & 0b1000) != 0
+    }
+
     fn emit_rex(&self, mem: &mut JitMemory, w: u8, r: u8, x: u8, b: u8) {
         assert_eq!((w | 1), 1);
         assert_eq!((r | 1), 1);
@@ -159,7 +163,7 @@ impl JitCompiler {
     // Emits a REX prefix with the top bit of src and dst.
     // Skipped if no bits would be set.
     fn emit_basic_rex(&self, mem: &mut JitMemory, w: u8, src: u8, dst: u8) {
-        if w != 0 || (src & 0b1000) != 0 || (dst & 0b1000) != 0 {
+        if self.basix_rex_would_set_bits(w, src, dst) {
             let is_masked = | val, mask | { match val & mask {
                 0 => 0,
                 _ => 1
@@ -397,7 +401,12 @@ impl JitCompiler {
 
             if div {
                 // No division by 0: skip next instructions
-                self.emit_direct_jcc(mem, 0x85, 7);
+                // Jump offset: emit_alu32 adds 2 to 3 bytes, emit_jmp adds 5
+                let offset = match self.basix_rex_would_set_bits(0, dst, dst) {
+                    true => 3 + 5,
+                    false => 2 + 5,
+                };
+                self.emit_direct_jcc(mem, 0x85, offset);
                 // Division by 0: set dst to 0 then go to next instruction
                 // Set register to 0: xor with itself
                 self.emit_alu32(mem, 0x31, dst, dst);

tests/ubpf_jit_x86_64.rs

@@ -392,6 +392,33 @@ fn test_jit_div64_reg() {
     unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0x300000000); }
 }
 
+// For some register numbers, we don't emit the same instructions for handling divisions by zero,
+// which means we don't use the same offset to skip these instructions when the divisor is not
+// zero. We've had a regression because of this before, make sure we test it.
+#[test]
+fn test_jit_div32_highreg() {
+    let prog = assemble("
+        mov r0, 2
+        mov r7, 4
+        div32 r7, r0
+        exit").unwrap();
+    let mut vm = rbpf::EbpfVmNoData::new(Some(&prog)).unwrap();
+    vm.jit_compile().unwrap();
+    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0x2); }
+}
+
+#[test]
+fn test_jit_div64_highreg() {
+    let prog = assemble("
+        mov r0, 2
+        mov r7, 4
+        div r7, r0
+        exit").unwrap();
+    let mut vm = rbpf::EbpfVmNoData::new(Some(&prog)).unwrap();
+    vm.jit_compile().unwrap();
+    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0x2); }
+}
+
 #[test]
 fn test_jit_early_exit() {
     let prog = assemble("