Merge pull request #31 from rex-rs/perf_array_map_impl

 - samples/trace_event: use the new TaskStruct::get_comm signature
 - rex/tracepoint: implement PerfEventStreamer trait
 - rex: add support for PERF_EVENT_ARRAY
 - rex-macros: add macro support for uprobe
 - rex/task_struct: update get_comm method

Signed-off-by: Jinghao Jia <[email protected]>

Author: jinghao-jia

rex-macros/src/kprobe.rs

@@ -1,9 +1,30 @@
+use std::fmt;
+
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 use syn::{parse2, ItemFn, Result};
 
 use crate::args::parse_string_args;
 
+#[allow(dead_code)]
+pub enum KprobeFlavor {
+    Kprobe,
+    Kretprobe,
+    Uprobe,
+    Uretprobe,
+}
+
+impl fmt::Display for KprobeFlavor {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            KprobeFlavor::Kprobe => write!(f, "kprobe"),
+            KprobeFlavor::Kretprobe => write!(f, "kretprobe"),
+            KprobeFlavor::Uprobe => write!(f, "uprobe"),
+            KprobeFlavor::Uretprobe => write!(f, "uretprobe"),
+        }
+    }
+}
+
 pub(crate) struct KProbe {
     function: Option<String>,
     item: ItemFn,
@@ -23,17 +44,17 @@ impl KProbe {
         Ok(KProbe { function, item })
     }
 
-    pub(crate) fn expand(&self) -> Result<TokenStream> {
+    pub(crate) fn expand(&self, flavor: KprobeFlavor) -> Result<TokenStream> {
         let fn_name = self.item.sig.ident.clone();
         let item = &self.item;
         let function_name = format!("{}", fn_name);
         let prog_ident =
             format_ident!("PROG_{}", fn_name.to_string().to_uppercase());
 
         let attached_function = if self.function.is_some() {
-            format!("rex/kprobe/{}", self.function.as_ref().unwrap())
+            format!("rex/{}/{}", flavor, self.function.as_ref().unwrap())
         } else {
-            "rex/kprobe".to_string()
+            format!("rex/{}", flavor)
         };
 
         let function_body_tokens = quote! {

rex-macros/src/lib.rs

@@ -8,7 +8,7 @@ mod xdp;
 
 use std::borrow::Cow;
 
-use kprobe::KProbe;
+use kprobe::{KProbe, KprobeFlavor};
 use perf_event::PerfEvent;
 use proc_macro::TokenStream;
 use proc_macro_error::{abort, proc_macro_error};
@@ -47,7 +47,19 @@ pub fn rex_tc(attrs: TokenStream, item: TokenStream) -> TokenStream {
 pub fn rex_kprobe(attrs: TokenStream, item: TokenStream) -> TokenStream {
     match KProbe::parse(attrs.into(), item.into()) {
         Ok(prog) => prog
-            .expand()
+            .expand(KprobeFlavor::Kprobe)
+            .unwrap_or_else(|err| abort!(err.span(), "{}", err))
+            .into(),
+        Err(err) => abort!(err.span(), "{}", err),
+    }
+}
+
+#[proc_macro_error]
+#[proc_macro_attribute]
+pub fn rex_uprobe(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    match KProbe::parse(attrs.into(), item.into()) {
+        Ok(prog) => prog
+            .expand(KprobeFlavor::Uprobe)
             .unwrap_or_else(|err| abort!(err.span(), "{}", err))
             .into(),
         Err(err) => abort!(err.span(), "{}", err),

rex/librexstub/lib.c

@@ -22,6 +22,10 @@ KSYM_FUNC(bpf_spin_unlock)
 KSYM_FUNC(just_return_func)
 KSYM_FUNC(bpf_get_stackid_pe)
 KSYM_FUNC(bpf_perf_prog_read_value)
+KSYM_FUNC(bpf_perf_event_output_tp)
+KSYM_FUNC(bpf_perf_event_read_value)
+KSYM_FUNC(bpf_skb_event_output)
+KSYM_FUNC(bpf_xdp_event_output)
 KSYM_FUNC(bpf_xdp_adjust_head)
 KSYM_FUNC(bpf_xdp_adjust_tail)
 KSYM_FUNC(bpf_clone_redirect)

rex/src/ffi.rs

@@ -132,6 +132,53 @@ unsafe extern "C" {
         size: u32,
     ) -> i64;
 
+    /// `long bpf_perf_event_output_tp(void *tp_buff, struct bpf_map *map, u64
+    /// flags, void *data, u64 size)`
+    pub(crate) fn bpf_perf_event_output_tp(
+        tp_buff: *const (),
+        map: *mut (),
+        flags: u64,
+        data: *const (),
+        size: u64,
+    ) -> i64;
+
+    /// `long bpf_perf_event_read_value(struct bpf_map *map, u64 flags,
+    /// struct bpf_perf_event_value *buf, u32 buf_size)`
+    /// same reason for use of improper_ctypes as bpf_perf_prog_read_value
+    #[allow(improper_ctypes)]
+    pub(crate) fn bpf_perf_event_read_value(
+        map: *mut (),
+        flags: u64,
+        buf: &mut bpf_perf_event_value,
+        buf_size: u32,
+    ) -> i64;
+
+    /// `long bpf_skb_event_output(struct sk_buff *skb, struct bpf_map *map, u64 flags,
+    /// void *meta, u64 meta_size)`
+    /// The compiler complains about some non-FFI safe type, but since the
+    /// kernel is using it fine it should be safe for an FFI call using C ABI
+    #[allow(improper_ctypes)]
+    pub(crate) fn bpf_skb_event_output(
+        skb: *const sk_buff,
+        map: *mut (),
+        flags: u64,
+        meta: *const (),
+        meta_size: u64,
+    ) -> i64;
+
+    /// `long bpf_xdp_event_output(struct xdp_buff *xdp, struct bpf_map *map, u64 flags,
+    /// void *meta, u64 meta_size)`
+    /// The compiler complains about some non-FFI safe type, but since the
+    /// kernel is using it fine it should be safe for an FFI call using C ABI
+    #[allow(improper_ctypes)]
+    pub(crate) fn bpf_xdp_event_output(
+        xdp: *const xdp_buff,
+        map: *mut (),
+        flags: u64,
+        meta: *const (),
+        meta_size: u64,
+    ) -> i64;
+
     /// `long bpf_xdp_adjust_head(struct xdp_buff *xdp, int offset)`
     ///
     /// The compiler complains about some non-FFI safe type, but since the

rex/src/map.rs

@@ -11,12 +11,15 @@ use crate::base_helper::{
 use crate::ffi;
 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,
+    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, PerfEventStreamer, Result,
+};
 
 /// 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`.
@@ -66,6 +69,8 @@ unsafe impl<const MT: bpf_map_type, K, V> Sync for RexMapHandle<MT, K, V> where
 
 pub type RexStackTrace<K, V> = RexMapHandle<BPF_MAP_TYPE_STACK_TRACE, K, V>;
 pub type RexPerCPUArrayMap<V> = RexMapHandle<BPF_MAP_TYPE_PERCPU_ARRAY, u32, V>;
+pub type RexPerfEventArray<V> =
+    RexMapHandle<BPF_MAP_TYPE_PERF_EVENT_ARRAY, u32, V>;
 pub type RexArrayMap<V> = RexMapHandle<BPF_MAP_TYPE_ARRAY, u32, V>;
 pub type RexHashMap<K, V> = RexMapHandle<BPF_MAP_TYPE_HASH, K, V>;
 pub type RexStack<V> = RexMapHandle<BPF_MAP_TYPE_STACK, (), V>;
@@ -114,6 +119,21 @@ where
     }
 }
 
+impl<V> RexPerfEventArray<V>
+where
+    V: Copy + NoRef,
+{
+    pub fn output<P: PerfEventStreamer>(
+        &'static self,
+        program: &P,
+        ctx: &P::Context,
+        data: &V,
+        cpu: PerfEventMaskedCPU,
+    ) -> Result {
+        program.output_event(ctx, self, data, cpu)
+    }
+}
+
 impl<V> RexStack<V>
 where
     V: Copy + NoRef,

rex/src/task_struct.rs

@@ -1,7 +1,7 @@
 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 as core_ffi, CStr};
 
 // Bindgen has problem generating these constants
 const TOP_OF_KERNEL_STACK_PADDING: u64 = 0;
@@ -43,27 +43,19 @@ impl TaskStruct {
         self.kptr.tgid
     }
 
-    // Design decision: the original BPF interface does not have type safety,
-    // since buf is just a buffer. But in Rust we can use const generics to
-    // restrict it to only [u8; N] given that comm is a cstring. This also
-    // automatically achieves size check, since N is a constexpr.
-    pub fn get_comm<const N: usize>(&self, buf: &mut [i8; N]) -> i32 {
-        if N == 0 {
-            return -(EINVAL as i32);
-        }
-
-        let size = core::cmp::min::<usize>(N, self.kptr.comm.len()) - 1;
-        if size == 0 {
-            return -(EINVAL as i32);
-        }
-
-        buf[..size].copy_from_slice(&self.kptr.comm[..size]);
-        buf[size] = 0;
-        0
+    // Design decision: the equivalent BPF helper writes the program name to
+    // a user-provided buffer, here we can take advantage of Rust's ownership by
+    // just providing a &CStr instead
+    pub fn get_comm(&self) -> Result<&CStr, core_ffi::FromBytesUntilNulError> {
+        // casting from c_char to u8 is sound, see:
+        // https://doc.rust-lang.org/src/core/ffi/c_str.rs.html#264
+        let comm_bytes =
+            unsafe { &*(&self.kptr.comm[..] as *const _ as *const [u8]) };
+        CStr::from_bytes_until_nul(comm_bytes)
     }
 
     pub fn get_pt_regs(&self) -> &'static PtRegs {
-        // X86 sepcific
+        // X86 specific
         // stack_top is actually bottom of the kernel stack, it refers to the
         // highest address of the stack pages
         let stack_top =

rex/src/tracepoint/tp_impl.rs

@@ -1,8 +1,12 @@
+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, PerfEventStreamer};
 use crate::Result;
 
 use super::{
@@ -65,3 +69,26 @@ impl<C: TracepointContext + 'static> rex_prog for tracepoint<C> {
         ((self.prog)(self, newctx)).unwrap_or_else(|e| e) as u32
     }
 }
+
+impl<C: TracepointContext + 'static> PerfEventStreamer for tracepoint<C> {
+    type Context = C;
+    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 = ctx as *const C as *const ();
+        termination_check!(unsafe {
+            to_result!(ffi::bpf_perf_event_output_tp(
+                ctx_ptr,
+                map_kptr,
+                cpu.masked_cpu,
+                data as *const T as *const (),
+                core::mem::size_of::<T>() as u64
+            ))
+        })
+    }
+}

rex/src/utils.rs

@@ -1,3 +1,5 @@
+use crate::bindings::uapi::linux::bpf::{BPF_F_CURRENT_CPU, BPF_F_INDEX_MASK};
+use crate::map::RexPerfEventArray;
 use core::ffi::{c_int, c_uchar};
 use core::mem;
 use core::ops::{Deref, DerefMut, Drop};
@@ -286,3 +288,39 @@ macro_rules! read_field {
         )
     }};
 }
+
+// For implementers, see tp_impl.rs for how to implement
+// this trait
+/// Programs that can stream data through a
+/// RexPerfEventArray will implement this trait
+pub trait PerfEventStreamer {
+    type Context;
+    fn output_event<T: Copy + NoRef>(
+        &self,
+        ctx: &Self::Context,
+        map: &'static RexPerfEventArray<T>,
+        data: &T,
+        cpu: PerfEventMaskedCPU,
+    ) -> Result;
+}
+
+/// Newtype for a cpu for perf event output to ensure
+/// type safety since the cpu must be masked with
+/// BPF_F_INDEX_MASK
+#[derive(Debug, Copy, Clone)]
+pub struct PerfEventMaskedCPU {
+    pub(crate) masked_cpu: u64,
+}
+
+impl PerfEventMaskedCPU {
+    pub fn current_cpu() -> Self {
+        PerfEventMaskedCPU {
+            masked_cpu: BPF_F_CURRENT_CPU,
+        }
+    }
+    pub fn any_cpu(cpu: u64) -> Self {
+        PerfEventMaskedCPU {
+            masked_cpu: cpu & BPF_F_INDEX_MASK,
+        }
+    }
+}

samples/trace_event/src/main.rs

@@ -15,7 +15,7 @@ pub const TASK_COMM_LEN: usize = 16;
 #[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,
 }
@@ -50,7 +50,10 @@ 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();
+            key.comm
+                .copy_from_slice(&prog_name.to_bytes()[..TASK_COMM_LEN]);
+            key.comm[TASK_COMM_LEN - 1] = 0;
             0u64
         })
         .ok_or(0i32)?;