From 9c167c714742c476ed78ba76de2733b9277e2244 Mon Sep 17 00:00:00 2001 From: pdliyan Date: Wed, 30 Aug 2023 15:11:11 +0800 Subject: [PATCH] merge sarg patch. --- bpf/aya-bpf/src/args.rs | 92 ++++++++++++++++++ bpf/aya-bpf/src/programs/probe.rs | 39 ++++++++ rustfmt.toml | 1 + test/integration-ebpf/src/stack_argument.rs | 70 ++++++++++++++ .../src/tests/stack_argument.rs | 95 +++++++++++++++++++ 5 files changed, 297 insertions(+) create mode 100644 test/integration-ebpf/src/stack_argument.rs create mode 100644 test/integration-test/src/tests/stack_argument.rs diff --git a/bpf/aya-bpf/src/args.rs b/bpf/aya-bpf/src/args.rs index 437fe329..93bed935 100644 --- a/bpf/aya-bpf/src/args.rs +++ b/bpf/aya-bpf/src/args.rs @@ -1,3 +1,5 @@ +use aya_bpf_bindings::bindings::__u64; + use crate::{cty::c_void, helpers::bpf_probe_read}; // aarch64 uses user_pt_regs instead of pt_regs @@ -75,6 +77,11 @@ impl PtRegs { T::from_argument(unsafe { &*self.regs }, n) } + /// Returns the value of the stack argument used to parss arg `n`. + pub fn stack_arg(&self, n: usize) -> Option { + T::from_stack_argument(unsafe { &*self.regs }, n) + } + /// Returns the value of the register used to pass the return value. pub fn ret(&self) -> Option { T::from_retval(unsafe { &*self.regs }) @@ -97,6 +104,10 @@ pub trait FromPtRegs: Sized { /// at 0 and increases by 1 for each successive argument. fn from_argument(ctx: &pt_regs, n: usize) -> Option; + /// Coerces a `T` from the `n`th stack argument of a pt_regs context where `n` + /// starts at 0 and increases by 1 for each successive argument. + fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option; + /// Coerces a `T` from the return value of a pt_regs context. fn from_retval(ctx: &pt_regs) -> Option; } @@ -115,6 +126,15 @@ impl FromPtRegs for *const T { } } + fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option { + unsafe { + let addr: __u64 = ctx.rsp + 8 * (n + 1) as __u64; + bpf_probe_read(addr as *const T) + .map(|v| &v as *const _) + .ok() + } + } + fn from_retval(ctx: &pt_regs) -> Option { unsafe { bpf_probe_read(&ctx.rax).map(|v| v as *const _).ok() } } @@ -130,6 +150,15 @@ impl FromPtRegs for *const T { } } + fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option { + unsafe { + let addr: __u64 = &ctx.uregs[13] + 8 * (n + 1) as __u64; + bpf_probe_read(addr as *const T) + .map(|v| &v as *const _) + .ok() + } + } + fn from_retval(ctx: &pt_regs) -> Option { unsafe { bpf_probe_read(&ctx.uregs[0]).map(|v| v as *const _).ok() } } @@ -145,6 +174,15 @@ impl FromPtRegs for *const T { } } + fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option { + unsafe { + let addr: __u64 = ctx.sp + 8 * (n + 1) as __u64; + bpf_probe_read(addr as *const T) + .map(|v| &v as *const _) + .ok() + } + } + fn from_retval(ctx: &pt_regs) -> Option { unsafe { bpf_probe_read(&ctx.regs[0]).map(|v| v as *const _).ok() } } @@ -185,6 +223,15 @@ impl FromPtRegs for *mut T { } } + fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option { + unsafe { + let addr: __u64 = ctx.rsp + 8 * (n + 1) as __u64; + bpf_probe_read(addr as *mut T) + .map(|mut v| &mut v as *mut _) + .ok() + } + } + fn from_retval(ctx: &pt_regs) -> Option { unsafe { bpf_probe_read(&ctx.rax).map(|v| v as *mut _).ok() } } @@ -200,6 +247,15 @@ impl FromPtRegs for *mut T { } } + fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option { + unsafe { + let addr: __u64 = ctx.uregs[13] + 8 * (n + 1) as __u64; + bpf_probe_read(addr as *mut T) + .map(|mut v| &mut v as *mut _) + .ok() + } + } + fn from_retval(ctx: &pt_regs) -> Option { unsafe { bpf_probe_read(&ctx.uregs[0]).map(|v| v as *mut _).ok() } } @@ -215,6 +271,15 @@ impl FromPtRegs for *mut T { } } + fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option { + unsafe { + let addr: __u64 = ctx.sp + 8 * (n + 1) as __u64; + bpf_probe_read(addr as *mut T) + .map(|mut v| &mut v as *mut _) + .ok() + } + } + fn from_retval(ctx: &pt_regs) -> Option { unsafe { bpf_probe_read(&ctx.regs[0]).map(|v| v as *mut _).ok() } } @@ -258,6 +323,15 @@ macro_rules! impl_from_pt_regs { } } + fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option { + unsafe { + let addr: __u64 = ctx.rsp + 8 * (n + 1) as __u64; + bpf_probe_read(addr as *const $type) + .map(|v| v as $type) + .ok() + } + } + fn from_retval(ctx: &pt_regs) -> Option { Some(ctx.rax as *const $type as _) } @@ -273,6 +347,15 @@ macro_rules! impl_from_pt_regs { } } + fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option { + unsafe { + let addr: __u64 = ctx.uregs[13] + 8 * (n + 1) as __u64; + bpf_probe_read(addr as *const $type) + .map(|v| v as $type) + .ok() + } + } + fn from_retval(ctx: &pt_regs) -> Option { Some(ctx.uregs[0] as *const $type as _) } @@ -288,6 +371,15 @@ macro_rules! impl_from_pt_regs { } } + fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option { + unsafe { + let addr: __u64 = ctx.sp + 8 * (n + 1) as __u64; + bpf_probe_read(addr as *const $type) + .map(|v| v as $type) + .ok() + } + } + fn from_retval(ctx: &pt_regs) -> Option { Some(ctx.regs[0] as *const $type as _) } diff --git a/bpf/aya-bpf/src/programs/probe.rs b/bpf/aya-bpf/src/programs/probe.rs index 25694845..82cab5bc 100644 --- a/bpf/aya-bpf/src/programs/probe.rs +++ b/bpf/aya-bpf/src/programs/probe.rs @@ -47,6 +47,45 @@ impl ProbeContext { T::from_argument(unsafe { &*self.regs }, n) } + /// Returns the `n`th stack argument to passed to the probe function, starting from 0. + /// + /// # Examples + /// + /// ```no_run + /// # # for c-function in x86_64 platform like: + /// # # void function_with_many_args(int64 a0, int64 a1, int64 a2, + /// # # int64 a3, int64 a4, int64 a5, int64 a6) + /// # #![allow(non_camel_case_types)] + /// # #![allow(dead_code)] + /// unsafe fn try_print_args(ctx: ProbeContext) -> Result { + /// let a_0: i64 = ctx.arg(0).ok_or(1u32)?; + /// let a_1: i64 = ctx.arg(1).ok_or(1u32)?; + /// let a_2: i64 = ctx.arg(2).ok_or(1u32)?; + /// let a_3: i64 = ctx.arg(3).ok_or(1u32)?; + /// let a_4: i64 = ctx.arg(4).ok_or(1u32)?; + /// let a_5: i64 = ctx.arg(5).ok_or(1u32)?; + /// let a_6: i64 = ctx.stack_arg(0).ok_or(1u32)?; + /// info!( + /// &ctx, + /// "arg 0-6: {}, {}, {}, {}, {}, {}, {}", + /// a_0, + /// a_1, + /// a_2, + /// a_3, + /// a_4, + /// a_5, + /// a_6 + /// ); + /// + /// // Do something with args + /// + /// Ok(0) + /// } + /// ``` + pub fn stack_arg(&self, n: usize) -> Option { + T::from_stack_argument(unsafe { &*self.regs }, n) + } + /// Returns the return value of the probed function. /// /// # Examples diff --git a/rustfmt.toml b/rustfmt.toml index 1cdefeb0..2eeb602b 100644 --- a/rustfmt.toml +++ b/rustfmt.toml @@ -1,3 +1,4 @@ +edition = "2021" unstable_features = true reorder_imports = true imports_granularity = "Crate" diff --git a/test/integration-ebpf/src/stack_argument.rs b/test/integration-ebpf/src/stack_argument.rs new file mode 100644 index 00000000..723847fc --- /dev/null +++ b/test/integration-ebpf/src/stack_argument.rs @@ -0,0 +1,70 @@ +#![no_std] +#![no_main] + +use aya_bpf::{ + macros::{uprobe, map}, + programs::{perf_event, ProbeContext}, maps::PerfEventArray, +}; +use aya_log_ebpf::{debug, info}; + +pub struct Args { + a_0: u64, + a_1: u64, + a_2: u64, + a_3: u64, + a_4: u64, + a_5: u64, + + a_6: u64, + a_7: i64, +} + +impl Args{ + fn new()->Self{ + Self{ + a_0: 0, + a_1: 0, + a_2: 0, + a_3: 0, + a_4: 0, + a_5: 0, + a_6: 0, + a_7: 0, + } + } +} + +#[map] +static EVENTS: PerfEventArray = PerfEventArray::with_max_entries(1024, 0); + +#[uprobe] +pub fn test_stack_argument(ctx: ProbeContext) -> i32 { + debug!(&ctx, "Hello from eBPF!"); + match try_stack_argument(ctx) { + Ok(ret) => ret, + Err(_) => 0, + } +} + +//read argument, and send event +fn try_stack_argument(ctx: ProbeContext) -> Result { + let args = Args::new(); + args.a_0 = ctx.arg(0).ok_or(255)?; + args.a_1 = ctx.arg(1).ok_or(255)?; + args.a_2 = ctx.arg(2).ok_or(255)?; + args.a_3 = ctx.arg(3).ok_or(255)?; + args.a_4 = ctx.arg(4).ok_or(255)?; + args.a_5 = ctx.arg(5).ok_or(255)?; + args.a_6 = ctx.stack_arg(0).ok_or(255)?; + args.a_7 = ctx.stack_arg(1).ok_or(255)?; + + + EVENTS.output(&ctx, &args, 0); + + Ok(0) +} + +#[panic_handler] +fn panic(_info: &core::panic::PanicInfo) -> ! { + unsafe { core::hint::unreachable_unchecked() } +} diff --git a/test/integration-test/src/tests/stack_argument.rs b/test/integration-test/src/tests/stack_argument.rs new file mode 100644 index 00000000..2d721b60 --- /dev/null +++ b/test/integration-test/src/tests/stack_argument.rs @@ -0,0 +1,95 @@ +use aya::{ + include_bytes_aligned, maps::AsyncPerfEventArray, programs::UProbe, util::online_cpus, Bpf, +}; +use integration_test_macros::tokio_integration_test; +use log::warn; +use tokio::{signal, task}; + +pub struct Args { + a_0: u64, + a_1: u64, + a_2: u64, + a_3: u64, + a_4: u64, + a_5: u64, + + a_6: u64, + a_7: i64, +} + +impl Args{ + fn new()->Self{ + Self{ + a_0: 0, + a_1: 0, + a_2: 0, + a_3: 0, + a_4: 0, + a_5: 0, + a_6: 0, + a_7: 0, + } + } +} +#[no_mangle] +#[inline(never)] +pub extern "C" fn trigger_stack_argument( + a_0: u64, + a_1: u64, + a_2: u64, + a_3: u64, + a_4: u64, + a_5: u64, + //from arg6, stack_argument would be used + a_6: u64, + a_7: i64, +) { +} + +#[tokio_integration_test] +async fn stack_argument() { + let bytes = + include_bytes_aligned!("../../../../target/bpfel-unknown-none/release/stack_argument"); + let mut bpf = Bpf::load(bytes).unwrap(); + + if let Err(e) = BpfLogger::init(&mut bpf) { + warn!("failed to initialize eBPF logger: {}", e); + } + let prog: &mut UProbe = bpf + .program_mut("test_stack_argument") + .unwrap() + .try_into() + .unwrap(); + prog.load().unwrap(); + prog.attach(Some("trigger_stack_argument"), 0, "/proc/self/exe", None) + .unwrap(); + let mut perf_array = AsyncPerfEventArray::try_from(bpf.take_map("EVENTS").unwrap())?; + for cpu_id in online_cpus()? { + let mut buf = perf_array.open(cpu_id, None)?; + + task::spawn(async move { + let mut buffers = (0..10) + .map(|_| BytesMut::with_capacity(1024)) + .collect::>(); + + loop { + let events = buf.read_events(&mut buffer).await.unwrap(); + for buf in buffers.iter_mut().task(events.read){ + let ptr = buf.as_ptr() as *const Args; + let data = unsafe{ptr.read_unaligned()}; + assert_eq!(data.a_0, 0); + assert_eq!(data.a_1, 1); + assert_eq!(data.a_2, 2); + assert_eq!(data.a_3, 3); + assert_eq!(data.a_4, 4); + assert_eq!(data.a_5, 5); + assert_eq!(data.a_6, 6); + assert_eq!(data.a_7, 7); + break; + } + } + }); + } + + trigger_stack_argument(0, 1, 2, 3, 4, 5, 6, 7); +}