feat: Use Range to indicate the memory range that is allowed to be accessed

Since some maps may store a lot of data, representing the address with a
single byte may waste a lot of space. Using range can simplify this
problem.

fix: fix the no_std error for HashSet

In no_std mode, BtreeMap is used instead of HashMap, but Range does not
implement the `Ord` trait. In order to use HashMap in no_std mode,
hashbrown crate is used here, which provides better performance than the
standard library.

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

Cargo.toml

@@ -41,6 +41,10 @@ cranelift-frontend = { version = "0.99", optional = true }
 cranelift-jit = { version = "0.99", optional = true }
 cranelift-native = { version = "0.99", optional = true }
 cranelift-module = { version = "0.99", optional = true }
+hashbrown = { version = "0.15", default-features = false, features = [
+    "alloc",
+    "default-hasher",
+] }
 
 [dev-dependencies]
 

README.md

@@ -210,7 +210,7 @@ useful for programs that should be compatible with the Linux kernel and
 therefore must use specific helper numbers.
 
 ```rust,ignore
-pub fn register_allowed_memory(&mut self,, addr: &[u64]) -> ()
+pub fn register_allowed_memory(&mut self,, addrs_range: Range<u64>) -> ()
 ```
 
 This function adds a list of memory addresses that the eBPF program is allowed

examples/allowed_memory.rs

@@ -2,7 +2,7 @@
 // Copyright 2024 Akenes SA <wouter.dullaert@exoscale.ch>
 
 extern crate elf;
-use std::{iter::FromIterator, ptr::addr_of};
+use std::{ptr::addr_of};
 
 extern crate rbpf;
 
@@ -40,8 +40,7 @@ fn main() {
         .unwrap();
 
     let start = addr_of!(MAP_VALUE) as u64;
-    let addrs = Vec::from_iter(start..start + size_of::<Value>() as u64);
-    vm.register_allowed_memory(&addrs);
+    vm.register_allowed_memory(start..start + size_of::<Value>() as u64);
 
     let res = vm.execute_program().unwrap();
     assert_eq!(res, 1);

src/interpreter.rs

@@ -9,6 +9,7 @@ use crate::ebpf;
 use crate::ebpf::MAX_CALL_DEPTH;
 use crate::lib::*;
 use crate::stack::{StackFrame, StackUsage};
+use core::ops::Range;
 
 #[allow(clippy::too_many_arguments)]
 fn check_mem(
@@ -19,7 +20,7 @@ fn check_mem(
     mbuff: &[u8],
     mem: &[u8],
     stack: &[u8],
-    allowed_memory: &HashSet<u64>,
+    allowed_memory: &HashSet<Range<u64>>,
 ) -> Result<(), Error> {
     if let Some(addr_end) = addr.checked_add(len as u64) {
         if mbuff.as_ptr() as u64 <= addr && addr_end <= mbuff.as_ptr() as u64 + mbuff.len() as u64 {
@@ -31,7 +32,7 @@ fn check_mem(
         if stack.as_ptr() as u64 <= addr && addr_end <= stack.as_ptr() as u64 + stack.len() as u64 {
             return Ok(());
         }
-        if allowed_memory.contains(&addr) {
+        if allowed_memory.iter().any(|range| range.contains(&addr)) {
             return Ok(());
         }
     }
@@ -51,7 +52,7 @@ pub fn execute_program(
     mem: &[u8],
     mbuff: &[u8],
     helpers: &HashMap<u32, ebpf::Helper>,
-    allowed_memory: &HashSet<u64>,
+    allowed_memory: &HashSet<Range<u64>>,
 ) -> Result<u64, Error> {
     const U32MAX: u64 = u32::MAX as u64;
     const SHIFT_MASK_64: u64 = 0x3f;

src/lib.rs

@@ -35,6 +35,7 @@ extern crate cranelift_native;
 
 use crate::lib::*;
 use byteorder::{ByteOrder, LittleEndian};
+use core::ops::Range;
 use stack::{StackUsage, StackVerifier};
 
 mod asm_parser;
@@ -70,6 +71,7 @@ pub mod lib {
     pub use self::core::mem;
     pub use self::core::mem::ManuallyDrop;
     pub use self::core::ptr;
+    pub use hashbrown::{HashMap,HashSet};
 
     #[cfg(feature = "std")]
     pub use std::println;
@@ -97,9 +99,9 @@ pub mod lib {
     // BTree-based implementations of Maps and Sets. The cranelift module uses
     // BTrees by default, hence we need to expose it twice here.
     #[cfg(not(feature = "std"))]
-    pub use alloc::collections::{BTreeMap as HashMap, BTreeMap, BTreeSet as HashSet, BTreeSet};
+    pub use alloc::collections::{BTreeMap};
     #[cfg(feature = "std")]
-    pub use std::collections::{BTreeMap, HashMap, HashSet};
+    pub use std::collections::{BTreeMap};
 
     /// In no_std we use a custom implementation of the error which acts as a
     /// replacement for the io Error.
@@ -180,7 +182,7 @@ pub struct EbpfVmMbuff<'a> {
     #[cfg(feature = "cranelift")]
     cranelift_prog: Option<cranelift::CraneliftProgram>,
     helpers: HashMap<u32, ebpf::Helper>,
-    allowed_memory: HashSet<u64>,
+    allowed_memory: HashSet<Range<u64>>,
     stack_usage: Option<StackUsage>,
     stack_verifier: StackVerifier,
 }
@@ -386,7 +388,6 @@ impl<'a> EbpfVmMbuff<'a> {
     /// # Examples
     ///
     /// ```
-    /// use std::iter::FromIterator;
     /// use std::ptr::addr_of;
     ///
     /// struct MapValue {
@@ -402,13 +403,10 @@ impl<'a> EbpfVmMbuff<'a> {
     /// // Instantiate a VM.
     /// let mut vm = rbpf::EbpfVmMbuff::new(Some(prog)).unwrap();
     /// let start = addr_of!(VALUE) as u64;
-    /// let addrs = Vec::from_iter(start..start+size_of::<MapValue>() as u64);
-    /// vm.register_allowed_memory(&addrs);
+    /// vm.register_allowed_memory(start..start+size_of::<MapValue>() as u64);
     /// ```
-    pub fn register_allowed_memory(&mut self, addrs: &[u64]) {
-        for i in addrs {
-            self.allowed_memory.insert(*i);
-        }
+    pub fn register_allowed_memory(&mut self, addrs_range: Range<u64>) {
+        self.allowed_memory.insert(addrs_range);
     }
 
     /// Execute the program loaded, with the given packet data and metadata buffer.
@@ -998,7 +996,6 @@ impl<'a> EbpfVmFixedMbuff<'a> {
     /// # Examples
     ///
     /// ```
-    /// use std::iter::FromIterator;
     /// use std::ptr::addr_of;
     ///
     /// struct MapValue {
@@ -1014,11 +1011,10 @@ impl<'a> EbpfVmFixedMbuff<'a> {
     /// // Instantiate a VM.
     /// let mut vm = rbpf::EbpfVmFixedMbuff::new(Some(prog), 0x40, 0x50).unwrap();
     /// let start = addr_of!(VALUE) as u64;
-    /// let addrs = Vec::from_iter(start..start+size_of::<MapValue>() as u64);
-    /// vm.register_allowed_memory(&addrs);
+    /// vm.register_allowed_memory(start..start+size_of::<MapValue>() as u64);
     /// ```
-    pub fn register_allowed_memory(&mut self, allowed: &[u64]) {
-        self.parent.register_allowed_memory(allowed)
+    pub fn register_allowed_memory(&mut self, addrs_range: Range<u64>) {
+        self.parent.register_allowed_memory(addrs_range)
     }
 
     /// Execute the program loaded, with the given packet data.
@@ -1528,7 +1524,6 @@ impl<'a> EbpfVmRaw<'a> {
     /// # Examples
     ///
     /// ```
-    /// use std::iter::FromIterator;
     /// use std::ptr::addr_of;
     ///
     /// struct MapValue {
@@ -1544,11 +1539,10 @@ impl<'a> EbpfVmRaw<'a> {
     /// // Instantiate a VM.
     /// let mut vm = rbpf::EbpfVmRaw::new(Some(prog)).unwrap();
     /// let start = addr_of!(VALUE) as u64;
-    /// let addrs = Vec::from_iter(start..start+size_of::<MapValue>() as u64);
-    /// vm.register_allowed_memory(&addrs);
+    /// vm.register_allowed_memory(start..start+size_of::<MapValue>() as u64);
     /// ```
-    pub fn register_allowed_memory(&mut self, allowed: &[u64]) {
-        self.parent.register_allowed_memory(allowed)
+    pub fn register_allowed_memory(&mut self, addrs_range: Range<u64>) {
+        self.parent.register_allowed_memory(addrs_range)
     }
 
     /// Execute the program loaded, with the given packet data.
@@ -1966,7 +1960,6 @@ impl<'a> EbpfVmNoData<'a> {
     /// # Examples
     ///
     /// ```
-    /// use std::iter::FromIterator;
     /// use std::ptr::addr_of;
     ///
     /// struct MapValue {
@@ -1982,11 +1975,10 @@ impl<'a> EbpfVmNoData<'a> {
     /// // Instantiate a VM.
     /// let mut vm = rbpf::EbpfVmNoData::new(Some(prog)).unwrap();
     /// let start = addr_of!(VALUE) as u64;
-    /// let addrs = Vec::from_iter(start..start+size_of::<MapValue>() as u64);
-    /// vm.register_allowed_memory(&addrs);
+    /// vm.register_allowed_memory(start..start+size_of::<MapValue>() as u64);
     /// ```
-    pub fn register_allowed_memory(&mut self, allowed: &[u64]) {
-        self.parent.register_allowed_memory(allowed)
+    pub fn register_allowed_memory(&mut self, addrs_range: Range<u64>) {
+        self.parent.register_allowed_memory(addrs_range)
     }
 
     /// JIT-compile the loaded program. No argument required for this.