feat: support JIT in no_std mode

In order to allow the use of jit in no_std mode, it is necessary to let
the user manually pass in the executable memory area.

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

Cargo.toml

@@ -28,7 +28,7 @@ include = [
 # Default features (std) are disabled so that the dependencies don't pull in the
 # standard library when the crate is compiled for no_std
 byteorder = { version = "1.2", default-features = false }
-log = {version = "0.4.21", default-features = false }
+log = { version = "0.4.21", default-features = false }
 combine = { version = "4.6", default-features = false }
 
 # Optional Dependencies when using the standard library

src/jit.rs

@@ -7,18 +7,12 @@
 
 #![allow(clippy::single_match)]
 
-use std::alloc;
-use std::collections::HashMap;
-use std::fmt::Error as FormatterError;
-use std::fmt::Formatter;
-use std::io::{Error, ErrorKind};
-use std::mem;
-use std::ops::{Index, IndexMut};
-use std::ptr;
-
-use crate::ebpf;
-
-extern crate libc;
+use crate::{ebpf, format, vec, Error, ErrorKind, HashMap, Vec};
+use core::fmt::Error as FormatterError;
+use core::fmt::Formatter;
+use core::mem;
+use core::ops::{Index, IndexMut};
+use core::ptr;
 
 type MachineCode = unsafe fn(*mut u8, usize, *mut u8, usize, usize, usize) -> u64;
 
@@ -1001,11 +995,11 @@ impl JitCompiler {
 
             // Assumes jump offset is at end of instruction
             unsafe {
-                let offset_loc = jump.offset_loc as i32 + std::mem::size_of::<i32>() as i32;
+                let offset_loc = jump.offset_loc as i32 + core::mem::size_of::<i32>() as i32;
                 let rel = &(target_loc as i32 - offset_loc) as *const i32;
 
                 let offset_ptr = mem.contents.as_ptr().add(jump.offset_loc) as *mut u8;
-                ptr::copy_nonoverlapping(rel.cast::<u8>(), offset_ptr, std::mem::size_of::<i32>());
+                ptr::copy_nonoverlapping(rel.cast::<u8>(), offset_ptr, core::mem::size_of::<i32>());
             }
         }
         Ok(())
@@ -1014,11 +1008,13 @@ impl JitCompiler {
 
 pub struct JitMemory<'a> {
     contents: &'a mut [u8],
-    layout: alloc::Layout,
+    #[cfg(feature = "std")]
+    layout: std::alloc::Layout,
     offset: usize,
 }
 
 impl<'a> JitMemory<'a> {
+    #[cfg(feature = "std")]
     pub fn new(
         prog: &[u8],
         helpers: &HashMap<u32, ebpf::Helper>,
@@ -1031,10 +1027,10 @@ impl<'a> JitMemory<'a> {
         let contents = unsafe {
             // Create a layout with the proper size and alignment.
             let size = NUM_PAGES * PAGE_SIZE;
-            layout = alloc::Layout::from_size_align_unchecked(size, PAGE_SIZE);
+            layout = std::alloc::Layout::from_size_align_unchecked(size, PAGE_SIZE);
 
             // Allocate the region of memory.
-            let ptr = alloc::alloc(layout);
+            let ptr = std::alloc::alloc(layout);
             if ptr.is_null() {
                 return Err(Error::from(std::io::ErrorKind::OutOfMemory));
             }
@@ -1059,6 +1055,40 @@ impl<'a> JitMemory<'a> {
         Ok(mem)
     }
 
+    #[cfg(not(feature = "std"))]
+    pub fn new(
+        prog: &[u8],
+        executable_memory: &'a mut [u8],
+        helpers: &HashMap<u32, ebpf::Helper>,
+        use_mbuff: bool,
+        update_data_ptr: bool,
+    ) -> Result<JitMemory<'a>, Error> {
+        let contents = executable_memory;
+        if contents.len() < NUM_PAGES * PAGE_SIZE {
+            return Err(Error::new(
+                ErrorKind::Other,
+                "Executable memory is too small",
+            ));
+        }
+        if contents.as_ptr() as usize % PAGE_SIZE != 0 {
+            return Err(Error::new(
+                ErrorKind::Other,
+                "Executable memory is not aligned",
+            ));
+        }
+
+        let mut mem = JitMemory {
+            contents,
+            offset: 0,
+        };
+
+        let mut jit = JitCompiler::new();
+        jit.jit_compile(&mut mem, prog, use_mbuff, update_data_ptr, helpers)?;
+        jit.resolve_jumps(&mut mem)?;
+
+        Ok(mem)
+    }
+
     pub fn get_prog(&self) -> MachineCode {
         unsafe { mem::transmute(self.contents.as_ptr()) }
     }
@@ -1078,15 +1108,16 @@ impl IndexMut<usize> for JitMemory<'_> {
     }
 }
 
+#[cfg(feature = "std")]
 impl Drop for JitMemory<'_> {
     fn drop(&mut self) {
         unsafe {
-            alloc::dealloc(self.contents.as_mut_ptr(), self.layout);
+            std::alloc::dealloc(self.contents.as_mut_ptr(), self.layout);
         }
     }
 }
 
-impl std::fmt::Debug for JitMemory<'_> {
+impl core::fmt::Debug for JitMemory<'_> {
     fn fmt(&self, fmt: &mut Formatter) -> Result<(), FormatterError> {
         fmt.write_str("JIT contents: [")?;
         fmt.write_str(" ] | ")?;

src/lib.rs

@@ -46,7 +46,7 @@ pub mod ebpf;
 pub mod helpers;
 pub mod insn_builder;
 mod interpreter;
-#[cfg(all(not(windows), feature = "std"))]
+#[cfg(not(windows))]
 mod jit;
 #[cfg(not(feature = "std"))]
 mod no_std_error;
@@ -79,6 +79,8 @@ pub mod lib {
     #[cfg(not(feature = "std"))]
     pub use alloc::vec::Vec;
     #[cfg(feature = "std")]
+    pub use std::vec;
+    #[cfg(feature = "std")]
     pub use std::vec::Vec;
 
     #[cfg(not(feature = "std"))]
@@ -108,6 +110,8 @@ pub mod lib {
 
     #[cfg(not(feature = "std"))]
     pub use alloc::format;
+    #[cfg(feature = "std")]
+    pub use std::format;
 }
 
 /// eBPF verification function that returns an error if the program does not meet its requirements.
@@ -171,7 +175,7 @@ struct MetaBuff {
 pub struct EbpfVmMbuff<'a> {
     prog: Option<&'a [u8]>,
     verifier: Verifier,
-    #[cfg(all(not(windows), feature = "std"))]
+    #[cfg(not(windows))]
     jit: Option<jit::JitMemory<'a>>,
     #[cfg(feature = "cranelift")]
     cranelift_prog: Option<cranelift::CraneliftProgram>,
@@ -209,7 +213,7 @@ impl<'a> EbpfVmMbuff<'a> {
         Ok(EbpfVmMbuff {
             prog,
             verifier: verifier::check,
-            #[cfg(all(not(windows), feature = "std"))]
+            #[cfg(not(windows))]
             jit: None,
             #[cfg(feature = "cranelift")]
             cranelift_prog: None,
@@ -487,6 +491,28 @@ impl<'a> EbpfVmMbuff<'a> {
         Ok(())
     }
 
+    /// JIT-compile the loaded program. The user has to provide a mutable slice of memory that
+    /// will be used for the JIT-compiled code. For more information, see the [EbpfVmMbuff::jit_compile]
+    /// function.
+    #[cfg(all(not(windows), not(feature = "std")))]
+    pub fn jit_compile(&mut self, executable_memory: &'a mut [u8]) -> Result<(), Error> {
+        let prog = match self.prog {
+            Some(prog) => prog,
+            None => Err(Error::new(
+                ErrorKind::Other,
+                "Error: No program set, call prog_set() to load one",
+            ))?,
+        };
+        self.jit = Some(jit::JitMemory::new(
+            prog,
+            executable_memory,
+            &self.helpers,
+            true,
+            false,
+        )?);
+        Ok(())
+    }
+
     /// Execute the previously JIT-compiled program, with the given packet data and metadata
     /// buffer, in a manner very similar to `execute_program()`.
     ///
@@ -539,7 +565,7 @@ impl<'a> EbpfVmMbuff<'a> {
     ///     assert_eq!(res, 0x2211);
     /// }
     /// ```
-    #[cfg(all(not(windows), feature = "std"))]
+    #[cfg(not(windows))]
     pub unsafe fn execute_program_jit(
         &self,
         mem: &mut [u8],
@@ -550,7 +576,7 @@ impl<'a> EbpfVmMbuff<'a> {
         //  in the kernel; anyway the verifier would prevent the use of uninitialized registers).
         //  See `mul_loop` test.
         let mem_ptr = match mem.len() {
-            0 => std::ptr::null_mut(),
+            0 => core::ptr::null_mut(),
             _ => mem.as_ptr() as *mut u8,
         };
         // The last two arguments are not used in this function. They would be used if there was a
@@ -1079,6 +1105,28 @@ impl<'a> EbpfVmFixedMbuff<'a> {
         Ok(())
     }
 
+    /// JIT-compile the loaded program. The user has to provide a mutable slice of memory that
+    /// will be used for the JIT-compiled code. For more information, see the [EbpfVmFixedMbuff::jit_compile]
+    /// function.
+    #[cfg(all(not(windows), not(feature = "std")))]
+    pub fn jit_compile(&mut self, executable_memory: &'a mut [u8]) -> Result<(), Error> {
+        let prog = match self.parent.prog {
+            Some(prog) => prog,
+            None => Err(Error::new(
+                ErrorKind::Other,
+                "Error: No program set, call prog_set() to load one",
+            ))?,
+        };
+        self.parent.jit = Some(jit::JitMemory::new(
+            prog,
+            executable_memory,
+            &self.parent.helpers,
+            true,
+            true,
+        )?);
+        Ok(())
+    }
+
     /// Execute the previously JIT-compiled program, with the given packet data, in a manner very
     /// similar to `execute_program()`.
     ///
@@ -1127,7 +1175,7 @@ impl<'a> EbpfVmFixedMbuff<'a> {
     /// ```
     // This struct redefines the `execute_program_jit()` function, in order to pass the offsets
     // associated with the fixed mbuff.
-    #[cfg(all(not(windows), feature = "std"))]
+    #[cfg(not(windows))]
     pub unsafe fn execute_program_jit(&mut self, mem: &'a mut [u8]) -> Result<u64, Error> {
         // If packet data is empty, do not send the address of an empty slice; send a null pointer
         //  as first argument instead, as this is uBPF's behavior (empty packet should not happen
@@ -1565,6 +1613,28 @@ impl<'a> EbpfVmRaw<'a> {
         Ok(())
     }
 
+    /// JIT-compile the loaded program. The user has to provide a mutable slice of memory that
+    /// will be used for the JIT-compiled code. For more information, see the [EbpfVmRaw::jit_compile]
+    /// function.
+    #[cfg(all(not(windows), not(feature = "std")))]
+    pub fn jit_compile(&mut self, executable_memory: &'a mut [u8]) -> Result<(), Error> {
+        let prog = match self.parent.prog {
+            Some(prog) => prog,
+            None => Err(Error::new(
+                ErrorKind::Other,
+                "Error: No program set, call prog_set() to load one",
+            ))?,
+        };
+        self.parent.jit = Some(jit::JitMemory::new(
+            prog,
+            executable_memory,
+            &self.parent.helpers,
+            false,
+            false,
+        )?);
+        Ok(())
+    }
+
     /// Execute the previously JIT-compiled program, with the given packet data, in a manner very
     /// similar to `execute_program()`.
     ///
@@ -1602,7 +1672,7 @@ impl<'a> EbpfVmRaw<'a> {
     ///     assert_eq!(res, 0x22cc);
     /// }
     /// ```
-    #[cfg(all(not(windows), feature = "std"))]
+    #[cfg(not(windows))]
     pub unsafe fn execute_program_jit(&self, mem: &'a mut [u8]) -> Result<u64, Error> {
         let mut mbuff = vec![];
         self.parent.execute_program_jit(mem, &mut mbuff)
@@ -1942,6 +2012,14 @@ impl<'a> EbpfVmNoData<'a> {
         self.parent.jit_compile()
     }
 
+    /// JIT-compile the loaded program. The user has to provide a mutable slice of memory that
+    /// will be used for the JIT-compiled code. For more information, see the [EbpfVmNoData::jit_compile]
+    /// function.
+    #[cfg(all(not(windows), not(feature = "std")))]
+    pub fn jit_compile(&mut self, executable_memory: &'a mut [u8]) -> Result<(), Error> {
+        self.parent.jit_compile(executable_memory)
+    }
+
     /// Execute the program loaded, without providing pointers to any memory area whatsoever.
     ///
     /// # Examples
@@ -1995,7 +2073,7 @@ impl<'a> EbpfVmNoData<'a> {
     ///     assert_eq!(res, 0x1122);
     /// }
     /// ```
-    #[cfg(all(not(windows), feature = "std"))]
+    #[cfg(not(windows))]
     pub unsafe fn execute_program_jit(&self) -> Result<u64, Error> {
         self.parent.execute_program_jit(&mut [])
     }