rex: clean up tracepoint and map impls

Signed-off-by: MinhPhan8803 <[email protected]>

Author: MinhPhan8803

rex-macros/src/kprobe.rs

@@ -1,6 +1,7 @@
+use std::fmt;
+
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
-use std::fmt;
 use syn::{parse2, ItemFn, Result};
 
 use crate::args::parse_string_args;

rex-macros/src/tracepoint.rs

@@ -34,8 +34,8 @@ impl TracePoint {
         if context_type_ref
             .lifetime
             .expect_or_abort("Context reference needs to be static")
-            .ident
-            != "static"
+            .ident !=
+            "static"
         {
             abort_call_site!("Context reference needs to be static");
         }
@@ -64,6 +64,7 @@ impl TracePoint {
             "SyscallsExitOpenCtx" => "syscalls/sys_exit_open",
             "SyscallsExitOpenatCtx" => "syscalls/sys_exit_openat",
             "SyscallsEnterDupCtx" => "syscalls/sys_enter_dup",
+            "RawSyscallsEnterCtx" => "raw_syscalls/sys_enter",
             _ => abort_call_site!("Please provide a valid context type. If your needed context isn't supported consider opening a PR!"),
         };
         let attached_name = format!("rex/tracepoint/{}", hook_point_name);

rex/src/lib.rs

@@ -31,7 +31,6 @@ mod random32;
 
 extern crate paste;
 
-use crate::bindings::uapi::linux::bpf::BPF_F_CURRENT_CPU;
 use paste::paste;
 pub use rex_macros::*;
 

rex/src/map.rs

@@ -9,32 +9,17 @@ use crate::base_helper::{
     termination_check,
 };
 use crate::ffi;
-use crate::utils::{NoRef, PerfEventMaskedCPU, Result, StreamableProgram};
 use crate::linux::bpf::{
     bpf_map_type, BPF_ANY, BPF_EXIST, BPF_MAP_TYPE_ARRAY, BPF_MAP_TYPE_HASH,
-    BPF_MAP_TYPE_PERCPU_ARRAY, BPF_MAP_TYPE_QUEUE, BPF_MAP_TYPE_RINGBUF,
-    BPF_MAP_TYPE_STACK, BPF_MAP_TYPE_STACK_TRACE, BPF_NOEXIST,
-    BPF_RB_AVAIL_DATA, BPF_RB_CONS_POS, BPF_RB_PROD_POS, BPF_RB_RING_SIZE}
-use crate::{
-    base_helper::{
-        bpf_map_delete_elem,
-        bpf_map_lookup_elem,
-        bpf_map_peek_elem,
-        bpf_map_pop_elem,
-        bpf_map_push_elem,
-        bpf_map_update_elem,
-        // bpf_ringbuf_discard, bpf_ringbuf_query, bpf_ringbuf_reserve,
-        // bpf_ringbuf_submit,
-    },
-    linux::bpf::{
-        bpf_map_type, BPF_ANY, BPF_EXIST, BPF_MAP_TYPE_ARRAY,
-        BPF_MAP_TYPE_HASH, BPF_MAP_TYPE_PERCPU_ARRAY,
-        BPF_MAP_TYPE_PERF_EVENT_ARRAY, BPF_MAP_TYPE_QUEUE, BPF_MAP_TYPE_STACK,
-        BPF_MAP_TYPE_STACK_TRACE, BPF_NOEXIST,
-    },
+    BPF_MAP_TYPE_PERCPU_ARRAY, BPF_MAP_TYPE_PERF_EVENT_ARRAY,
+    BPF_MAP_TYPE_QUEUE, BPF_MAP_TYPE_RINGBUF, BPF_MAP_TYPE_STACK,
+    BPF_MAP_TYPE_STACK_TRACE, BPF_NOEXIST, BPF_RB_AVAIL_DATA, BPF_RB_CONS_POS,
+    BPF_RB_PROD_POS, BPF_RB_RING_SIZE,
 };
 use crate::linux::errno::EINVAL;
-use crate::utils::{to_result, NoRef, Result};
+use crate::utils::{
+    to_result, NoRef, PerfEventMaskedCPU, Result, StreamableProgram,
+};
 
 /// Rex equivalent to be used for map APIs in place of the `struct bpf_map`.
 /// The key and the value type are encoded as generics types `K` and `V`.

rex/src/task_struct.rs

@@ -1,5 +1,4 @@
 use crate::bindings::linux::kernel::task_struct;
-use crate::bindings::uapi::linux::errno::EINVAL;
 use crate::per_cpu::{current_task, this_cpu_read};
 use crate::pt_regs::PtRegs;
 use core::ffi::{self, CStr};

rex/src/tracepoint/binding.rs

@@ -45,7 +45,7 @@ pub struct SyscallsEnterDupCtx {
 
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
-pub struct RawSyscallsEnterArgs {
+pub struct RawSyscallsEnterCtx {
     pub unused: u64,
     pub id: i64,
     pub args: [u64; 6],

rex/src/tracepoint/tp_impl.rs

@@ -1,11 +1,50 @@
+use crate::base_helper::termination_check;
 use crate::bindings::uapi::linux::bpf::{
     bpf_map_type, BPF_PROG_TYPE_TRACEPOINT,
 };
 use crate::ffi;
+use crate::map::RexPerfEventArray;
 use crate::prog_type::rex_prog;
 use crate::task_struct::TaskStruct;
+use crate::utils::{to_result, NoRef, PerfEventMaskedCPU, StreamableProgram};
 use crate::Result;
 
+use super::binding::*;
+
+pub enum TracepointContext {
+    SyscallsEnterOpen(&'static SyscallsEnterOpenCtx),
+    SyscallsEnterOpenat(&'static SyscallsEnterOpenatCtx),
+    SyscallsExitOpen(&'static SyscallsExitOpenCtx),
+    SyscallsExitOpenat(&'static SyscallsExitOpenatCtx),
+    SyscallsEnterDup(&'static SyscallsEnterDupCtx),
+    RawSyscallsEnter(&'static RawSyscallsEnterCtx),
+}
+
+impl TracepointContext {
+    unsafe fn get_ptr(&self) -> *const () {
+        match self {
+            TracepointContext::SyscallsEnterOpen(ctx) => {
+                *ctx as *const SyscallsEnterOpenCtx as *const ()
+            }
+            TracepointContext::SyscallsEnterOpenat(ctx) => {
+                *ctx as *const SyscallsEnterOpenatCtx as *const ()
+            }
+            TracepointContext::SyscallsExitOpen(ctx) => {
+                *ctx as *const SyscallsExitOpenCtx as *const ()
+            }
+            TracepointContext::SyscallsExitOpenat(ctx) => {
+                *ctx as *const SyscallsExitOpenatCtx as *const ()
+            }
+            TracepointContext::SyscallsEnterDup(ctx) => {
+                *ctx as *const SyscallsEnterDupCtx as *const ()
+            }
+            TracepointContext::RawSyscallsEnter(ctx) => {
+                *ctx as *const RawSyscallsEnterCtx as *const ()
+            }
+        }
+    }
+}
+
 /// First 3 fields should always be rtti, prog_fn, and name
 ///
 /// rtti should be u64, therefore after compiling the
@@ -50,30 +89,7 @@ impl rex_prog for tracepoint {
 }
 
 impl StreamableProgram for tracepoint {
-    type Context = tp_ctx;
-    fn output_event<T: Copy + NoRef>(
-        &self,
-        ctx: &Self::Context,
-        map: &'static RexPerfEventArray<T>,
-        data: &T,
-        cpu: PerfEventMaskedCPU,
-    ) -> Result {
-        let map_kptr = unsafe { core::ptr::read_volatile(&map.kptr) };
-        let ctx_ptr = unsafe { ctx.get_ptr() };
-        termination_check!(unsafe {
-            to_result!(ffi::bpf_perf_event_output_tp(
-                ctx_ptr,
-                map_kptr,
-                cpu.masked_cpu,
-                data as *const T as *const (),
-                mem::size_of::<T>() as u64,
-            ))
-        })
-    }
-}
-
-impl StreamableProgram for tracepoint {
-    type Context = tp_ctx;
+    type Context = TracepointContext;
     fn output_event<T: Copy + NoRef>(
         &self,
         ctx: &Self::Context,
@@ -89,7 +105,7 @@ impl StreamableProgram for tracepoint {
                 map_kptr,
                 cpu.masked_cpu,
                 data as *const T as *const (),
-                mem::size_of::<T>() as u64,
+                core::mem::size_of::<T>() as u64,
             ))
         })
     }

rex/src/utils.rs

@@ -1,7 +1,6 @@
 use crate::bindings::uapi::linux::bpf::{BPF_F_CURRENT_CPU, BPF_F_INDEX_MASK};
 use crate::bindings::uapi::linux::errno::EINVAL;
 use crate::map::RexPerfEventArray;
-use crate::tracepoint::{tp_ctx, tracepoint};
 use core::ffi::{c_int, c_uchar, CStr};
 use core::mem;
 use core::ops::{Deref, DerefMut, Drop};

samples/harpoon/Cargo.lock

@@ -1,5 +1,5 @@
 [package]
-name = "hello"
+name = "harpoon"
 version = "0.1.0"
 edition = "2024"
 

samples/harpoon/Cargo.toml

@@ -13,8 +13,8 @@ env.set('LINUX_OBJ', kbuild_dir)
 env.set('LINUX_SRC', join_paths(meson.project_source_root(), './linux'))
 env.set('CARGO_TARGET_DIR', join_paths(build_dir, 'target'))
 
-hello_clippy = custom_target(
-  'hello-clippy',
+harpoon_clippy = custom_target(
+  'harpoon-clippy',
   output: ['target'],
   command: [
     cargo_wrapper, rust_bin, '-Z',
@@ -27,9 +27,9 @@ hello_clippy = custom_target(
   build_by_default: true
 )
 
-hello_build = custom_target(
-  'hello-build',
-  output: ['hello'],
+harpoon_build = custom_target(
+  'harpoon-build',
+  output: ['harpoon'],
   command: [
     cargo_wrapper, rust_bin, '-Z',
     'unstable-options',
@@ -43,15 +43,15 @@ hello_build = custom_target(
   build_by_default: true
 )
 
-hello_loader = executable(
+harpoon_loader = executable(
   'loader',
   'loader.c',
   build_by_default: true,
   dependencies: [librex_dep, libbpf_dep, kernel_dep],
   pie: true
 )
 
-hello_trigger = executable(
+harpoon_trigger = executable(
   'event-trigger',
   'event-trigger.c',
   build_by_default: true,
@@ -76,7 +76,7 @@ sanity_test_env.set('Q_SCRIPT',
 )
 sanity_test_env.set('KERNEL_PATH', kbuild_dir)
 
-test('hello_test', 
+test('harpoon_test', 
   python3_bin, 
   args: [sanity_test_scripts],
   env: sanity_test_env,

samples/harpoon/meson.build

@@ -3,20 +3,20 @@
 
 extern crate rex;
 
+use core::ffi::CStr;
 use rex::Result;
-use rex::rex_printk;
-use rex::{rex_tracepoint, rex_map, rex_uprobe};
-use rex::tracepoint::*;
 use rex::kprobe::kprobe;
 use rex::map::{RexArrayMap, RexHashMap, RexPerfEventArray};
 use rex::pt_regs::PtRegs;
+use rex::rex_printk;
+use rex::tracepoint::*;
 use rex::utils::PerfEventMaskedCPU;
-use core::ffi::CStr;
+use rex::{rex_map, rex_tracepoint, rex_uprobe};
 
 #[repr(C)]
 #[derive(Clone, Copy)]
 struct Config {
-    pub values: [u8, 25],
+    pub values: [u8; 25],
 }
 
 #[repr(C)]
@@ -29,7 +29,7 @@ struct SyscallData {
 #[derive(Clone, Copy)]
 struct Tracing {
     status: u32,
-};
+}
 
 #[rex_map]
 static CONFIG_MAP: RexArrayMap<Config> = RexArrayMap::new(1, 0);
@@ -56,16 +56,16 @@ fn exit_function(obj: &kprobe, ctx: &mut PtRegs) -> Result {
     rex_printk!("Exit function.\n");
 }
 
-#[rex_tracepoint(name = "raw_syscalls/sys_enter", tp_type = "RawSyscallsEnter")]
-fn rex_prog1(obj: &tracepoint, ctx: tp_ctx) -> Result {
-    let mut data = SyscallData::new();
+#[rex_tracepoint]
+fn rex_prog1(obj: &tracepoint, ctx: &'static RawSyscallsEnterCtx) -> Result {
+    let mut data = SyscallData { id: 0 };
     let key_config = 0;
     let key_trace = 0;
 
     let Some(tc) = TRACING_STATUS.get_mut(&key_trace) else {
         rex_printk!("Error getting tracing status.\n");
         return Err(1);
-    }
+    };
 
     if tc.status == 1 {
         rex_printk!("Tracing is not active.\n");
@@ -75,35 +75,37 @@ fn rex_prog1(obj: &tracepoint, ctx: tp_ctx) -> Result {
     let Some(task) = obj.bpf_get_current_task() else {
         rex_printk!("Unable to get current task.\n");
         return Err(1);
-    }
+    };
 
     let Ok(command) = task.get_comm() else {
         rex_printk!("Unable to read current program name.\n");
         return Err(1);
-    }
+    };
 
     let Some(input_command_raw) = CONFIG_MAP.get_mut(&key_config) else {
         rex_printk!("Unable to get config.\n");
         return Err(1);
-    }
+    };
 
-    let Ok(input_command) = CStr::from_bytes_until_nul(input_command_raw.values) else {
+    let Ok(input_command) =
+        CStr::from_bytes_until_nul(input_command_raw.values)
+    else {
         rex_printk!("Unable to read input command.\n");
         return Err(1);
-    }
+    };
 
     if command != input_command {
         return Err(1);
     }
 
-    let id = match ctx {
-        RawSyscallsEnter(args) => args.id,
-        _ => 0,
-    };
-
-    data.id = id;
+    data.id = ctx.id;
 
-    EVENTS.output(obj, ctx, data, PerfEventMaskedCPU::current_cpu());
+    EVENTS.output(
+        obj,
+        TracepointContext::RawSyscallsEnter(ctx),
+        data,
+        PerfEventMaskedCPU::current_cpu(),
+    );
 
     rex_printk!("Sending syscall id {}.\n", id);
 

samples/harpoon/src/main.rs

@@ -8,7 +8,7 @@ use rex::linux::perf_event::PERF_MAX_STACK_DEPTH;
 use rex::map::*;
 use rex::perf_event::*;
 use rex::utils::copy_cstr_to_array;
-use rex::{rex_map, rex_perf_event, rex_printk, Result};
+use rex::{Result, rex_map, rex_perf_event, rex_printk};
 
 pub const TASK_COMM_LEN: usize = 16;