Harpoon port

rex-macros/src/kprobe.rs

@@ -1,9 +1,29 @@
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
+use std::fmt;
 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,15 +43,15 @@ 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()
         };

rex-macros/src/lib.rs

@@ -12,7 +12,7 @@ use quote::quote;
 use std::borrow::Cow;
 use syn::ItemStatic;
 
-use kprobe::KProbe;
+use kprobe::{KProbe, KprobeFlavor};
 use perf_event::PerfEvent;
 use tc::SchedCls;
 use tracepoint::TracePoint;
@@ -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/lib.rs

@@ -36,6 +36,7 @@ mod stub;
 
 extern crate paste;
 
+use crate::bindings::uapi::linux::bpf::BPF_F_CURRENT_CPU;
 use crate::prog_type::rex_prog;
 use core::panic::PanicInfo;
 pub use rex_macros::*;
@@ -68,3 +69,5 @@ 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,4 +1,7 @@
-use crate::utils::{to_result, NoRef, Result};
+use crate::utils::{
+    to_result, NoRef, PerfEventMaskedCPU, Result, StreamableProgram,
+};
+use crate::CURRENT_CPU;
 use crate::{
     base_helper::{
         bpf_map_delete_elem,
@@ -12,7 +15,8 @@ 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_HASH, 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,
@@ -65,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>;
@@ -124,6 +130,21 @@ where
     }
 }
 
+impl<'a, V> RexPerfEventArray<V>
+where
+    V: Copy + NoRef,
+{
+    pub fn output<P: StreamableProgram>(
+        &'static self,
+        program: &P,
+        ctx: &P::Context,
+        data: &V,
+        cpu: PerfEventMaskedCPU,
+    ) -> Result {
+        program.output_event(ctx, self, data, cpu)
+    }
+}
+
 // impl RexRingBuf {
 //     pub const fn new(ms: u32, mf: u32) -> RexRingBuf {
 //         RexRingBuf {

rex/src/stub.rs

@@ -140,6 +140,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/task_struct.rs

@@ -3,6 +3,7 @@ use crate::bindings::uapi::linux::errno::EINVAL;
 use crate::per_cpu::{current_task, this_cpu_read};
 use crate::pt_regs::PtRegs;
 use crate::stub;
+use core::ffi::{self, CStr};
 
 // Bindgen has problem generating these constants
 const TOP_OF_KERNEL_STACK_PADDING: u64 = 0;
@@ -44,27 +45,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 reference to the CString instead
+    pub fn get_comm(&self) -> Result<&CStr, ffi::FromBytesUntilNulError> {
+        // casting from c_char to u8 is sound, see:
+        // https://doc.rust-lang.org/stable/src/core/ffi/c_str.rs.html#276-288
+        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,10 +1,12 @@
 use crate::bindings::uapi::linux::bpf::{
     bpf_map_type, BPF_PROG_TYPE_TRACEPOINT,
 };
-use crate::map::*;
 use crate::prog_type::rex_prog;
+use crate::stub;
 use crate::task_struct::TaskStruct;
-use crate::Result;
+use crate::utils::{NoRef, PerfEventMaskedCPU, StreamableProgram};
+use crate::{base_helper::termination_check, map::*, to_result, Result};
+use core::{mem, ptr::null};
 
 use super::binding::*;
 
@@ -75,3 +77,45 @@ impl rex_prog for tracepoint {
         ((self.prog)(self, newctx)).unwrap_or_else(|e| e) as u32
     }
 }
+
+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) };
+        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,
+                    )
+                }
+            })
+        })
+    }
+}

rex/src/utils.rs

@@ -1,3 +1,6 @@
+use crate::bindings::uapi::linux::bpf::BPF_F_INDEX_MASK;
+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};
@@ -258,3 +261,40 @@ where
         unsafe { AlignedMut::from_val(ptr.read_unaligned(), slice) }
     }
 }
+
+pub enum ProgramContextPair {
+    Tracepoint(),
+}
+
+/// programs that can stream data through a
+/// RexPerfEventArray will implement this trait
+pub trait StreamableProgram {
+    type Context: ?Sized;
+    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
+pub struct PerfEventMaskedCPU {
+    pub(crate) masked_cpu: u64,
+}
+
+impl PerfEventMaskedCPU {
+    pub fn current_cpu() -> Self {
+        PerfEventMaskedCPU {
+            masked_cpu: CURRENT_CPU,
+        }
+    }
+    pub fn any_cpu(cpu: u64) -> Self {
+        PerfEventMaskedCPU {
+            masked_cpu: cpu & BPF_F_INDEX_MASK,
+        }
+    }
+}

samples/harpoon/.cargo/config.toml

@@ -0,0 +1,16 @@
+[build]
+target = "x86_64-unknown-none"
+
+[target.x86_64-unknown-none]
+linker = "ld.mold"
+rustflags = [
+  "-Zthreads=8",
+  "-Cforce-frame-pointers=y",
+  "-Csymbol-mangling-version=v0",
+  "-Ccodegen-units=1",
+  "-Crelocation-model=pie",
+  "-Crelro-level=full",
+]
+
+[unstable]
+build-std = ["core"]

samples/harpoon/.gitignore

@@ -0,0 +1 @@
+target