rex: add support for PERF_EVENT_ARRAY

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

Author: MinhPhan8803

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/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,
+        }
+    }
+}