src/jit.rs: Update behaviour for divisions by zero (JIT)

Following the behaviour of Linux, uBPF, and the (in-progress)
specification for the eBPF instruction set, rework the divisions and
modulo remainder when the divisor is zero:

- Dividing by 0 sets the destination register to 0.
- When taking the remainder for the modulo of a division by 0, the
  destination register is unchanged.

See kernel commit f6b1b3bf0d5f ("bpf: fix subprog verifier bypass by
div/mod by 0 exception") for context on this behaviour.

This commit only updates the x86-64 JIT-compiler.

Signed-off-by: Quentin Monnet <quentin@isovalent.com>

src/jit.rs

@@ -253,6 +253,13 @@ fn set_anchor(jit: &mut JitMemory, target: isize) {
     jit.special_targets.insert(target, jit.offset);
 }
 
+#[inline]
+fn emit_direct_jcc(jit: &mut JitMemory, code: u8, offset: u32) {
+    emit1(jit, 0x0f);
+    emit1(jit, code);
+    emit_bytes!(jit, offset, u32);
+}
+
 // Load [src + offset] into dst
 #[inline]
 fn emit_load(jit: &mut JitMemory, size: OperandSize, src: u8, dst: u8, offset: i32) {
@@ -360,8 +367,19 @@ fn muldivmod(jit: &mut JitMemory, pc: u16, opc: u8, src: u8, dst: u8, imm: i32)
     let div = (opc & ebpf::BPF_ALU_OP_MASK) == (ebpf::DIV32_IMM & ebpf::BPF_ALU_OP_MASK);
     let modrm = (opc & ebpf::BPF_ALU_OP_MASK) == (ebpf::MOD32_IMM & ebpf::BPF_ALU_OP_MASK);
     let is64 = (opc & ebpf::BPF_CLS_MASK) == ebpf::BPF_ALU64;
+    let is_reg = (opc & ebpf::BPF_X) == ebpf::BPF_X;
 
-    if div || modrm {
+    if div && !is_reg && imm == 0 {
+        // Division by zero returns 0
+        // Set register to 0: xor with itself
+        emit_alu32(jit, 0x31, dst, dst);
+        return;
+    }
+    if modrm && !is_reg && imm == 0 {
+        // Modulo remainder of division by zero keeps destination register unchanged
+        return;
+    }
+    if (div || modrm) && is_reg {
         emit_load_imm(jit, RCX, pc as i64);
 
         // test src,src
@@ -371,8 +389,18 @@ fn muldivmod(jit: &mut JitMemory, pc: u16, opc: u8, src: u8, dst: u8, imm: i32)
             emit_alu32(jit, 0x85, src, src);
         }
 
-        // jz div_by_zero
-        emit_jcc(jit, 0x84, TARGET_PC_DIV_BY_ZERO);
+        if div {
+            // No division by 0: skip next instructions
+            emit_direct_jcc(jit, 0x85, 7);
+            // Division by 0: set dst to 0 then go to next instruction
+            // Set register to 0: xor with itself
+            emit_alu32(jit, 0x31, dst, dst);
+            emit_jmp(jit, (pc + 1) as isize);
+        }
+        if modrm {
+            // Modulo by zero: keep destination register unchanged
+            emit_jcc(jit, 0x84, (pc + 1) as isize);
+        }
     }
 
     if dst != RAX {
@@ -390,7 +418,7 @@ fn muldivmod(jit: &mut JitMemory, pc: u16, opc: u8, src: u8, dst: u8, imm: i32)
     emit_mov(jit, dst, RAX);
 
     if div || modrm {
-        // xor %edx,%edx
+        // Set register to 0: xor %edx,%edx
         emit_alu32(jit, 0x31, RDX, RDX);
     }
 
@@ -651,7 +679,7 @@ impl<'a> JitMemory<'a> {
                 ebpf::SUB64_REG  => emit_alu64(self, 0x29, src, dst),
                 ebpf::MUL64_IMM | ebpf::MUL64_REG |
                     ebpf::DIV64_IMM | ebpf::DIV64_REG |
-                    ebpf::MOD64_IMM | ebpf::MOD64_REG  =>
+                    ebpf::MOD64_IMM | ebpf::MOD64_REG =>
                     muldivmod(self, insn_ptr as u16, insn.opc, src, dst, insn.imm),
                 ebpf::OR64_IMM   => emit_alu64_imm32(self, 0x81, 1, dst, insn.imm),
                 ebpf::OR64_REG   => emit_alu64(self, 0x09, src, dst),
@@ -662,7 +690,7 @@ impl<'a> JitMemory<'a> {
                     emit_mov(self, src, RCX);
                     emit_alu64(self, 0xd3, 4, dst);
                 },
-                ebpf::RSH64_IMM  =>  emit_alu64_imm8(self, 0xc1, 5, dst, insn.imm as i8),
+                ebpf::RSH64_IMM  => emit_alu64_imm8(self, 0xc1, 5, dst, insn.imm as i8),
                 ebpf::RSH64_REG  => {
                     emit_mov(self, src, RCX);
                     emit_alu64(self, 0xd3, 5, dst);

tests/ubpf_jit_x86_64.rs

@@ -425,11 +425,52 @@ fn test_jit_err_call_unreg() {
     unsafe { vm.execute_program_jit().unwrap(); }
 }
 
-// TODO: Should panic!() instead, but I could not make it panic in JIT-compiled code, so the
-// program returns -1 instead. We can make it write on stderr, though.
 #[test]
-//#[should_panic(expected = "[JIT] Error: division by 0")]
-fn test_jit_err_div64_by_zero_reg() {
+fn test_jit_div64_by_zero_imm() {
+    let prog = assemble("
+        mov32 r0, 1
+        div r0, 0
+        exit").unwrap();
+    let mut vm = rbpf::EbpfVmNoData::new(Some(&prog)).unwrap();
+    vm.jit_compile().unwrap();
+    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0x0); }
+}
+
+#[test]
+fn test_jit_div_by_zero_imm() {
+    let prog = assemble("
+        mov32 r0, 1
+        div32 r0, 0
+        exit").unwrap();
+    let mut vm = rbpf::EbpfVmNoData::new(Some(&prog)).unwrap();
+    vm.jit_compile().unwrap();
+    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0x0); }
+}
+
+#[test]
+fn test_jit_mod64_by_zero_imm() {
+    let prog = assemble("
+        mov32 r0, 1
+        mod r0, 0
+        exit").unwrap();
+    let mut vm = rbpf::EbpfVmNoData::new(Some(&prog)).unwrap();
+    vm.jit_compile().unwrap();
+    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0x1); }
+}
+
+#[test]
+fn test_jit_mod_by_zero_imm() {
+    let prog = assemble("
+        mov32 r0, 1
+        mod32 r0, 0
+        exit").unwrap();
+    let mut vm = rbpf::EbpfVmNoData::new(Some(&prog)).unwrap();
+    vm.jit_compile().unwrap();
+    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0x1); }
+}
+
+#[test]
+fn test_jit_div64_by_zero_reg() {
     let prog = assemble("
         mov32 r0, 1
         mov32 r1, 0
@@ -437,13 +478,11 @@ fn test_jit_err_div64_by_zero_reg() {
         exit").unwrap();
     let mut vm = rbpf::EbpfVmNoData::new(Some(&prog)).unwrap();
     vm.jit_compile().unwrap();
-    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0xffffffffffffffff); }
+    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0x0); }
 }
 
-// TODO: Same remark as above
 #[test]
-//#[should_panic(expected = "[JIT] Error: division by 0")]
-fn test_jit_err_div_by_zero_reg() {
+fn test_jit_div_by_zero_reg() {
     let prog = assemble("
         mov32 r0, 1
         mov32 r1, 0
@@ -451,13 +490,11 @@ fn test_jit_err_div_by_zero_reg() {
         exit").unwrap();
     let mut vm = rbpf::EbpfVmNoData::new(Some(&prog)).unwrap();
     vm.jit_compile().unwrap();
-    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0xffffffffffffffff); }
+    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0x0); }
 }
 
-// TODO: Same remark as above
 #[test]
-//#[should_panic(expected = "[JIT] Error: division by 0")]
-fn test_jit_err_mod64_by_zero_reg() {
+fn test_jit_mod64_by_zero_reg() {
     let prog = assemble("
         mov32 r0, 1
         mov32 r1, 0
@@ -465,13 +502,11 @@ fn test_jit_err_mod64_by_zero_reg() {
         exit").unwrap();
     let mut vm = rbpf::EbpfVmNoData::new(Some(&prog)).unwrap();
     vm.jit_compile().unwrap();
-    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0xffffffffffffffff); }
+    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0x1); }
 }
 
-// TODO: Same remark as above
 #[test]
-//#[should_panic(expected = "[JIT] Error: division by 0")]
-fn test_jit_err_mod_by_zero_reg() {
+fn test_jit_mod_by_zero_reg() {
     let prog = assemble("
         mov32 r0, 1
         mov32 r1, 0
@@ -479,7 +514,7 @@ fn test_jit_err_mod_by_zero_reg() {
         exit").unwrap();
     let mut vm = rbpf::EbpfVmNoData::new(Some(&prog)).unwrap();
     vm.jit_compile().unwrap();
-    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0xffffffffffffffff); }
+    unsafe { assert_eq!(vm.execute_program_jit().unwrap(), 0x1); }
 }
 
 // TODO SKIP: JIT disabled for this testcase (stack oob check not implemented)