Skip to content

Commit

Permalink
fix stack_argument test.
Browse files Browse the repository at this point in the history
  • Loading branch information
pdliyan committed Sep 26, 2023
1 parent 7060235 commit b3b8d44
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 78 deletions.
4 changes: 2 additions & 2 deletions bpf/aya-bpf/src/args.rs
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ impl PtRegs {
T::from_argument(unsafe { &*self.regs }, n)
}

/// Returns the value of the stack argument used to parss arg `n`.
/// Returns the value of the stack argument used to pass arg `n`.
pub fn stack_arg<T: FromPtRegs>(&self, n: usize) -> Option<T> {
T::from_stack_argument(unsafe { &*self.regs }, n)
}
Expand Down Expand Up @@ -152,7 +152,7 @@ impl<T> FromPtRegs for *const T {

fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
unsafe {
let addr: c_ulonglong = &ctx.uregs[13] + 8 * (n + 1) as c_ulonglong;
let addr: c_ulonglong = ctx.uregs[13] + 8 * (n + 1) as c_ulonglong;
bpf_probe_read(addr as *const T)
.map(|v| &v as *const _)
.ok()
Expand Down
55 changes: 12 additions & 43 deletions test/integration-ebpf/src/stack_argument.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,44 +3,15 @@

use aya_bpf::{
macros::{map, uprobe},
maps::PerfEventArray,
maps::HashMap,
programs::ProbeContext,
};
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<Args> = PerfEventArray::with_max_entries(1024, 0);
static ARGS: HashMap<u32, u64> = HashMap::with_max_entries(24, 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,
Expand All @@ -49,22 +20,20 @@ pub fn test_stack_argument(ctx: ProbeContext) -> i32 {

//read argument, and send event
fn try_stack_argument(ctx: ProbeContext) -> Result<i32, i64> {
let mut 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);
let _ = ARGS.insert(&0, &ctx.arg(0).ok_or(255)?, 0);
let _ = ARGS.insert(&1, &ctx.arg(1).ok_or(255)?, 0);
let _ = ARGS.insert(&2, &ctx.arg(2).ok_or(255)?, 0);
let _ = ARGS.insert(&3, &ctx.arg(3).ok_or(255)?, 0);
let _ = ARGS.insert(&4, &ctx.arg(4).ok_or(255)?, 0);
let _ = ARGS.insert(&5, &ctx.arg(5).ok_or(255)?, 0);
let _ = ARGS.insert(&6, &ctx.stack_arg(0).ok_or(255)?, 0);
let _ = ARGS.insert(&7, &ctx.stack_arg(1).ok_or(255)?, 0);

Ok(0)
}

#[cfg(not(test))]
#[panic_handler]
fn panic(_info: &core::panic::PanicInfo) -> ! {
unsafe { core::hint::unreachable_unchecked() }
loop {}
}
42 changes: 9 additions & 33 deletions test/integration-test/src/tests/stack_argument.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use aya::{
include_bytes_aligned, maps::AsyncPerfEventArray, programs::UProbe, util::online_cpus, Bpf,
include_bytes_aligned, maps::{AsyncPerfEventArray, HashMap}, programs::UProbe, util::online_cpus, Bpf,
};
use aya_log::BpfLogger;
use bytes::BytesMut;
Expand Down Expand Up @@ -29,7 +29,7 @@ pub extern "C" fn trigger_stack_argument(
a_3: u64,
a_4: u64,
a_5: u64,
//from arg6, stack_argument would be used
//in x86_64, from arg6, stack_argument would be used
a_6: u64,
a_7: i64,
) {
Expand All @@ -40,9 +40,6 @@ async fn stack_argument() {
event_logger::init();
let mut bpf = Bpf::load(crate::STACK_ARGUMENT).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()
Expand All @@ -51,34 +48,13 @@ async fn stack_argument() {
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)?;
perf_buffers.push(perf_array.open(cpu_id, None)?);

task::spawn(async move {
let mut buffers = (0..10)
.map(|_| BytesMut::with_capacity(1024))
.collect::<Vec<_>>();
let mut args_map = HashMap::try_from(bpf.take_map("ARGS").unwrap())?;
trigger_stack_argument(0, 1, 2, 3, 4, 5, 6, 7);

loop {
let events = buf.read_events(&mut buffers).await.unwrap();
for buf in buffers.iter_mut().take(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;
}
}
});
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!(args_map.keys().count(), 8);
for iter in args_map.iter(){
let iter_v = iter.unwrap();
assert_eq!(iter_v.0, iter_v.1);
}

trigger_stack_argument(0, 1, 2, 3, 4, 5, 6, 7);
}

0 comments on commit b3b8d44

Please sign in to comment.