Add support for stack argument

bpf/aya-bpf/src/args.rs

@@ -1,3 +1,5 @@
+use aya_bpf_cty::c_ulonglong;
+
 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<T: FromPtRegs>(&self, n: usize) -> Option<T> {
+        T::from_stack_argument(unsafe { &*self.regs }, n)
+    }
+
     /// Returns the value of the register used to pass the return value.
     pub fn ret<T: FromPtRegs>(&self) -> Option<T> {
         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<Self>;
 
+    /// 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<Self>;
+
     /// Coerces a `T` from the return value of a pt_regs context.
     fn from_retval(ctx: &pt_regs) -> Option<Self>;
 }
@@ -115,6 +126,15 @@ impl<T> FromPtRegs for *const T {
         }
     }
 
+    fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        unsafe {
+            let addr: c_ulonglong = ctx.rsp + 8 * (n + 1) as c_ulonglong;
+            bpf_probe_read(addr as *const T)
+                .map(|v| &v as *const _)
+                .ok()
+        }
+    }
+
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.rax).map(|v| v as *const _).ok() }
     }
@@ -130,6 +150,15 @@ 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;
+            bpf_probe_read(addr as *const T)
+                .map(|v| &v as *const _)
+                .ok()
+        }
+    }
+
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.uregs[0]).map(|v| v as *const _).ok() }
     }
@@ -145,6 +174,15 @@ impl<T> FromPtRegs for *const T {
         }
     }
 
+    fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        unsafe {
+            let addr: c_ulonglong = ctx.sp + 8 * (n + 1) as c_ulonglong;
+            bpf_probe_read(addr as *const T)
+                .map(|v| &v as *const _)
+                .ok()
+        }
+    }
+
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.regs[0]).map(|v| v as *const _).ok() }
     }
@@ -185,6 +223,15 @@ impl<T> FromPtRegs for *mut T {
         }
     }
 
+    fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        unsafe {
+            let addr: c_ulonglong = ctx.rsp + 8 * (n + 1) as c_ulonglong;
+            bpf_probe_read(addr as *mut T)
+                .map(|mut v| &mut v as *mut _)
+                .ok()
+        }
+    }
+
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.rax).map(|v| v as *mut _).ok() }
     }
@@ -200,6 +247,15 @@ impl<T> FromPtRegs for *mut 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;
+            bpf_probe_read(addr as *mut T)
+                .map(|mut v| &mut v as *mut _)
+                .ok()
+        }
+    }
+
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.uregs[0]).map(|v| v as *mut _).ok() }
     }
@@ -215,6 +271,15 @@ impl<T> FromPtRegs for *mut T {
         }
     }
 
+    fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        unsafe {
+            let addr: c_ulonglong = ctx.sp + 8 * (n + 1) as c_ulonglong;
+            bpf_probe_read(addr as *mut T)
+                .map(|mut v| &mut v as *mut _)
+                .ok()
+        }
+    }
+
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         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<Self> {
+                unsafe {
+                    let addr: c_ulonglong = ctx.rsp + 8 * (n + 1) as c_ulonglong;
+                    bpf_probe_read(addr as *const $type)
+                        .map(|v| v as $type)
+                        .ok()
+                }
+            }
+
             fn from_retval(ctx: &pt_regs) -> Option<Self> {
                 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<Self> {
+                unsafe {
+                    let addr: c_ulonglong = ctx.uregs[13] + 8 * (n + 1) as c_ulonglong;
+                    bpf_probe_read(addr as *const $type)
+                        .map(|v| v as $type)
+                        .ok()
+                }
+            }
+
             fn from_retval(ctx: &pt_regs) -> Option<Self> {
                 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<Self> {
+                unsafe {
+                    let addr: c_ulonglong = ctx.sp + 8 * (n + 1) as c_ulonglong;
+                    bpf_probe_read(addr as *const $type)
+                        .map(|v| v as $type)
+                        .ok()
+                }
+            }
+
             fn from_retval(ctx: &pt_regs) -> Option<Self> {
                 Some(ctx.regs[0] as *const $type as _)
             }

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<u32, u32> {
+    ///     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<T: FromPtRegs>(&self, n: usize) -> Option<T> {
+        T::from_stack_argument(unsafe { &*self.regs }, n)
+    }
+
     /// Returns the return value of the probed function.
     ///
     /// # Examples

rustfmt.toml

@@ -1,3 +1,4 @@
+edition = "2021"
 unstable_features = true
 reorder_imports = true
 imports_granularity = "Crate"

test/integration-ebpf/src/stack_argument.rs

@@ -0,0 +1,70 @@
+#![no_std]
+#![no_main]
+
+use aya_bpf::{
+    macros::{map, uprobe},
+    maps::PerfEventArray,
+    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);
+
+#[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<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);
+
+    Ok(0)
+}
+
+#[panic_handler]
+fn panic(_info: &core::panic::PanicInfo) -> ! {
+    unsafe { core::hint::unreachable_unchecked() }
+}

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::<Vec<_>>();
+
+            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);
+}