rex: cleanup: hide BPF_F_CURRENT_CPU, remove ProgramContextPair, add method to grab pointer from tp_ctx, fix KproveFlavor use, fix TaskStruct::get_comm use

Signed-off-by: MinhPhan8803 <nhatminhvinh8803@gmail.com>

Author: MinhPhan8803

rex-macros/src/kprobe.rs

@@ -53,7 +53,7 @@ impl KProbe {
         let attached_function = if self.function.is_some() {
             format!("rex/{}/{}", flavor, self.function.as_ref().unwrap())
         } else {
-            "rex/kprobe".to_string()
+            format!("rex/{}", flavor)
         };
 
         let function_body_tokens = quote! {

rex/src/lib.rs

@@ -69,5 +69,3 @@ define_prog_entry!(sched_cls);
 
 pub use bindings::uapi::*;
 pub use utils::Result;
-
-pub static CURRENT_CPU: u64 = BPF_F_CURRENT_CPU;

rex/src/map.rs

@@ -1,7 +1,6 @@
 use crate::utils::{
     to_result, NoRef, PerfEventMaskedCPU, Result, StreamableProgram,
 };
-use crate::CURRENT_CPU;
 use crate::{
     base_helper::{
         bpf_map_delete_elem,

rex/src/tracepoint/tp_impl.rs

@@ -21,6 +21,20 @@ pub enum tp_ctx {
     SyscallsExitOpen(&'static SyscallsExitOpenArgs),
 }
 
+impl tp_ctx {
+    unsafe fn get_ptr(&self) -> *const () {
+        match self {
+            tp_ctx::Void => null(),
+            tp_ctx::SyscallsEnterOpen(args) => {
+                *args as *const SyscallsEnterOpenArgs as *const ()
+            }
+            tp_ctx::SyscallsExitOpen(args) => {
+                *args as *const SyscallsExitOpenArgs as *const ()
+            }
+        }
+    }
+}
+
 /// First 3 fields should always be rtti, prog_fn, and name
 ///
 /// rtti should be u64, therefore after compiling the
@@ -88,34 +102,15 @@ impl StreamableProgram for tracepoint {
         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!(match ctx {
-                tp_ctx::Void => stub::bpf_perf_event_output_tp(
-                    null(),
-                    map_kptr,
-                    cpu.masked_cpu,
-                    data as *const T as *const (),
-                    mem::size_of::<T>() as u64,
-                ),
-                tp_ctx::SyscallsEnterOpen(args) => {
-                    stub::bpf_perf_event_output_tp(
-                        *args as *const SyscallsEnterOpenArgs as *const (),
-                        map_kptr,
-                        cpu.masked_cpu,
-                        data as *const T as *const (),
-                        mem::size_of::<T>() as u64,
-                    )
-                }
-                tp_ctx::SyscallsExitOpen(args) => {
-                    stub::bpf_perf_event_output_tp(
-                        *args as *const SyscallsExitOpenArgs as *const (),
-                        map_kptr,
-                        cpu.masked_cpu,
-                        data as *const T as *const (),
-                        mem::size_of::<T>() as u64,
-                    )
-                }
-            })
+            to_result!(stub::bpf_perf_event_output_tp(
+                ctx_ptr,
+                map_kptr,
+                cpu.masked_cpu,
+                data as *const T as *const (),
+                mem::size_of::<T>() as u64,
+            ))
         })
     }
 }

rex/src/utils.rs

@@ -1,6 +1,6 @@
-use crate::bindings::uapi::linux::bpf::BPF_F_INDEX_MASK;
+use crate::bindings::uapi::linux::bpf::{BPF_F_CURRENT_CPU, BPF_F_INDEX_MASK};
+use crate::map::RexPerfEventArray;
 use crate::tracepoint::{tp_ctx, tracepoint};
-use crate::{map::RexPerfEventArray, CURRENT_CPU};
 use core::ffi::{c_int, c_uchar};
 use core::mem;
 use core::ops::{Deref, DerefMut, Drop};
@@ -262,10 +262,6 @@ where
     }
 }
 
-pub enum ProgramContextPair {
-    Tracepoint(),
-}
-
 /// programs that can stream data through a
 /// RexPerfEventArray will implement this trait
 pub trait StreamableProgram {
@@ -289,7 +285,7 @@ pub struct PerfEventMaskedCPU {
 impl PerfEventMaskedCPU {
     pub fn current_cpu() -> Self {
         PerfEventMaskedCPU {
-            masked_cpu: CURRENT_CPU,
+            masked_cpu: BPF_F_CURRENT_CPU,
         }
     }
     pub fn any_cpu(cpu: u64) -> Self {

samples/trace_event/src/main.rs

@@ -8,15 +8,15 @@ use rex::linux::bpf::*;
 use rex::linux::perf_event::PERF_MAX_STACK_DEPTH;
 use rex::map::*;
 use rex::perf_event::*;
-use rex::{Result, rex_map, rex_perf_event};
+use rex::{rex_map, rex_perf_event, Result};
 
 pub const TASK_COMM_LEN: usize = 16;
 
 // What if user does not use repr(C)?
 #[repr(C)]
 #[derive(Copy, Clone)]
 pub struct KeyT {
-    pub comm: [i8; TASK_COMM_LEN],
+    pub comm: [u8; TASK_COMM_LEN],
     pub kernstack: u32,
     pub userstack: u32,
 }
@@ -51,7 +51,12 @@ fn rex_prog1(obj: &perf_event, ctx: &bpf_perf_event_data) -> Result {
 
     obj.bpf_get_current_task()
         .map(|t| {
-            t.get_comm(&mut key.comm);
+            let prog_name =
+                t.get_comm().unwrap_or_default().to_bytes_with_nul();
+            let size =
+                core::cmp::min::<usize>(TASK_COMM_LEN, prog_name.len()) - 1;
+            key.comm[..size].copy_from_slice(&prog_name[..size]);
+            key.comm[size] = 0;
             0u64
         })
         .ok_or(0i32)?;