Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor - Make FunctionRegistry accept dense indexes #617

Closed
wants to merge 2 commits into from
Closed

Refactor - Make FunctionRegistry accept dense indexes #617

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a3b5fae
New data structure
LucasSte Oct 17, 2024
1c847bb
Rename new_sparse to default
LucasSte Oct 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/elf.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,9 @@ pub enum ElfError {
/// Invalid program header
#[error("Invalid ELF program header")]
InvalidProgramHeader,
/// Invalid syscall code
#[error("Invalid function index")]
InvalidDenseFunctionIndex,
}

impl From<ElfParserError> for ElfError {
Expand Down
180 changes: 140 additions & 40 deletions src/program.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,8 @@ use {
crate::{
ebpf,
elf::ElfError,
vm::{Config, ContextObject, EbpfVm},
error::{EbpfError, ProgramResult},
vm::{get_runtime_environment_key, Config, ContextObject, EbpfVm},
},
std::collections::{btree_map::Entry, BTreeMap},
};
Expand Down Expand Up @@ -85,14 +86,40 @@ impl SBPFVersion {

/// Holds the function symbols of an Executable
#[derive(Debug, PartialEq, Eq)]
pub struct FunctionRegistry<T> {
pub(crate) map: BTreeMap<u32, (Vec<u8>, T)>,
pub enum FunctionRegistry<T> {
/// Sparse allows the registration of hashed symbols
Sparse(BTreeMap<u32, (Vec<u8>, T)>),
/// Dense allows the registration of symbols indexed by integers only
Dense(Vec<(Vec<u8>, T)>),
}

impl<T> Default for FunctionRegistry<T> {
/// Default for FunctionRegistry returns a sparse registry
fn default() -> Self {
Self {
map: BTreeMap::new(),
FunctionRegistry::Sparse(BTreeMap::new())
}
}

impl<T: Copy + PartialEq + ContextObject> FunctionRegistry<BuiltinFunction<T>> {
Copy link
Author

@LucasSte LucasSte Oct 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I still want to add a few unit tests for these functions.

/// Create a dense function registry with pre-allocated spaces.
pub fn new_dense(size: usize) -> Self {
FunctionRegistry::Dense(vec![(b"tombstone".to_vec(), SyscallTombstone::vm); size])
}
}

impl<T: Default> FunctionRegistry<T> {
/// Unregister a symbol again
pub fn unregister_function(&mut self, key: u32) {
match self {
FunctionRegistry::Sparse(map) => {
map.remove(&key);
}
FunctionRegistry::Dense(vec) => {
if key == 0 || key > vec.len() as u32 {
return;
}
vec[key.saturating_sub(1) as usize] = (b"tombstone".to_vec(), T::default());
}
}
}
}
Expand All @@ -105,16 +132,25 @@ impl<T: Copy + PartialEq> FunctionRegistry<T> {
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));
match self {
FunctionRegistry::Sparse(map) => match map.entry(key) {
Entry::Vacant(entry) => {
entry.insert((name.into(), value));
}
Entry::Occupied(entry) => {
if entry.get().1 != value {
return Err(ElfError::SymbolHashCollision(key));
}
}
},
FunctionRegistry::Dense(vec) => {
if key == 0 || key > vec.len() as u32 {
return Err(ElfError::InvalidDenseFunctionIndex);
}
vec[key.saturating_sub(1) as usize] = (name.into(), value);
}
}

Ok(())
}

Expand Down Expand Up @@ -168,51 +204,80 @@ impl<T: Copy + PartialEq> FunctionRegistry<T> {
Ok(key)
}

/// Unregister a symbol again
pub fn unregister_function(&mut self, key: u32) {
self.map.remove(&key);
}

/// Iterate over all keys
pub fn keys(&self) -> impl Iterator<Item = u32> + '_ {
self.map.keys().copied()
pub fn keys(&self) -> Box<dyn Iterator<Item = u32> + '_> {
match self {
FunctionRegistry::Sparse(map) => Box::new(map.keys().copied()),
FunctionRegistry::Dense(vec) => Box::new(1..=(vec.len() as u32)),
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Needs to filter out the tombstones

}
}

/// Iterate over all entries
pub fn iter(&self) -> impl Iterator<Item = (u32, (&[u8], T))> + '_ {
self.map
.iter()
.map(|(key, (name, value))| (*key, (name.as_slice(), *value)))
#[allow(clippy::type_complexity)]
pub fn iter(&self) -> Box<dyn Iterator<Item = (u32, (&[u8], T))> + '_> {
match self {
FunctionRegistry::Sparse(map) => Box::new(
map.iter()
.map(|(key, (name, value))| (*key, (name.as_slice(), *value))),
),
FunctionRegistry::Dense(vec) => {
Box::new(vec.iter().enumerate().map(|(idx, value)| {
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Needs to filter out the tombstones

(idx.saturating_add(1) as u32, (value.0.as_slice(), value.1))
}))
}
}
}

/// Get a function by its key
pub fn lookup_by_key(&self, key: u32) -> Option<(&[u8], T)> {
// String::from_utf8_lossy(function_name).as_str()
self.map
.get(&key)
.map(|(function_name, value)| (function_name.as_slice(), *value))
match self {
FunctionRegistry::Sparse(map) => map
.get(&key)
.map(|(function_name, value)| (function_name.as_slice(), *value)),
FunctionRegistry::Dense(vec) => {
if key == 0 {
return None;
}
vec.get(key.saturating_sub(1) as usize)
.map(|(function_name, value)| (function_name.as_slice(), *value))
}
}
}

/// Get a function by its name
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))
match self {
FunctionRegistry::Sparse(map) => map
.values()
.find(|(function_name, _value)| function_name == name)
.map(|(function_name, value)| (function_name.as_slice(), *value)),
FunctionRegistry::Dense(vec) => vec
.iter()
.find(|(funtion_name, _)| funtion_name == name)
.map(|(function_name, value)| (function_name.as_slice(), *value)),
}
}

/// Calculate memory size
pub fn mem_size(&self) -> usize {
std::mem::size_of::<Self>().saturating_add(self.map.iter().fold(
0,
|state: usize, (_, (name, value))| {
state.saturating_add(
std::mem::size_of_val(value).saturating_add(
let self_size = std::mem::size_of::<Self>();
let content_size =
match &self {
FunctionRegistry::Sparse(map) => {
map.iter().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()),
))
})
}
FunctionRegistry::Dense(vec) => vec.iter().fold(0, |acc: usize, (name, value)| {
acc.saturating_add(std::mem::size_of_val(value).saturating_add(
std::mem::size_of_val(name).saturating_add(name.capacity()),
),
)
},
))
))
}),
};

self_size.saturating_add(content_size)
}
}

Expand Down Expand Up @@ -350,6 +415,41 @@ macro_rules! declare_builtin_function {
};
}

/// Tombstone function for when we call an invalid syscall (e.g. a syscall deactivated through a
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"when we call" makes it sound like runtime when it is about verification instead. Maybe "for leaving gaps in a sparse function registry"

/// feature gate
#[derive(Clone, PartialEq)]
pub struct SyscallTombstone {}
impl SyscallTombstone {
/// Tombstone function entrypoint
#[allow(clippy::arithmetic_side_effects)]
pub fn vm<T: ContextObject>(
ctx: *mut EbpfVm<T>,
_arg1: u64,
_arg2: u64,
_arg3: u64,
_arg4: u64,
_arg5: u64,
) {
let vm = unsafe {
&mut *(ctx
.cast::<u64>()
.offset(-(get_runtime_environment_key() as isize))
.cast::<EbpfVm<T>>())
};
let config = vm.loader.get_config();
if config.enable_instruction_meter {
vm.context_object_pointer
.consume(vm.previous_instruction_meter - vm.due_insn_count);
}
vm.program_result = ProgramResult::Err(EbpfError::SyscallError(Box::new(
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could also add a debug_assert!(false); to indicate that this is never supposed to actually run.

ElfError::InvalidDenseFunctionIndex,
)));
if config.enable_instruction_meter {
vm.previous_instruction_meter = vm.context_object_pointer.get_remaining();
}
}
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down