feat: Add ebpf to ebpf call support

Like ubpf, we can check the correctness of ebpf calls and add support for
ebpf to ebpf calls.
To do this, a `StackFrame` data structure is added to save the context
when making function calls. Currently, the maximum nesting level of
function calls is 8.
In addition, function calls require the stack to be protected.
Therefore, when entering a new function, the stack pointer will be
moved. Since it is impossible to accurately calculate the stack size
used by the current function, the current stack usage calculator gives a
default value: 256 bytes. Users can customize `StackUsageCalculator` to
override the current default value. For a simple example: when executing
multiple simple function calls, since these functions usually do not use
the stack, in theory the stack space will not be exhausted, the user can
return 0 in the customized `StackUsageCalculator`, so that the execution
of the ebpf program will not cause a stack overflow.

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

src/ebpf.rs

@@ -17,6 +17,10 @@
 use byteorder::{ByteOrder, LittleEndian};
 use crate::lib::*;
 
+/// The default stack size for the eBPF program if there is some bpf to bpf calls.
+pub const RBPF_EBPF_LOCAL_FUNCTION_STACK_SIZE: u16 = 256;
+/// Maximum number of bpf to bpf call depth.
+pub const RBPF_MAX_CALL_DEPTH: usize = 8;
 /// Maximum number of instructions in an eBPF program.
 pub const PROG_MAX_INSNS: usize = 1000000;
 /// Size of an eBPF instructions, in bytes.

src/interpreter.rs

@@ -6,7 +6,9 @@
 //      (Translation to Rust, MetaBuff/multiple classes addition, hashmaps for helpers)
 
 use crate::ebpf;
+use crate::ebpf::RBPF_MAX_CALL_DEPTH;
 use crate::lib::*;
+use crate::stack::{StackFrame, StackUsage};
 
 #[allow(clippy::too_many_arguments)]
 fn check_mem(
@@ -45,6 +47,7 @@ fn check_mem(
 
 pub fn execute_program(
     prog_: Option<&[u8]>,
+    stack_usage: Option<&StackUsage>,
     mem: &[u8],
     mbuff: &[u8],
     helpers: &HashMap<u32, ebpf::Helper>,
@@ -53,13 +56,14 @@ pub fn execute_program(
     const U32MAX: u64 = u32::MAX as u64;
     const SHIFT_MASK_64: u64 = 0x3f;
 
-    let prog = match prog_ {
-        Some(prog) => prog,
+    let (prog,stack_usage) = match prog_ {
+        Some(prog) => (prog, stack_usage.unwrap()),
         None => Err(Error::new(ErrorKind::Other,
                     "Error: No program set, call prog_set() to load one"))?,
     };
     let stack = vec![0u8;ebpf::STACK_SIZE];
-
+    let mut stacks = [StackFrame::new();RBPF_MAX_CALL_DEPTH];
+    let mut stack_frame_idx = 0;
     // R1 points to beginning of memory area, R10 to stack
     let mut reg: [u64;11] = [
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, stack.as_ptr() as u64 + stack.len() as u64
@@ -82,6 +86,11 @@ pub fn execute_program(
     let mut insn_ptr:usize = 0;
     while insn_ptr * ebpf::INSN_SIZE < prog.len() {
         let insn = ebpf::get_insn(prog, insn_ptr);
+        if stack_frame_idx < RBPF_MAX_CALL_DEPTH{
+            if let Some(usage) = stack_usage.stack_usage_for_local_func(insn_ptr) {
+                stacks[stack_frame_idx].set_stack_usage(usage);
+            }
+        }
         insn_ptr += 1;
         let _dst = insn.dst as usize;
         let _src = insn.src as usize;
@@ -356,13 +365,43 @@ pub fn execute_program(
 
             // Do not delegate the check to the verifier, since registered functions can be
             // changed after the program has been verified.
-            ebpf::CALL       => if let Some(function) = helpers.get(&(insn.imm as u32)) {
-                reg[0] = function(reg[1], reg[2], reg[3], reg[4], reg[5]);
-            } else {
-                Err(Error::new(ErrorKind::Other, format!("Error: unknown helper function (id: {:#x})", insn.imm as u32)))?;
-            },
+            ebpf::CALL       => {
+                match _src {
+                    // Call helper function
+                    0 => {
+                        if let Some(function) = helpers.get(&(insn.imm as u32)) {
+                            reg[0] = function(reg[1], reg[2], reg[3], reg[4], reg[5]);
+                        } else {
+                            Err(Error::new(ErrorKind::Other, format!("Error: unknown helper function (id: {:#x})", insn.imm as u32)))?;
+                        }
+                    }
+                    // BPF To BPF call
+                    1 => {
+                        if stack_frame_idx >= RBPF_MAX_CALL_DEPTH {
+                            Err(Error::new(ErrorKind::Other, format!("Error: too many nested calls (max: {RBPF_MAX_CALL_DEPTH})")))?;
+                        }
+                        stacks[stack_frame_idx].save_registers(&reg[6..=9]);
+                        stacks[stack_frame_idx].save_return_address(insn_ptr);
+                        reg[10] -= stacks[stack_frame_idx].get_stack_usage().stack_usage() as u64;
+                        stack_frame_idx += 1;
+                        insn_ptr += insn.imm as usize;
+                    }
+                    _ => {
+                        Err(Error::new(ErrorKind::Other, format!("Error: invalid call to function #{:?} (insn #{insn_ptr:?})", insn.imm)))?;
+                    }
+                }
+            }
             ebpf::TAIL_CALL  => unimplemented!(),
-            ebpf::EXIT       => return Ok(reg[0]),
+            ebpf::EXIT       => {
+                if stack_frame_idx > 0 {
+                    stack_frame_idx -= 1;
+                    reg[6..=9].copy_from_slice(&stacks[stack_frame_idx].get_registers());
+                    insn_ptr = stacks[stack_frame_idx].get_return_address();
+                    reg[10] += stacks[stack_frame_idx].get_stack_usage().stack_usage() as u64;
+                } else {
+                    return Ok(reg[0]);
+                }
+            }
 
             _                => unreachable!()
         }

src/lib.rs

@@ -35,6 +35,7 @@ extern crate cranelift_native;
 
 use crate::lib::*;
 use byteorder::{ByteOrder, LittleEndian};
+use stack::{StackUsage, StackVerifier};
 
 mod asm_parser;
 pub mod assembler;
@@ -45,6 +46,7 @@ pub mod ebpf;
 pub mod helpers;
 pub mod insn_builder;
 mod interpreter;
+mod stack;
 #[cfg(all(not(windows), feature = "std"))]
 mod jit;
 #[cfg(not(feature = "std"))]
@@ -66,7 +68,7 @@ pub mod lib {
     pub use self::core::mem;
     pub use self::core::mem::ManuallyDrop;
     pub use self::core::ptr;
-
+    pub use core::any::Any;
     pub use self::core::f64;
 
     #[cfg(feature = "std")]
@@ -116,6 +118,9 @@ pub type Verifier = fn(prog: &[u8]) -> Result<(), Error>;
 /// eBPF helper function.
 pub type Helper = fn(u64, u64, u64, u64, u64) -> u64;
 
+/// eBPF stack usage calculator function.
+pub type StackUsageCalculator = fn(prog:&[u8], pc:usize, data:&mut dyn Any) -> u16;
+
 // A metadata buffer with two offset indications. It can be used in one kind of eBPF VM to simulate
 // the use of a metadata buffer each time the program is executed, without the user having to
 // actually handle it. The offsets are used to tell the VM where in the buffer the pointers to
@@ -167,6 +172,8 @@ pub struct EbpfVmMbuff<'a> {
     cranelift_prog: Option<cranelift::CraneliftProgram>,
     helpers: HashMap<u32, ebpf::Helper>,
     allowed_memory: HashSet<u64>,
+    stack_usage: Option<StackUsage>,
+    stack_verifier: StackVerifier,
 }
 
 impl<'a> EbpfVmMbuff<'a> {
@@ -186,9 +193,13 @@ impl<'a> EbpfVmMbuff<'a> {
     /// let mut vm = rbpf::EbpfVmMbuff::new(Some(prog)).unwrap();
     /// ```
     pub fn new(prog: Option<&'a [u8]>) -> Result<EbpfVmMbuff<'a>, Error> {
-        if let Some(prog) = prog {
+        let mut stack_verifier = StackVerifier::new(None, None);
+        let stack_usage = if let Some(prog) = prog {
             verifier::check(prog)?;
-        }
+            Some(stack_verifier.stack_validate(prog)?)
+        } else {
+            None
+        };
 
         Ok(EbpfVmMbuff {
             prog,
@@ -199,6 +210,8 @@ impl<'a> EbpfVmMbuff<'a> {
             cranelift_prog: None,
             helpers: HashMap::new(),
             allowed_memory: HashSet::new(),
+            stack_usage,
+            stack_verifier
         })
     }
 
@@ -223,7 +236,9 @@ impl<'a> EbpfVmMbuff<'a> {
     /// ```
     pub fn set_program(&mut self, prog: &'a [u8]) -> Result<(), Error> {
         (self.verifier)(prog)?;
+        let stack_usage = self.stack_verifier.stack_validate(prog)?;
         self.prog = Some(prog);
+        self.stack_usage = Some(stack_usage);
         Ok(())
     }
 
@@ -265,6 +280,47 @@ impl<'a> EbpfVmMbuff<'a> {
         Ok(())
     }
 
+    /// Set a new stack usage calculator function. The function should return the stack usage
+    /// of the program in bytes. If a program has been loaded to the VM already, the calculator
+    /// is immediately run.
+    /// 
+    /// # Examples
+    /// 
+    /// ```
+    /// use rbpf::lib::{Error, ErrorKind};
+    /// use rbpf::ebpf;
+    /// use core::any::Any;
+    /// // Define a simple stack usage calculator function.
+    /// fn calculator(prog: &[u8], pc: usize, data: &mut dyn Any) -> u16 {
+    ///    // This is a dummy implementation, just for the example.
+    ///    // In a real implementation, you would calculate the stack usage based on the program.
+    ///    // Here we just return a fixed value.
+    ///    16
+    /// }
+    /// 
+    /// let prog1 = &[
+    ///     0xb7, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // mov r0, 0
+    ///     0x95, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00  // exit
+    /// ];
+    /// 
+    /// // Instantiate a VM.
+    /// let mut vm = rbpf::EbpfVmMbuff::new(Some(prog1)).unwrap();
+    /// // Change the stack usage calculator.
+    /// vm.set_stack_usage_calculator(calculator, Box::new(())).unwrap();
+    /// ```
+    pub fn set_stack_usage_calculator(
+        &mut self,
+        calculator: StackUsageCalculator,
+        data: Box<dyn Any>,
+    ) -> Result<(), Error> {
+        let mut stack_verifier = StackVerifier::new(Some(calculator), Some(data));
+        if let Some(prog) = self.prog {
+            self.stack_usage = Some(stack_verifier.stack_validate(prog)?);
+        }
+        self.stack_verifier = stack_verifier;
+        Ok(())
+    }
+
     /// Register a built-in or user-defined helper function in order to use it later from within
     /// the eBPF program. The helper is registered into a hashmap, so the `key` can be any `u32`.
     ///
@@ -383,7 +439,8 @@ impl<'a> EbpfVmMbuff<'a> {
     /// assert_eq!(res, 0x2211);
     /// ```
     pub fn execute_program(&self, mem: &[u8], mbuff: &[u8]) -> Result<u64, Error> {
-        interpreter::execute_program(self.prog, mem, mbuff, &self.helpers, &self.allowed_memory)
+        let stack_usage = self.stack_usage.as_ref();
+        interpreter::execute_program(self.prog,stack_usage, mem, mbuff, &self.helpers, &self.allowed_memory)
     }
 
     /// JIT-compile the loaded program. No argument required for this.

src/stack.rs

@@ -0,0 +1,140 @@
+use core::any::Any;
+
+use crate::{
+    ebpf::{self, RBPF_EBPF_LOCAL_FUNCTION_STACK_SIZE},
+    lib::*,
+    StackUsageCalculator,
+};
+
+#[derive(Debug, Copy, Clone)]
+pub struct StackFrame {
+    return_address: usize,
+    saved_registers: [u64; 4],
+    stack_usage: StackUsageType,
+}
+
+impl StackFrame {
+    /// Create a new stack frame
+    pub const fn new() -> Self {
+        Self {
+            return_address: 0,
+            saved_registers: [0; 4],
+            stack_usage: StackUsageType::Default,
+        }
+    }
+    /// Save the callee-saved registers
+    pub fn save_registers(&mut self, regs: &[u64]) {
+        self.saved_registers.copy_from_slice(regs);
+    }
+
+    /// Get the callee-saved registers
+    pub fn get_registers(&self) -> [u64; 4] {
+        self.saved_registers
+    }
+
+    /// Save the return address
+    pub fn save_return_address(&mut self, address: usize) {
+        self.return_address = address;
+    }
+
+    /// Get the return address
+    pub fn get_return_address(&self) -> usize {
+        self.return_address
+    }
+
+    /// Set the stack usage
+    pub fn set_stack_usage(&mut self, usage: StackUsageType) {
+        self.stack_usage = usage;
+    }
+    /// Get the stack usage
+    pub fn get_stack_usage(&self) -> StackUsageType {
+        self.stack_usage
+    }
+}
+
+#[derive(Debug, Copy, Clone)]
+pub enum StackUsageType {
+    Default,
+    Custom(u16),
+}
+
+impl StackUsageType {
+    pub fn stack_usage(&self) -> u16 {
+        match self {
+            StackUsageType::Default => RBPF_EBPF_LOCAL_FUNCTION_STACK_SIZE,
+            StackUsageType::Custom(size) => *size,
+        }
+    }
+}
+
+pub struct StackVerifier {
+    calculator: Option<StackUsageCalculator>,
+    data: Option<Box<dyn Any>>,
+}
+
+impl StackVerifier {
+    pub fn new(
+        stack_usage_calculator: Option<StackUsageCalculator>,
+        data: Option<Box<dyn Any>>,
+    ) -> Self {
+        Self {
+            calculator: stack_usage_calculator,
+            data,
+        }
+    }
+    /// Validate the stack usage of a program
+    ///
+    /// This function checks the stack usage of a program and returns a `StackUsage` object
+    /// containing the stack usage for each local function in the program.
+    ///
+    /// # Returns
+    /// - `Ok(StackUsage)` if the stack usage is valid
+    /// - `Err(Error)` if the stack usage is invalid
+    pub fn stack_validate(&mut self, prog: &[u8]) -> Result<StackUsage, Error> {
+        let mut stack_usage = HashMap::new();
+        let ty = self.calculate_stack_usage_for_local_func(prog, 0)?;
+        stack_usage.insert(0, ty);
+        for idx in 0..prog.len() / ebpf::INSN_SIZE {
+            let insn = ebpf::get_insn(prog, idx);
+            if insn.opc == ebpf::CALL {
+                let dst_insn_ptr = idx as isize + 1 + insn.imm as isize;
+                let ty = self.calculate_stack_usage_for_local_func(prog, dst_insn_ptr as usize)?;
+                stack_usage.insert(dst_insn_ptr as usize, ty);
+            }
+        }
+        Ok(StackUsage(stack_usage))
+    }
+
+    /// Calculate the stack usage for a local function
+    fn calculate_stack_usage_for_local_func(
+        &mut self,
+        prog: &[u8],
+        pc: usize,
+    ) -> Result<StackUsageType, Error> {
+        let mut ty = StackUsageType::Default;
+        if let Some(calculator) = self.calculator {
+            let size = calculator(prog, pc, self.data.as_mut().unwrap());
+            ty = StackUsageType::Custom(size);
+        }
+        if ty.stack_usage() % 16 > 0 {
+            Err(Error::new(
+                ErrorKind::Other,
+                format!(
+                    "local function (at PC {}) has improperly sized stack use ({})",
+                    pc,
+                    ty.stack_usage()
+                ),
+            ))?;
+        }
+        Ok(ty)
+    }
+}
+
+pub struct StackUsage(HashMap<usize, StackUsageType>);
+
+impl StackUsage {
+    /// Get the stack usage for a local function
+    pub fn stack_usage_for_local_func(&self, pc: usize) -> Option<StackUsageType> {
+        self.0.get(&pc).cloned()
+    }
+}

src/verifier.rs

@@ -248,7 +248,21 @@ pub fn check(prog: &[u8]) -> Result<(), Error> {
             ebpf::JSLE_IMM32 => { check_jmp_offset(prog, insn_ptr)?; },
             ebpf::JSLE_REG32 => { check_jmp_offset(prog, insn_ptr)?; },
 
-            ebpf::CALL       => {},
+            ebpf::CALL       => {
+                let src = insn.src;
+                match src {
+                    0 => {
+                        if insn.imm < 0 { reject(format!("invalid call to function #{:?} (insn #{insn_ptr:?})", insn.imm))?; }
+                    }
+                    1 => {
+                        let dst_insn_ptr = insn_ptr as isize + 1 + insn.imm as isize;
+                        if dst_insn_ptr < 0 || dst_insn_ptr as usize >= (prog.len() / ebpf::INSN_SIZE) {
+                            reject(format!("call out of code to #{dst_insn_ptr:?} (insn #{insn_ptr:?})"))?;
+                        }
+                    }
+                    _ => { reject(format!("invalid call to function #{:?} (insn #{insn_ptr:?})", insn.imm))?; }
+                }
+            },
             ebpf::TAIL_CALL  => { unimplemented!() },
             ebpf::EXIT       => {},
 

tests/ubpf_vm.rs

@@ -19,7 +19,8 @@
 extern crate rbpf;
 mod common;
 
-use rbpf::helpers;
+use rbpf::insn_builder::{Arch, Instruction, IntoBytes, Source};
+use rbpf::{helpers, insn_builder::BpfCode};
 use rbpf::assembler::assemble;
 use common::{TCP_SACK_ASM, TCP_SACK_MATCH, TCP_SACK_NOMATCH};
 
@@ -2227,6 +2228,19 @@ fn test_vm_stxw() {
     assert_eq!(vm.execute_program(mem).unwrap(), 0x44332211);
 }
 
+#[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);
+}
+
 #[test]
 fn test_vm_subnet() {
     let prog = assemble("