feat: add func call for jit

Signed-off-by: Godones <chenlinfeng25@outlook.com>

src/jit.rs

@@ -461,6 +461,20 @@ impl JitCompiler {
         }
     }
 
+    fn emit_local_call(&mut self, mem: &mut JitMemory, target_pc: isize) {
+        self.emit_push(mem, map_register(6));
+        self.emit_push(mem, map_register(7));
+        self.emit_push(mem, map_register(8));
+        self.emit_push(mem, map_register(9));
+        // e8 is the opcode for a CALL
+        self.emit1(mem, 0xe8); 
+        self.emit_jump_offset(mem, target_pc);
+        self.emit_pop(mem, map_register(9));
+        self.emit_pop(mem, map_register(8));
+        self.emit_pop(mem, map_register(7));
+        self.emit_pop(mem, map_register(6));
+    }
+
     fn jit_compile(&mut self, mem: &mut JitMemory, prog: &[u8], use_mbuff: bool, update_data_ptr: bool,
                    helpers: &HashMap<u32, ebpf::Helper>) -> Result<(), Error> {
         self.emit_push(mem, RBP);
@@ -516,6 +530,10 @@ impl JitCompiler {
         // Allocate stack space
         self.emit_alu64_imm32(mem, 0x81, 5, RSP, ebpf::STACK_SIZE as i32);
 
+        self.emit1(mem, 0xe8);
+        self.emit4(mem, 5);
+        self.emit_jmp(mem, TARGET_PC_EXIT);
+
         self.pc_locs = vec![0; prog.len() / ebpf::INSN_SIZE + 1];
 
         let mut insn_ptr:usize = 0;
@@ -870,24 +888,33 @@ impl JitCompiler {
                 },
 
                 ebpf::CALL       => {
-                    // For JIT, helpers in use MUST be registered at compile time. They can be
-                    // updated later, but not created after compiling (we need the address of the
-                    // helper function in the JIT-compiled program).
-                    if let Some(helper) = helpers.get(&(insn.imm as u32)) {
-                        // We reserve RCX for shifts
-                        self.emit_mov(mem, R9, RCX);
-                        self.emit_call(mem, *helper as usize);
-                    } else {
-                        Err(Error::new(ErrorKind::Other,
-                                       format!("[JIT] Error: unknown helper function (id: {:#x})",
-                                               insn.imm as u32)))?;
-                    };
+                    match insn.src {
+                        0x0 => {
+                            // For JIT, helpers in use MUST be registered at compile time. They can be
+                            // updated later, but not created after compiling (we need the address of the
+                            // helper function in the JIT-compiled program).
+                            if let Some(helper) = helpers.get(&(insn.imm as u32)) {
+                                // We reserve RCX for shifts
+                                self.emit_mov(mem, R9, RCX);
+                                self.emit_call(mem, *helper as usize);
+                            } else {
+                                Err(Error::new(ErrorKind::Other,
+                                            format!("[JIT] Error: unknown helper function (id: {:#x})",
+                                                    insn.imm as u32)))?;
+                            };
+                        }
+                        0x1 => {
+                            let target_pc = insn_ptr as isize + insn.imm as isize + 1;
+                            self.emit_local_call(mem, target_pc);
+                        }   
+                        _ => {
+                            Err(Error::new(ErrorKind::Other, format!("Error: unexpected call type {:#x}",src)))?;
+                        }
+                    }
                 },
                 ebpf::TAIL_CALL  => { unimplemented!() },
                 ebpf::EXIT       => {
-                    if insn_ptr != prog.len() / ebpf::INSN_SIZE - 1 {
-                        self.emit_jmp(mem, TARGET_PC_EXIT);
-                    };
+                    self.emit1(mem, 0xc3); // ret
                 },
 
                 _                => {

tests/ubpf_jit_x86_64.rs

@@ -21,7 +21,8 @@
 extern crate rbpf;
 mod common;
 
-use rbpf::helpers;
+use rbpf::ebpf::CALL;
+use rbpf::{ebpf::to_insn_vec, helpers};
 use rbpf::assembler::assemble;
 use common::{TCP_SACK_ASM, TCP_SACK_MATCH, TCP_SACK_NOMATCH};
 
@@ -2356,3 +2357,37 @@ fn test_jit_tcp_sack_nomatch() {
     vm.jit_compile().unwrap();
     unsafe { assert_eq!(vm.execute_program_jit(mem.as_mut_slice()).unwrap(), 0x0); }
 }
+
+#[test]
+fn test_bpf_to_bpf_call(){
+    let test_code = assemble(
+        "
+    mov64 r1, 0x10
+    mov64 r2, 0x1
+    call 0x4
+    mov64 r1, 0x1
+    mov64 r2, r0
+    call 0x4
+    exit
+    mov64 r0, r1
+    sub64 r0, r2
+    exit
+    mov64 r0, r2
+    add64 r0, r1
+    exit
+    ",
+    )
+    .unwrap();
+    let mut code = to_insn_vec(&test_code);
+    let mut real_code = Vec::new();
+    code.iter_mut().for_each(|insn| {
+        if insn.opc == CALL {
+            insn.src = 0x1;
+        }
+        real_code.extend_from_slice(&insn.to_array());
+    });
+    let mut vm = rbpf::EbpfVmNoData::new(Some(&real_code)).unwrap();
+    vm.jit_compile().unwrap();
+    let vm_res= unsafe { vm.execute_program_jit().unwrap() };    
+    assert_eq!(vm_res, 0x10);
+}

tests/ubpf_vm.rs

@@ -19,8 +19,8 @@
 extern crate rbpf;
 mod common;
 
-use rbpf::insn_builder::{Arch, Instruction, IntoBytes, Source};
-use rbpf::{helpers, insn_builder::BpfCode};
+use rbpf::ebpf::{to_insn_vec, CALL};
+use rbpf::helpers;
 use rbpf::assembler::assemble;
 use common::{TCP_SACK_ASM, TCP_SACK_MATCH, TCP_SACK_NOMATCH};
 
@@ -2230,15 +2230,35 @@ fn test_vm_stxw() {
 
 #[test]
 fn test_bpf_to_bpf_call(){
-    let mut program = BpfCode::new();
-    program.mov(Source::Imm, Arch::X64).set_dst(0x1).set_imm(0xff).push()
-    .call().set_src(1).set_imm(1).push()
-    .mov(Source::Imm, Arch::X64).set_dst(0x2).set_imm(0x1).push()
-    .mov(Source::Imm, Arch::X64).set_dst(0).set_imm(0xf).push()
-    .exit().push();
-    let prog = program.into_bytes();
-    let vm = rbpf::EbpfVmNoData::new(Some(&prog)).unwrap();
-    assert_eq!(vm.execute_program().unwrap(), 0xf);
+    let test_code = assemble(
+        "
+    mov64 r1, 0x10
+    mov64 r2, 0x1
+    call 0x4
+    mov64 r1, 0x1
+    mov64 r2, r0
+    call 0x4
+    exit
+    mov64 r0, r1
+    sub64 r0, r2
+    exit
+    mov64 r0, r2
+    add64 r0, r1
+    exit
+    ",
+    )
+    .unwrap();
+    let mut code = to_insn_vec(&test_code);
+    let mut real_code = Vec::new();
+    code.iter_mut().for_each(|insn| {
+        if insn.opc == CALL {
+            insn.src = 0x1;
+        }
+        real_code.extend_from_slice(&insn.to_array());
+    });
+    let vm = rbpf::EbpfVmNoData::new(Some(&real_code)).unwrap();
+    let vm_res= vm.execute_program().unwrap();    
+    assert_eq!(vm_res, 0x10);
 }
 
 #[test]