diff --git a/rex-macros/src/kprobe.rs b/rex-macros/src/kprobe.rs index 9b145181..be3380c9 100644 --- a/rex-macros/src/kprobe.rs +++ b/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, item: ItemFn, @@ -23,7 +44,7 @@ impl KProbe { Ok(KProbe { function, item }) } - pub(crate) fn expand(&self) -> Result { + pub(crate) fn expand(&self, flavor: KprobeFlavor) -> Result { let fn_name = self.item.sig.ident.clone(); let item = &self.item; let function_name = format!("{}", fn_name); @@ -31,9 +52,9 @@ impl KProbe { 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! { diff --git a/rex-macros/src/lib.rs b/rex-macros/src/lib.rs index a32db6d5..aead3cf1 100644 --- a/rex-macros/src/lib.rs +++ b/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), diff --git a/rex/librexstub/lib.c b/rex/librexstub/lib.c index 117f29fe..24d15030 100644 --- a/rex/librexstub/lib.c +++ b/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) diff --git a/rex/src/ffi.rs b/rex/src/ffi.rs index d6d9a133..a528fe48 100644 --- a/rex/src/ffi.rs +++ b/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 diff --git a/rex/src/map.rs b/rex/src/map.rs index f9d941f7..1847ec88 100644 --- a/rex/src/map.rs +++ b/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 Sync for RexMapHandle where pub type RexStackTrace = RexMapHandle; pub type RexPerCPUArrayMap = RexMapHandle; +pub type RexPerfEventArray = + RexMapHandle; pub type RexArrayMap = RexMapHandle; pub type RexHashMap = RexMapHandle; pub type RexStack = RexMapHandle; @@ -114,6 +119,21 @@ where } } +impl RexPerfEventArray +where + V: Copy + NoRef, +{ + pub fn output( + &'static self, + program: &P, + ctx: &P::Context, + data: &V, + cpu: PerfEventMaskedCPU, + ) -> Result { + program.output_event(ctx, self, data, cpu) + } +} + impl RexStack where V: Copy + NoRef, diff --git a/rex/src/task_struct.rs b/rex/src/task_struct.rs index 1e756606..74a3ea8c 100644 --- a/rex/src/task_struct.rs +++ b/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(&self, buf: &mut [i8; N]) -> i32 { - if N == 0 { - return -(EINVAL as i32); - } - - let size = core::cmp::min::(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 = diff --git a/rex/src/tracepoint/tp_impl.rs b/rex/src/tracepoint/tp_impl.rs index 58fd199c..604a1dba 100644 --- a/rex/src/tracepoint/tp_impl.rs +++ b/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 rex_prog for tracepoint { ((self.prog)(self, newctx)).unwrap_or_else(|e| e) as u32 } } + +impl PerfEventStreamer for tracepoint { + type Context = C; + fn output_event( + &self, + ctx: &Self::Context, + map: &'static RexPerfEventArray, + 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::() as u64 + )) + }) + } +} diff --git a/rex/src/utils.rs b/rex/src/utils.rs index 9558ee0e..3f06085e 100644 --- a/rex/src/utils.rs +++ b/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( + &self, + ctx: &Self::Context, + map: &'static RexPerfEventArray, + 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, + } + } +} diff --git a/samples/trace_event/src/main.rs b/samples/trace_event/src/main.rs index b4b93696..8fe3a192 100644 --- a/samples/trace_event/src/main.rs +++ b/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)?;