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 pass 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/Cargo.toml

@@ -51,3 +51,7 @@ path = "src/redirect.rs"
 [[bin]]
 name = "xdp_sec"
 path = "src/xdp_sec.rs"
+
+[[bin]]
+name = "stack_argument"
+path = "src/stack_argument.rs"

test/integration-ebpf/src/stack_argument.rs

@@ -0,0 +1,39 @@
+#![no_std]
+#![no_main]
+
+use aya_bpf::{
+    macros::{map, uprobe},
+    maps::HashMap,
+    programs::ProbeContext,
+};
+
+#[map]
+static ARGS: HashMap<u32, u64> = HashMap::with_max_entries(24, 0);
+
+#[uprobe]
+pub fn test_stack_argument(ctx: ProbeContext) -> i32 {
+    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 _ = 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) -> ! {
+    loop {}
+}

test/integration-test/Cargo.toml

@@ -10,6 +10,7 @@ assert_matches = { workspace = true }
 aya = { workspace = true }
 aya-log = { workspace = true }
 aya-obj = { workspace = true }
+bytes = { workspace = true }
 libc = { workspace = true }
 log = { workspace = true }
 netns-rs = { workspace = true }

test/integration-test/src/lib.rs

@@ -11,6 +11,8 @@ pub const TEXT_64_64_RELOC: &[u8] =
     include_bytes_aligned!(concat!(env!("OUT_DIR"), "/text_64_64_reloc.o"));
 
 pub const LOG: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/log"));
+pub const STACK_ARGUMENT: &[u8] =
+    include_bytes_aligned!(concat!(env!("OUT_DIR"), "/stack_argument"));
 pub const MAP_TEST: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/map_test"));
 pub const NAME_TEST: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/name_test"));
 pub const PASS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/pass"));

test/integration-test/src/tests.rs

@@ -6,4 +6,5 @@ mod log;
 mod rbpf;
 mod relocations;
 mod smoke;
+mod stack_argument;
 mod xdp;

test/integration-test/src/tests/stack_argument.rs

@@ -0,0 +1,52 @@
+use aya::{
+    include_bytes_aligned,
+    maps::{AsyncPerfEventArray, HashMap},
+    programs::UProbe,
+    util::online_cpus,
+    Bpf,
+};
+use aya_log::BpfLogger;
+use bytes::BytesMut;
+use log::warn;
+
+use crate::STACK_ARGUMENT;
+
+#[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,
+    //in x86_64, from arg6, stack_argument would be used
+    a_6: u64,
+    a_7: i64,
+) {
+}
+
+#[tokio::test]
+async fn stack_argument() {
+    event_logger::init();
+    let mut bpf = Bpf::load(crate::STACK_ARGUMENT).unwrap();
+
+    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 args_map: HashMap<_, u32, u64> =
+        HashMap::try_from(bpf.take_map("ARGS").unwrap()).unwrap();
+    trigger_stack_argument(0, 1, 2, 3, 4, 5, 6, 7);
+
+    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 as u64, iter_v.1);
+    }
+}