refactor(ebpf): clean trait api

Signed-off-by: Felix Hilgers <[email protected]>

rust/playground/ebpf-refactored/src/event_local.rs

@@ -2,14 +2,102 @@
 //
 // SPDX-License-Identifier: MIT
 
-use aya_ebpf::maps::HashMap;
+use core::ops::Deref;
+
+use aya_ebpf::{bindings::BPF_NOEXIST, maps::HashMap};
+use ebpf_types::{EventData, EventKind};
 
 #[repr(C)]
-pub struct Args<T> {
-    pub args: [u64; 6],
-    pub extra: T,
+pub struct EventLocal<T: EventLocalData> {
+    atomic_held: u64,
+    pub data: T::Data,
+}
+
+pub struct EventLocalStorage<T: 'static + EventLocalData>(&'static HashMap<u64, EventLocal<T>>);
+
+pub trait EventLocalData: EventData + 'static {
+    type Data;
+}
+
+pub struct PlaceHolder;
+
+impl EventData for PlaceHolder {
+    const EVENT_KIND: EventKind = EventKind::MAX;
+}
+
+impl EventLocalData for PlaceHolder {
+    type Data = [u8; 8192];
 }
 
-pub struct EventLocal<T: 'static>(&'static HashMap<u64, Args<T>>);
+impl<T: EventLocalData> EventLocalStorage<T> {
+    /// # Safety
+    ///
+    /// You must ensure that the map is only accessed throught this class
+    /// for it to be safe.
+    pub const unsafe fn new(map: &'static HashMap<u64, EventLocal<T>>) -> Self {
+        Self(map)
+    }
+
+    #[inline(always)]
+    pub fn cast<U: 'static + EventLocalData>(&self) -> &EventLocalStorage<U> {
+        const {
+            assert!(size_of::<T::Data>() >= size_of::<U::Data>());
+        };
+        unsafe { &*(self as *const _ as *const _) }
+    }
+
+    fn get_key<D: EventLocalData>(key: u32) -> u64 {
+        ((D::EVENT_KIND as u64) << 32) | key as u64
+    }
+
+    pub fn set(&self, key: u32, data: &mut EventLocal<T>) -> Result<(), i64> {
+        data.atomic_held = 0;
 
-impl<T> EventLocal<T> {}
+        let full_key = Self::get_key::<T>(key);
+        if self.0.insert(&full_key, data, BPF_NOEXIST as u64).is_ok() {
+            return Ok(());
+        }
+
+        let _ = self.take(key)?;
+
+        self.0.insert(&full_key, data, BPF_NOEXIST as u64)
+    }
+
+    #[inline(always)]
+    pub fn take(&self, key: u32) -> Result<EventLocalValue<T>, i64> {
+        let key = Self::get_key::<T>(key);
+        let ptr = self.0.get_ptr_mut(&key).ok_or(-1)?;
+        let is_held =
+            unsafe { core::intrinsics::atomic_xchg_seqcst(&raw mut (*ptr).atomic_held, 1) };
+        if is_held != 0 {
+            return Err(-1);
+        }
+
+        unsafe {
+            Ok(EventLocalValue {
+                key,
+                map: self.0,
+                value: &*ptr,
+            })
+        }
+    }
+}
+
+pub struct EventLocalValue<T: EventLocalData + 'static> {
+    key: u64,
+    map: &'static HashMap<u64, EventLocal<T>>,
+    value: &'static EventLocal<T>,
+}
+
+impl<T: EventLocalData + 'static> Deref for EventLocalValue<T> {
+    type Target = EventLocal<T>;
+    fn deref(&self) -> &Self::Target {
+        self.value
+    }
+}
+
+impl<T: EventLocalData + 'static> Drop for EventLocalValue<T> {
+    fn drop(&mut self) {
+        let _ = self.map.remove(&self.key);
+    }
+}

rust/playground/ebpf-refactored/src/events/blocking.rs

@@ -2,18 +2,70 @@
 //
 // SPDX-License-Identifier: MIT
 
+use core::mem::MaybeUninit;
+
 use aya_ebpf::helpers::bpf_ktime_get_ns;
 use ebpf_types::Blocking;
 
-pub fn initialize_blocking_enter(syscall_id: i64, blocking_data: &mut Blocking) -> Option<()> {
-    blocking_data.syscall_id = syscall_id as u64;
-    blocking_data.duration = unsafe { bpf_ktime_get_ns() };
+use super::SyscallProg;
+use crate::{
+    event_local::{EventLocal, EventLocalData, EventLocalValue},
+    pipeline::{ProgramInfo, SysEnterInfo, SysExitInfo},
+};
+
+#[repr(C)]
+pub struct BlockingEntryData {
+    pub syscall_id: u64,
+    pub start_time: u64,
+}
+
+impl EventLocalData for Blocking {
+    type Data = BlockingEntryData;
+}
+
+impl SyscallProg for Blocking {
+    fn enter<'a>(
+        sys_enter: SysEnterInfo,
+        _: ProgramInfo,
+        mem: &'a mut MaybeUninit<EventLocal<Self>>,
+    ) -> Option<&'a mut EventLocal<Self>> {
+        initialize_blocking_enter(sys_enter.syscall_id, mem)
+    }
 
-    Some(())
+    fn exit<'a>(
+        _: SysExitInfo,
+        _: ProgramInfo,
+        entry: &EventLocalValue<Self>,
+        mem: &'a mut MaybeUninit<Self>,
+    ) -> Option<&'a Self> {
+        initialize_blocking_exit(entry, mem)
+    }
 }
 
-pub fn initialize_blocking_exit(blocking_data: &mut Blocking) -> Option<()> {
-    blocking_data.duration = unsafe { bpf_ktime_get_ns() - blocking_data.duration };
+fn initialize_blocking_enter(
+    syscall_id: i64,
+    blocking_data: &mut MaybeUninit<EventLocal<Blocking>>,
+) -> Option<&mut EventLocal<Blocking>> {
+    let ptr = blocking_data.as_mut_ptr();
+
+    unsafe {
+        (&raw mut (*ptr).data.syscall_id).write(syscall_id as u64);
+        (&raw mut (*ptr).data.start_time).write(bpf_ktime_get_ns());
+
+        Some(blocking_data.assume_init_mut())
+    }
+}
+
+fn initialize_blocking_exit<'a>(
+    blocking_entry: &EventLocalValue<Blocking>,
+    blocking_data: &'a mut MaybeUninit<Blocking>,
+) -> Option<&'a Blocking> {
+    let ptr = blocking_data.as_mut_ptr();
+
+    unsafe {
+        (&raw mut (*ptr).duration).write(bpf_ktime_get_ns() - blocking_entry.data.start_time);
+        (&raw mut (*ptr).syscall_id).write(blocking_entry.data.syscall_id);
 
-    Some(())
+        Some(blocking_data.assume_init_mut())
+    }
 }

rust/playground/ebpf-refactored/src/events/fdtracking.rs

@@ -2,28 +2,75 @@
 //
 // SPDX-License-Identifier: MIT
 
+use core::mem::MaybeUninit;
+
 use aya_ebpf::helpers::bpf_probe_read_kernel;
 use ebpf_relocation_helpers::TaskStruct;
 use ebpf_types::{FileDescriptorChange, FileDescriptorOp};
 
-use crate::syscalls;
+use super::SyscallProg;
+use crate::{
+    event_local::{EventLocal, EventLocalData, EventLocalValue},
+    pipeline::{ProgramInfo, SysEnterInfo, SysExitInfo},
+    syscalls,
+};
+
+#[repr(C)]
+pub struct FdTrackingEnter {
+    operation: FileDescriptorOp,
+}
+
+impl EventLocalData for FileDescriptorChange {
+    type Data = FdTrackingEnter;
+}
+
+impl SyscallProg for FileDescriptorChange {
+    fn enter<'a>(
+        sys_enter: SysEnterInfo,
+        _: ProgramInfo,
+        mem: &'a mut MaybeUninit<EventLocal<Self>>,
+    ) -> Option<&'a mut EventLocal<Self>> {
+        initialize_fdtracking_enter(sys_enter.syscall_id, mem)
+    }
+
+    fn exit<'a>(
+        sys_exit: SysExitInfo,
+        _: ProgramInfo,
+        entry: &EventLocalValue<Self>,
+        mem: &'a mut MaybeUninit<Self>,
+    ) -> Option<&'a Self> {
+        initialize_fdtracking_exit(sys_exit.task, entry, mem)
+    }
+}
 
-pub fn initialize_fdtracking_enter(
+#[inline(always)]
+fn initialize_fdtracking_enter(
     syscall_id: i64,
-    fdtracking_data: &mut FileDescriptorChange,
-) -> Option<()> {
-    fdtracking_data.operation = get_file_op(syscall_id)?;
+    fdtracking_data: &mut MaybeUninit<EventLocal<FileDescriptorChange>>,
+) -> Option<&mut EventLocal<FileDescriptorChange>> {
+    let ptr = fdtracking_data.as_mut_ptr();
+
+    unsafe {
+        (&raw mut (*ptr).data.operation).write(get_file_op(syscall_id)?);
 
-    Some(())
+        Some(fdtracking_data.assume_init_mut())
+    }
 }
 
-pub fn initialize_fdtracking_exit(
+#[inline(always)]
+fn initialize_fdtracking_exit<'a>(
     task: TaskStruct,
-    fdtracking_data: &mut FileDescriptorChange,
-) -> Option<()> {
-    fdtracking_data.open_fds = get_open_fds(task)?;
+    fdtracking_enter: &EventLocalValue<FileDescriptorChange>,
+    fdtracking_data: &'a mut MaybeUninit<FileDescriptorChange>,
+) -> Option<&'a FileDescriptorChange> {
+    let ptr = fdtracking_data.as_mut_ptr();
 
-    Some(())
+    unsafe {
+        (&raw mut (*ptr).open_fds).write(get_open_fds(task)?);
+        (&raw mut (*ptr).operation).write(fdtracking_enter.data.operation);
+
+        Some(fdtracking_data.assume_init_ref())
+    }
 }
 
 pub fn get_file_op(syscall_number: i64) -> Option<FileDescriptorOp> {

rust/playground/ebpf-refactored/src/events/mod.rs

@@ -2,7 +2,28 @@
 //
 // SPDX-License-Identifier: MIT
 
+use core::mem::MaybeUninit;
+
+use crate::{
+    event_local::{EventLocal, EventLocalData, EventLocalValue},
+    pipeline::{ProgramInfo, SysEnterInfo, SysExitInfo},
+};
+
 pub mod blocking;
 pub mod fdtracking;
 pub mod signal;
 pub mod write;
+
+pub trait SyscallProg: EventLocalData + Sized {
+    fn enter<'a>(
+        sys_enter: SysEnterInfo,
+        program_info: ProgramInfo,
+        mem: &'a mut MaybeUninit<EventLocal<Self>>,
+    ) -> Option<&'a mut EventLocal<Self>>;
+    fn exit<'a>(
+        sys_exit: SysExitInfo,
+        program_info: ProgramInfo,
+        entry: &EventLocalValue<Self>,
+        mem: &'a mut MaybeUninit<Self>,
+    ) -> Option<&'a Self>;
+}

rust/playground/ebpf-refactored/src/events/signal.rs

@@ -2,22 +2,70 @@
 //
 // SPDX-License-Identifier: MIT
 
+use core::mem::MaybeUninit;
+
 use aya_ebpf::PtRegs;
 use ebpf_types::Signal;
 
-use crate::syscalls;
+use super::SyscallProg;
+use crate::{
+    event_local::{EventLocal, EventLocalData, EventLocalValue},
+    pipeline::{ProgramInfo, SysEnterInfo, SysExitInfo},
+    syscalls,
+};
+
+impl EventLocalData for Signal {
+    type Data = Signal;
+}
+
+impl SyscallProg for Signal {
+    fn enter<'a>(
+        sys_enter: SysEnterInfo,
+        _: ProgramInfo,
+        mem: &'a mut MaybeUninit<EventLocal<Self>>,
+    ) -> Option<&'a mut EventLocal<Self>> {
+        initialize_signal_enter(sys_enter.syscall_id, sys_enter.pt_regs, mem)
+    }
+
+    fn exit<'a>(
+        _: SysExitInfo,
+        _: ProgramInfo,
+        entry: &EventLocalValue<Self>,
+        mem: &'a mut MaybeUninit<Self>,
+    ) -> Option<&'a Self> {
+        initialize_signal_exit(entry, mem)
+    }
+}
 
-pub fn initialize_signal_enter(
+#[inline(always)]
+fn initialize_signal_enter(
     syscall_id: i64,
     pt_regs: PtRegs,
-    signal_data: &mut Signal,
-) -> Option<()> {
+    signal_data: &mut MaybeUninit<EventLocal<Signal>>,
+) -> Option<&mut EventLocal<Signal>> {
     if syscall_id != syscalls::SYS_kill {
         return None;
     }
 
-    signal_data.target_pid = pt_regs.arg::<*const u64>(0)? as i32;
-    signal_data.signal = pt_regs.arg::<*const u64>(1)? as u32;
+    let ptr = signal_data.as_mut_ptr();
+    unsafe {
+        (&raw mut (*ptr).data.target_pid).write(pt_regs.arg::<*const u64>(0)? as i32);
+        (&raw mut (*ptr).data.signal).write(pt_regs.arg::<*const u64>(1)? as u32);
+
+        Some(signal_data.assume_init_mut())
+    }
+}
 
-    Some(())
+#[inline(always)]
+fn initialize_signal_exit<'a>(
+    signal_entry: &EventLocalValue<Signal>,
+    signal_data: &'a mut MaybeUninit<Signal>,
+) -> Option<&'a Signal> {
+    let ptr = signal_data.as_mut_ptr();
+
+    unsafe {
+        ptr.write(signal_entry.data);
+
+        Some(signal_data.assume_init_mut())
+    }
 }