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-macros/src/tracepoint.rs b/rex-macros/src/tracepoint.rs index 8ac213e7..3f6e419c 100644 --- a/rex-macros/src/tracepoint.rs +++ b/rex-macros/src/tracepoint.rs @@ -64,6 +64,7 @@ impl TracePoint { "SyscallsExitOpenCtx" => "syscalls/sys_exit_open", "SyscallsExitOpenatCtx" => "syscalls/sys_exit_openat", "SyscallsEnterDupCtx" => "syscalls/sys_enter_dup", + "RawSyscallsEnterCtx" => "raw_syscalls/sys_enter", _ => abort_call_site!("Please provide a valid context type. If your needed context isn't supported consider opening a PR!"), }; let attached_name = format!("rex/tracepoint/{}", hook_point_name); 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..e94a70f5 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..e9fa58dc 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, Result, StreamableProgram, +}; /// 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<'a, V> 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..b5455565 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, 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 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 = diff --git a/rex/src/tracepoint/binding.rs b/rex/src/tracepoint/binding.rs index 59e2eb8e..12aeb9fa 100644 --- a/rex/src/tracepoint/binding.rs +++ b/rex/src/tracepoint/binding.rs @@ -42,3 +42,11 @@ pub struct SyscallsEnterDupCtx { pub syscall_nr: i64, pub fildes: u64, } + +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct RawSyscallsEnterCtx { + pub unused: u64, + pub id: i64, + pub args: [u64; 6], +} diff --git a/rex/src/tracepoint/tp_impl.rs b/rex/src/tracepoint/tp_impl.rs index 5feaf71a..d324d1d3 100644 --- a/rex/src/tracepoint/tp_impl.rs +++ b/rex/src/tracepoint/tp_impl.rs @@ -1,10 +1,50 @@ +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, StreamableProgram}; use crate::Result; +use super::binding::*; + +pub enum TracepointContext { + SyscallsEnterOpen(&'static SyscallsEnterOpenCtx), + SyscallsEnterOpenat(&'static SyscallsEnterOpenatCtx), + SyscallsExitOpen(&'static SyscallsExitOpenCtx), + SyscallsExitOpenat(&'static SyscallsExitOpenatCtx), + SyscallsEnterDup(&'static SyscallsEnterDupCtx), + RawSyscallsEnter(&'static RawSyscallsEnterCtx), +} + +impl TracepointContext { + unsafe fn get_ptr(&self) -> *const () { + match self { + TracepointContext::SyscallsEnterOpen(ctx) => { + *ctx as *const SyscallsEnterOpenCtx as *const () + } + TracepointContext::SyscallsEnterOpenat(ctx) => { + *ctx as *const SyscallsEnterOpenatCtx as *const () + } + TracepointContext::SyscallsExitOpen(ctx) => { + *ctx as *const SyscallsExitOpenCtx as *const () + } + TracepointContext::SyscallsExitOpenat(ctx) => { + *ctx as *const SyscallsExitOpenatCtx as *const () + } + TracepointContext::SyscallsEnterDup(ctx) => { + *ctx as *const SyscallsEnterDupCtx as *const () + } + TracepointContext::RawSyscallsEnter(ctx) => { + *ctx as *const RawSyscallsEnterCtx as *const () + } + } + } +} + /// First 3 fields should always be rtti, prog_fn, and name /// /// rtti should be u64, therefore after compiling the @@ -47,3 +87,26 @@ impl rex_prog for tracepoint { ((self.prog)(self, ctx)).unwrap_or_else(|e| e) as u32 } } + +impl StreamableProgram for tracepoint { + type Context = TracepointContext; + 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 = unsafe { ctx.get_ptr() }; + 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 b1200dde..c936a7ed 100644 --- a/rex/src/utils.rs +++ b/rex/src/utils.rs @@ -1,4 +1,7 @@ -use core::ffi::{c_int, c_uchar}; +use crate::bindings::uapi::linux::bpf::{BPF_F_CURRENT_CPU, BPF_F_INDEX_MASK}; +use crate::bindings::uapi::linux::errno::EINVAL; +use crate::map::RexPerfEventArray; +use core::ffi::{c_int, c_uchar, CStr}; use core::mem; use core::ops::{Deref, DerefMut, Drop}; @@ -258,3 +261,54 @@ where unsafe { AlignedMut::from_val(ptr.read_unaligned(), slice) } } } + +/// programs that can stream data through a +/// RexPerfEventArray will implement this trait +pub trait StreamableProgram { + type Context: ?Sized; + 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, + } + } +} + +pub fn copy_cstr_to_array( + src: &CStr, + dst: &mut [u8; N], +) -> Result { + if N == 0 { + return Err(-(EINVAL as i32)); + } + let src_bytes = src.to_bytes_with_nul(); + let size = core::cmp::min::(N, src_bytes.len()) - 1; + if size == 0 { + return Err(-(EINVAL as i32)); + } + dst[..size].copy_from_slice(&src_bytes[..size]); + dst[size] = 0; + Ok(0) +} diff --git a/samples/harpoon/.cargo/config.toml b/samples/harpoon/.cargo/config.toml new file mode 100644 index 00000000..113a4b5a --- /dev/null +++ b/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"] diff --git a/samples/harpoon/.gitignore b/samples/harpoon/.gitignore new file mode 100644 index 00000000..eb5a316c --- /dev/null +++ b/samples/harpoon/.gitignore @@ -0,0 +1 @@ +target diff --git a/samples/harpoon/Cargo.lock b/samples/harpoon/Cargo.lock new file mode 100644 index 00000000..3391185f --- /dev/null +++ b/samples/harpoon/Cargo.lock @@ -0,0 +1,98 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "harpoon" +version = "0.1.0" +dependencies = [ + "rex", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "proc-macro-error" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" +dependencies = [ + "proc-macro-error-attr", + "proc-macro2", + "quote", + "version_check", +] + +[[package]] +name = "proc-macro-error-attr" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" +dependencies = [ + "proc-macro2", + "quote", + "version_check", +] + +[[package]] +name = "proc-macro2" +version = "1.0.94" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31971752e70b8b2686d7e46ec17fb38dad4051d94024c88df49b667caea9c84" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.39" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1f1914ce909e1658d9907913b4b91947430c7d9be598b15a1912935b8c04801" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rex" +version = "0.2.0" +dependencies = [ + "paste", + "rex-macros", +] + +[[package]] +name = "rex-macros" +version = "0.2.0" +dependencies = [ + "proc-macro-error", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "syn" +version = "2.0.99" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e02e925281e18ffd9d640e234264753c43edc62d64b2d4cf898f1bc5e75f3fc2" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "unicode-ident" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" diff --git a/samples/harpoon/Cargo.toml b/samples/harpoon/Cargo.toml new file mode 100644 index 00000000..03c43ded --- /dev/null +++ b/samples/harpoon/Cargo.toml @@ -0,0 +1,31 @@ +[package] +name = "harpoon" +version = "0.1.0" +edition = "2024" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +num_cpus = "1.16.0" + +[dependencies.rex] +path = "../../rex" + +[lints.clippy] +disallowed_methods = "forbid" +disallowed_types = "forbid" + +[lints.rust] +incomplete_features = "forbid" +internal_features = "forbid" +unsafe_code = "forbid" +unstable_features = "forbid" + +[profile.dev] +panic = "abort" +debug = false + +[profile.release] +panic = "abort" +debug = false +lto = true diff --git a/samples/harpoon/clippy.toml b/samples/harpoon/clippy.toml new file mode 100644 index 00000000..1c485c92 --- /dev/null +++ b/samples/harpoon/clippy.toml @@ -0,0 +1,7 @@ +disallowed-methods = [ + "core::mem::forget", +] + +disallowed-types = [ + "core::mem::ManuallyDrop", +] diff --git a/samples/harpoon/event-trigger.c b/samples/harpoon/event-trigger.c new file mode 100644 index 00000000..6923e7a7 --- /dev/null +++ b/samples/harpoon/event-trigger.c @@ -0,0 +1,7 @@ +#include +#include + +int main(void) +{ + return syscall(__NR_dup, 1); +} diff --git a/samples/harpoon/loader.c b/samples/harpoon/loader.c new file mode 100644 index 00000000..3dc38e69 --- /dev/null +++ b/samples/harpoon/loader.c @@ -0,0 +1,54 @@ +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +#define EXE "./target/x86_64-unknown-none/release/hello" + +int main(void) +{ + int trace_pipe_fd; + struct bpf_object *obj; + struct bpf_program *prog; + struct bpf_link *link = NULL; + + obj = rex_obj_get_bpf(rex_obj_load(EXE)); + if (!obj) { + fprintf(stderr, "Object could not be opened\n"); + return 1; + } + + prog = bpf_object__find_program_by_name(obj, "rex_prog1"); + if (!prog) { + fprintf(stderr, "_start not found\n"); + return 1; + } + + link = bpf_program__attach(prog); + if (libbpf_get_error(link)) { + fprintf(stderr, "ERROR: bpf_program__attach failed\n"); + link = NULL; + return 1; + } + + trace_pipe_fd = openat(AT_FDCWD, "/sys/kernel/debug/tracing/trace_pipe", + O_RDONLY); + + for (;;) { + char c; + fflush(stdout); + if (read(trace_pipe_fd, &c, 1) == 1) + putchar(c); + } + + bpf_link__destroy(link); + return 0; +} diff --git a/samples/harpoon/meson.build b/samples/harpoon/meson.build new file mode 100644 index 00000000..93a5d3a0 --- /dev/null +++ b/samples/harpoon/meson.build @@ -0,0 +1,86 @@ +build_dir = run_command( + realpath, + '--relative-to', + meson.current_source_dir(), + meson.current_build_dir(), + capture: true, + check: true +).stdout().strip() + +env = environment() +env.prepend('PATH', rust_bin) +env.set('LINUX_OBJ', kbuild_dir) +env.set('LINUX_SRC', join_paths(meson.project_source_root(), './linux')) +env.set('CARGO_TARGET_DIR', join_paths(build_dir, 'target')) + +harpoon_clippy = custom_target( + 'harpoon-clippy', + output: ['target'], + command: [ + cargo_wrapper, rust_bin, '-Z', + 'unstable-options', + '-C', meson.current_source_dir(), + 'clippy', '-qr' + ], + env: env, + console: false, + build_by_default: true +) + +harpoon_build = custom_target( + 'harpoon-build', + output: ['harpoon'], + command: [ + cargo_wrapper, rust_bin, '-Z', + 'unstable-options', + '-C', meson.current_source_dir(), + 'rustc', '-qr', '--', + '-Cenable_rex' + ], + depends: sample_clippy, + env: env, + console: false, + build_by_default: true +) + +harpoon_loader = executable( + 'loader', + 'loader.c', + build_by_default: true, + dependencies: [librex_dep, libbpf_dep, kernel_dep], + pie: true +) + +harpoon_trigger = executable( + 'event-trigger', + 'event-trigger.c', + build_by_default: true, + dependencies: [kernel_dep], + pie: true +) + +sanity_test = custom_target( + 'sanity_test', + output: ['runtest.py'], + command: [ + 'cp', '@CURRENT_SOURCE_DIR@/tests/runtest.py', '@OUTPUT@', + ] + ) + +runtest_deps += [hello_build, hello_loader, hello_trigger, sanity_test] + +sanity_test_env = environment() +sanity_test_env.set('SAMPLE_PATH', meson.current_build_dir()) +sanity_test_env.set('Q_SCRIPT', + join_paths(meson.project_source_root(), 'scripts/q-script/sanity-test-q') +) +sanity_test_env.set('KERNEL_PATH', kbuild_dir) + +test('harpoon_test', + python3_bin, + args: [sanity_test_scripts], + env: sanity_test_env, + depends: runtest_deps, + is_parallel: false, + workdir: meson.current_build_dir() +) diff --git a/samples/harpoon/rustfmt.toml b/samples/harpoon/rustfmt.toml new file mode 100644 index 00000000..84ecd476 --- /dev/null +++ b/samples/harpoon/rustfmt.toml @@ -0,0 +1,4 @@ +max_width = 80 +binop_separator = "Back" +reorder_impl_items = true +wrap_comments = true diff --git a/samples/harpoon/src/main.rs b/samples/harpoon/src/main.rs new file mode 100644 index 00000000..62aac1e4 --- /dev/null +++ b/samples/harpoon/src/main.rs @@ -0,0 +1,113 @@ +#![no_std] +#![no_main] + +extern crate rex; + +use core::ffi::CStr; +use rex::Result; +use rex::kprobe::kprobe; +use rex::map::{RexArrayMap, RexHashMap, RexPerfEventArray}; +use rex::pt_regs::PtRegs; +use rex::rex_printk; +use rex::tracepoint::*; +use rex::utils::PerfEventMaskedCPU; +use rex::{rex_map, rex_tracepoint, rex_uprobe}; + +#[repr(C)] +#[derive(Clone, Copy)] +struct Config { + pub values: [u8; 25], +} + +#[repr(C)] +#[derive(Clone, Copy, core::Default)] +struct SyscallData { + id: i64, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct Tracing { + status: u32, +} + +#[rex_map] +static CONFIG_MAP: RexArrayMap = RexArrayMap::new(1, 0); + +#[rex_map] +static EVENTS: RexPerfEventArray = RexPerfEventArray::new(64, 0); + +#[rex_map] +static TRACING_STATUS: RexHashMap = RexHashMap::new(1, 0); + +#[rex_uprobe] +fn enter_function(obj: &kprobe, ctx: &mut PtRegs) -> Result { + let tc = Tracing { status: 0 }; + let key = 0; + TRACING_STATUS.insert(&key, &tc); + rex_printk!("Enter function.\n"); +} + +#[rex_uprobe] +fn exit_function(obj: &kprobe, ctx: &mut PtRegs) -> Result { + let tc = Tracing { status: 1 }; + let key = 0; + TRACING_STATUS.insert(&key, &tc); + rex_printk!("Exit function.\n"); +} + +#[rex_tracepoint] +fn rex_prog1(obj: &tracepoint, ctx: &'static RawSyscallsEnterCtx) -> Result { + let mut data = SyscallData { id: 0 }; + let key_config = 0; + let key_trace = 0; + + let Some(tc) = TRACING_STATUS.get_mut(&key_trace) else { + rex_printk!("Error getting tracing status.\n"); + return Err(1); + }; + + if tc.status == 1 { + rex_printk!("Tracing is not active.\n"); + return Err(1); + } + + let Some(task) = obj.bpf_get_current_task() else { + rex_printk!("Unable to get current task.\n"); + return Err(1); + }; + + let Ok(command) = task.get_comm() else { + rex_printk!("Unable to read current program name.\n"); + return Err(1); + }; + + let Some(input_command_raw) = CONFIG_MAP.get_mut(&key_config) else { + rex_printk!("Unable to get config.\n"); + return Err(1); + }; + + let Ok(input_command) = + CStr::from_bytes_until_nul(input_command_raw.values) + else { + rex_printk!("Unable to read input command.\n"); + return Err(1); + }; + + if command != input_command { + return Err(1); + } + + data.id = ctx.id; + + EVENTS.output( + obj, + TracepointContext::RawSyscallsEnter(ctx), + data, + PerfEventMaskedCPU::current_cpu(), + ); + + rex_printk!("Sending syscall id {}.\n", id); + + Ok(0) +} diff --git a/samples/harpoon/tests/runtest.py b/samples/harpoon/tests/runtest.py new file mode 100755 index 00000000..2012b07c --- /dev/null +++ b/samples/harpoon/tests/runtest.py @@ -0,0 +1,97 @@ +#!/bin/python + +import re +import subprocess +from time import sleep + +process = 0 + + +def count_bpf_programs(): + try: + # Run bpftool to list all loaded BPF programs + result = subprocess.run( + "bpftool prog show", + capture_output=True, + shell=True, + text=True, + ) + + # Process the output to count programs + if result.stdout: + # Each program details start on a new line + output = result.stdout.strip().split("\n") + programs = [line for line in output if "name" in line] + return len(programs) + else: + return 0 + except FileNotFoundError: + print("bpftool is not installed or not found in the PATH.") + return 0 + except Exception as e: + print(f"An error occurred: {e}") + return 0 + + +def run_loader(): + global process + process = subprocess.Popen( + ["./loader"], + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + +def trigger_prog(): + try: + subprocess.run("./event-trigger", shell=True) + except subprocess.CalledProcessError: + print("CalledProcessError") + + +def capture_output() -> bool: + try: + global process + trigger_prog() + trigger_prog() + + sleep(2) + process.kill() + std_out, std_err = process.communicate(timeout=7) + re_match = re.findall( + r"bpf_trace_printk: Rust triggered from PID \d+ on CPU .+", std_out, re.M + ) + if len(re_match) == 2: + print("Success") + return True + else: + print("Failed") + return False + + except subprocess.CalledProcessError: + return False + except subprocess.TimeoutExpired: + process.kill() + return False + + +def main(): + + old_prog_num = count_bpf_programs() + run_loader() + count = 0 + while old_prog_num == count_bpf_programs(): + sleep(1) + count += 1 + if count == 5: + break + grade_file = open("auto_grade.txt", "w") + if capture_output(): + grade_file.write("success") + else: + grade_file.write("fail") + + +if __name__ == "__main__": + main() diff --git a/samples/trace_event/src/main.rs b/samples/trace_event/src/main.rs index b4b93696..ced7e050 100644 --- a/samples/trace_event/src/main.rs +++ b/samples/trace_event/src/main.rs @@ -7,6 +7,7 @@ use rex::linux::bpf::*; use rex::linux::perf_event::PERF_MAX_STACK_DEPTH; use rex::map::*; use rex::perf_event::*; +use rex::utils::copy_cstr_to_array; use rex::{Result, rex_map, rex_perf_event, rex_printk}; pub const TASK_COMM_LEN: usize = 16; @@ -15,7 +16,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 +51,8 @@ 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(); + copy_cstr_to_array(prog_name, &mut key.comm); 0u64 }) .ok_or(0i32)?;