#[cfg(feature = "shuttle-test")]
use shuttle::rand;
use {
crate::{
ebpf,
elf::ElfError,
error::EbpfError,
memory_management::{get_system_page_size, round_to_page_size, Mapping, Read, ReadExecute},
vm::{Config, ContextObject, EncryptedHostAddressToEbpfVm},
},
rand::{thread_rng, Rng},
std::collections::{hash_map::Entry, HashMap},
std::hash::{BuildHasher, Hasher},
};
#[derive(Debug, PartialEq, PartialOrd, Eq, Clone, Copy)]
pub enum SBPFVersion {
V0,
V1,
V2,
V3,
V4,
Reserved,
}
impl SBPFVersion {
pub const ALL: &'static [SBPFVersion] = &[
SBPFVersion::V0,
SBPFVersion::V1,
SBPFVersion::V2,
SBPFVersion::V3,
SBPFVersion::V4,
];
pub const fn manual_stack_frame_bump(self) -> bool {
matches!(self, SBPFVersion::V1 | SBPFVersion::V2)
}
pub const fn stack_frame_gaps(self) -> bool {
matches!(self, SBPFVersion::V0)
}
pub const fn enable_pqr(self) -> bool {
matches!(self, SBPFVersion::V2)
}
pub const fn explicit_sign_extension_of_results(self) -> bool {
matches!(self, SBPFVersion::V2)
}
pub const fn swap_sub_reg_imm_operands(self) -> bool {
matches!(self, SBPFVersion::V2)
}
pub const fn disable_neg(self) -> bool {
matches!(self, SBPFVersion::V2)
}
pub const fn callx_uses_src_reg(self) -> bool {
matches!(self, SBPFVersion::V2)
}
pub const fn disable_lddw(self) -> bool {
matches!(self, SBPFVersion::V2)
}
pub const fn disable_le(self) -> bool {
matches!(self, SBPFVersion::V2)
}
pub const fn move_memory_instruction_classes(self) -> bool {
matches!(self, SBPFVersion::V2)
}
pub const fn static_syscalls(self) -> bool {
matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
}
pub const fn enable_stricter_elf_headers(self) -> bool {
matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
}
pub const fn enable_lower_rodata_vaddr(self) -> bool {
matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
}
pub const fn enable_jmp32(self) -> bool {
matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
}
pub const fn callx_uses_dst_reg(self) -> bool {
matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
}
pub const fn calculate_call_imm_target_pc(self, pc: usize, imm: i64) -> u32 {
if self.static_syscalls() {
(pc as i64).saturating_add(imm).saturating_add(1) as u32
} else {
imm as u32
}
}
}
#[derive(Debug, PartialEq, Eq)]
pub struct FunctionRegistry<T> {
map: HashMap<u32, (Vec<u8>, T), KeyHasherRandomState>,
}
impl<T> Default for FunctionRegistry<T> {
fn default() -> Self {
Self {
map: HashMap::default(),
}
}
}
#[derive(Clone, Copy)]
struct KeyHasherRandomState {
multiplier: u64,
}
impl Default for KeyHasherRandomState {
fn default() -> Self {
Self {
multiplier: thread_rng().gen::<u64>() | 1,
}
}
}
impl BuildHasher for KeyHasherRandomState {
type Hasher = KeyHasher;
fn build_hasher(&self) -> KeyHasher {
KeyHasher {
multiplier: self.multiplier,
hash: 0,
}
}
}
struct KeyHasher {
multiplier: u64,
hash: u64,
}
impl Hasher for KeyHasher {
fn finish(&self) -> u64 {
self.hash
}
fn write(&mut self, bytes: &[u8]) {
for &byte in bytes {
self.write_u32(self.hash as u32 ^ u32::from(byte));
}
}
fn write_u32(&mut self, key: u32) {
let product = u64::from(key).wrapping_mul(self.multiplier);
self.hash = product ^ product >> 32;
}
}
impl<T: Copy + PartialEq> FunctionRegistry<T> {
pub fn register_function(
&mut self,
key: u32,
name: impl Into<Vec<u8>>,
value: T,
) -> Result<(), ElfError> {
match self.map.entry(key) {
Entry::Vacant(entry) => {
entry.insert((name.into(), value));
}
Entry::Occupied(entry) => {
if entry.get().1 != value {
return Err(ElfError::SymbolHashCollision(key));
}
}
}
Ok(())
}
pub(crate) fn register_function_hashed_legacy<C: ContextObject>(
&mut self,
loader: &BuiltinProgram<C>,
hash_symbol_name: bool,
name: impl Into<Vec<u8>>,
value: T,
) -> Result<u32, ElfError>
where
usize: From<T>,
{
let name = name.into();
let config = loader.get_config();
let key = if hash_symbol_name {
let hash = if name == b"entrypoint" {
ebpf::hash_symbol_name(b"entrypoint")
} else {
ebpf::hash_symbol_name(&usize::from(value).to_le_bytes())
};
if loader.get_function_registry().lookup_by_key(hash).is_some() {
return Err(ElfError::SymbolHashCollision(hash));
}
hash
} else {
usize::from(value) as u32
};
self.register_function(
key,
if config.enable_symbol_and_section_labels || name == b"entrypoint" {
name
} else {
Vec::default()
},
value,
)?;
Ok(key)
}
pub fn unregister_function(&mut self, key: u32) -> bool {
self.map.remove(&key).is_some()
}
pub fn keys(&self) -> impl Iterator<Item = u32> + '_ {
self.map.keys().copied()
}
pub fn iter(&self) -> impl Iterator<Item = (u32, (&[u8], T))> + '_ {
self.map
.iter()
.map(|(key, (name, value))| (*key, (name.as_slice(), *value)))
}
#[inline]
pub fn lookup_by_key(&self, key: u32) -> Option<(&[u8], T)> {
self.map
.get(&key)
.map(|(function_name, value)| (function_name.as_slice(), *value))
}
pub fn lookup_by_name(&self, name: &[u8]) -> Option<(&[u8], T)> {
self.map
.values()
.find(|(function_name, _value)| function_name == name)
.map(|(function_name, value)| (function_name.as_slice(), *value))
}
pub fn mem_size(&self) -> usize {
std::mem::size_of::<Self>().saturating_add(self.map.values().fold(
0,
|state: usize, (name, value)| {
state.saturating_add(
std::mem::size_of_val(value).saturating_add(
std::mem::size_of_val(name).saturating_add(name.capacity()),
),
)
},
))
}
}
pub struct UnsealedJitProgram {
pc_section: Mapping,
text_section: Mapping,
}
impl UnsealedJitProgram {
pub fn new(pc: usize, code_capacity: usize) -> Self {
let page_size = get_system_page_size();
let pc_len = pc.saturating_mul(std::mem::size_of::<u32>());
let pc_size = round_to_page_size(pc_len, page_size);
let text_capacity = round_to_page_size(code_capacity, page_size);
let mapping = Mapping::pooled(pc_size.saturating_add(text_capacity));
let Ok([pc_section, text_section]) =
mapping.split([0..pc_len, pc_size..pc_size.saturating_add(text_capacity)])
else {
unreachable!("mapping split failed for some reason");
};
let mut program = Self {
pc_section,
text_section,
};
program.pc_section.fill(0);
program
}
#[inline]
pub fn pc_section(&self) -> &[u32] {
unsafe {
let (&[], out, &[]) = self.pc_section.align_to() else {
unreachable!("pc_section can't be unaligned to u32");
};
out
}
}
#[inline]
pub fn text_section(&self) -> &[u8] {
&self.text_section
}
#[inline]
pub fn pc_section_mut(&mut self) -> &mut [u32] {
unsafe {
let (&mut [], out, &mut []) = self.pc_section.align_to_mut() else {
unreachable!("pc_section can't be unaligned to u32");
};
out
}
}
#[inline]
pub fn text_section_mut(&mut self) -> &mut [u8] {
&mut self.text_section
}
pub fn seal(self, used: usize) -> Result<JitProgram, EbpfError> {
let Self {
pc_section,
mut text_section,
} = self;
if used > text_section.len() {
return Err(EbpfError::ExhaustedTextSegment(used));
}
let code_size = round_to_page_size(used, get_system_page_size());
text_section[used..code_size].fill(0xcc);
text_section.truncate(used);
Ok(JitProgram {
pc_section: pc_section.protect()?,
text_section: text_section.protect()?,
})
}
}
pub struct JitProgram {
pc_section: Mapping<Read>,
text_section: Mapping<ReadExecute>,
}
impl JitProgram {
#[inline]
pub fn pc_section(&self) -> &[u32] {
unsafe {
let (&[], out, &[]) = self.pc_section.align_to() else {
unreachable!("pc_section can't be unaligned to u32");
};
out
}
}
#[inline]
pub fn text_section(&self) -> &[u8] {
&self.text_section
}
pub fn machine_code_length(&self) -> usize {
self.text_section().len()
}
pub fn mem_size(&self) -> usize {
self.text_section.mapped_len()
}
}
impl std::fmt::Debug for JitProgram {
fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
fmt.write_fmt(format_args!("JitProgram {:?}", self as *const _))
}
}
impl PartialEq for JitProgram {
fn eq(&self, other: &Self) -> bool {
std::ptr::eq(self, other)
}
}
pub type BuiltinFunction<C> = fn(EncryptedHostAddressToEbpfVm<C>, u64, u64, u64, u64, u64);
#[cfg(all(feature = "jit", not(target_os = "windows"), target_arch = "x86_64"))]
pub type JitCompiler<'a, C> = crate::jit::JitCompiler<'a, C>;
#[cfg(not(all(feature = "jit", not(target_os = "windows"), target_arch = "x86_64")))]
pub struct JitCompiler<'a, C> {
_phantom: std::marker::PhantomData<&'a C>,
}
#[cfg(not(all(feature = "jit", not(target_os = "windows"), target_arch = "x86_64")))]
impl<'a, C: ContextObject> JitCompiler<'a, C> {
#[allow(dead_code)]
pub fn emit_external_call(&mut self, _function: BuiltinFunction<C>) {}
}
pub type BuiltinCodegen<C> = fn(&mut JitCompiler<C>);
#[derive(Eq)]
pub struct BuiltinProgram<C: ContextObject> {
config: Option<Box<Config>>,
sparse_registry: FunctionRegistry<(BuiltinFunction<C>, BuiltinCodegen<C>)>,
}
impl<C: ContextObject> PartialEq for BuiltinProgram<C> {
fn eq(&self, other: &Self) -> bool {
self.config.eq(&other.config) && self.sparse_registry.eq(&other.sparse_registry)
}
}
impl<C: ContextObject> BuiltinProgram<C> {
pub fn new_loader(config: Config) -> Self {
Self {
config: Some(Box::new(config)),
sparse_registry: FunctionRegistry::default(),
}
}
pub fn new_builtin() -> Self {
Self {
config: None,
sparse_registry: FunctionRegistry::default(),
}
}
pub fn new_mock() -> Self {
Self {
config: Some(Box::default()),
sparse_registry: FunctionRegistry::default(),
}
}
pub fn get_config(&self) -> &Config {
self.config.as_ref().unwrap()
}
pub fn get_function_registry(
&self,
) -> &FunctionRegistry<(BuiltinFunction<C>, BuiltinCodegen<C>)> {
&self.sparse_registry
}
pub fn mem_size(&self) -> usize {
std::mem::size_of::<Self>()
.saturating_add(if self.config.is_some() {
std::mem::size_of::<Config>()
} else {
0
})
.saturating_add(self.sparse_registry.mem_size())
}
pub fn register_function(
&mut self,
name: &str,
entry: (BuiltinFunction<C>, BuiltinCodegen<C>),
) -> Result<(), ElfError> {
let key = ebpf::hash_symbol_name(name.as_bytes());
self.sparse_registry
.register_function(key, name, entry)
.map(|_| ())
}
pub fn register_definition<BFD: BuiltinFunctionDefinition<C>>(
&mut self,
name: &str,
) -> Result<(), ElfError> {
self.register_function(name, (BFD::vm, BFD::codegen))
}
pub fn unregister_function(&mut self, name: &str) -> bool {
let key = ebpf::hash_symbol_name(name.as_bytes());
self.sparse_registry.unregister_function(key)
}
}
pub trait BuiltinFunctionDefinition<C>
where
C: crate::vm::ContextObject,
{
type Error: Into<Box<dyn core::error::Error>>;
fn rust(
vm: &mut C,
arg_a: u64,
arg_b: u64,
arg_c: u64,
arg_d: u64,
arg_e: u64,
) -> Result<u64, Self::Error>;
#[expect(clippy::arithmetic_side_effects)]
fn vm(mut vm: EncryptedHostAddressToEbpfVm<C>, a: u64, b: u64, c: u64, d: u64, e: u64) {
unsafe {
vm.with_vm(|vm| {
let enable_insn_meter = vm.loader.get_config().enable_instruction_meter;
if enable_insn_meter {
let used_cus = vm.previous_instruction_meter - vm.due_insn_count;
vm.context().consume(used_cus);
}
let converted_result: crate::error::ProgramResult =
Self::rust(vm.context(), a, b, c, d, e)
.map_err(|err| crate::error::EbpfError::SyscallError(err.into()))
.into();
vm.program_result = converted_result;
if enable_insn_meter {
vm.previous_instruction_meter = vm.context().get_remaining();
}
})
}
}
fn codegen(jit: &mut JitCompiler<C>) {
jit.emit_external_call(Self::vm);
}
fn register(program: &mut BuiltinProgram<C>, name: &str) -> Result<(), ElfError>
where
Self: Sized,
{
program.register_definition::<Self>(name)
}
}
impl<C: ContextObject> std::fmt::Debug for BuiltinProgram<C> {
fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
f.debug_struct("BuiltinProgram")
.field("registry", unsafe {
std::mem::transmute::<
&FunctionRegistry<(BuiltinFunction<C>, BuiltinCodegen<C>)>,
&FunctionRegistry<(usize, usize)>,
>(&self.sparse_registry)
})
.finish()
}
}
#[macro_export]
macro_rules! declare_builtin_function {
(
$(#[$attr:meta])*
$name:ident $(<$($generic_ident:tt : $generic_type:tt),+>)?,
fn rust(
$vm:ident : &mut $ContextObject:ty,
$arg_a:ident : u64,
$arg_b:ident : u64,
$arg_c:ident : u64,
$arg_d:ident : u64,
$arg_e:ident : u64,
) -> Result<$Ok:ty, $Err:ty> {
$($rust:tt)*
}
$(fn codegen($jit:ident : &mut JitCompiler<$ContextObject2:ty>) {
$($codegen:tt)*
})?
) => {
$(#[$attr])*
pub struct $name $(<$($generic_ident),+>)? (
$(std::marker::PhantomData<($($generic_ident,)+)>)?
);
impl $(<$($generic_ident : $generic_type),+>)?
$crate::program::BuiltinFunctionDefinition<$ContextObject> for
$name $(<$($generic_ident),+>)?
{
type Error = $Err;
fn rust(
$vm: &mut $ContextObject,
$arg_a: u64,
$arg_b: u64,
$arg_c: u64,
$arg_d: u64,
$arg_e: u64,
) -> core::result::Result<$Ok, $Err> {
$($rust)*
}
$(fn codegen(
$jit: &mut $crate::program::JitCompiler<$ContextObject2>,
) {
$($codegen)*
})?
}
};
}
#[cfg(test)]
mod tests {
use super::*;
macro_rules! check_lookups {
($registry:ident, $lens:expr) => {
let key = |i: u32| ebpf::hash_symbol_name(&i.to_le_bytes());
for len in $lens {
let mut registry = $registry::default();
for i in 0..len {
registry
.register_function(key(i), format!("f{i}"), i)
.unwrap();
}
let found = |registry: &$registry<u32>, i: u32| {
registry
.lookup_by_key(key(i))
.map(|(name, value)| (String::from_utf8(name.to_vec()).unwrap(), value))
};
for i in 0..len {
assert_eq!(found(®istry, i), Some((format!("f{i}"), i)));
}
for i in len..len.saturating_add(1000) {
assert_eq!(found(®istry, i), None);
}
let mut keys: Vec<_> = (0..len).map(key).collect();
keys.sort();
let mut registered: Vec<_> = registry.keys().collect();
registered.sort();
assert_eq!(registered, keys);
let mut iterated: Vec<_> = registry.iter().map(|(key, _)| key).collect();
iterated.sort();
assert_eq!(iterated, keys);
if len > 0 {
assert!(registry.register_function(key(0), "f0", 0).is_ok());
assert!(matches!(
registry.register_function(key(0), "f0", 1),
Err(ElfError::SymbolHashCollision(_))
));
assert!(registry.unregister_function(key(0)));
assert!(!registry.unregister_function(key(0)));
assert_eq!(found(®istry, 0), None);
for i in 1..len {
assert_eq!(found(®istry, i), Some((format!("f{i}"), i)));
}
assert_eq!(registry.keys().count(), (len as usize).saturating_sub(1));
registry.register_function(key(0), "f0", 0).unwrap();
assert_eq!(found(®istry, 0), Some(("f0".to_string(), 0)));
for i in 0..len {
assert!(registry.unregister_function(key(i)));
}
assert_eq!(registry.keys().count(), 0);
for i in 0..len.saturating_add(1000) {
assert_eq!(found(®istry, i), None);
}
}
}
};
}
#[test]
fn function_registry_lookups() {
check_lookups!(FunctionRegistry, [0u32, 1, 3, 64, 1000]);
}
#[test]
fn jit_program_sections() {
let page_size = get_system_page_size();
let mut program = UnsealedJitProgram::new(3, page_size + 1);
assert_eq!(program.pc_section(), [0; 3]);
assert_eq!(program.text_section().len(), 2 * page_size);
program.pc_section_mut().copy_from_slice(&[1, 2, 3]);
program.text_section_mut()[..2].copy_from_slice(&[0x90, 0xc3]);
assert!(matches!(
UnsealedJitProgram::new(3, page_size + 1).seal(2 * page_size + 1),
Err(EbpfError::ExhaustedTextSegment(_))
));
let program = program.seal(2).unwrap();
assert_eq!(program.pc_section(), [1, 2, 3]);
assert_eq!(program.text_section(), [0x90, 0xc3]);
assert_eq!(program.machine_code_length(), 2);
assert!(program.mem_size() >= 3 * page_size);
}
#[test]
fn jit_program_pc_section_zeroed() {
for _ in 0..4 {
let mut program = UnsealedJitProgram::new(16, 16);
assert!(program.pc_section().iter().all(|&pc| pc == 0));
program.pc_section_mut().fill(u32::MAX);
program.seal(0).unwrap();
}
}
}