rbpf 0.4.1

Virtual machine and JIT compiler for eBPF programs
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

// SPDX-License-Identifier: (Apache-2.0 OR MIT)

use core::any::Any;

use crate::{
    ebpf::{self, 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 => 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;
        match self.calculator {
            Some(calculator) => {
                ty = StackUsageType::Custom(calculator(prog, pc, self.data.as_mut().unwrap()));
            }
            None => return Ok(ty),
        }
        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()
    }
}