Merge pull request #1133 from ajwerner/uprobe-cookie

Fixes #1132.

Note that this change does not add support in the public API for kprobes
or tracepoints, but it's a trivial matter of plumbing.

Along the way, the Uprobe::attach API is cleaned up to make the
attachment location more coherent. The logic being: if we're going to be
breaking the API anyway, may as well clean it up a bit.

Furthermore, the aya::sys::bpf_link_attach function is cleaned up by
properly modeling the the union in the final field with a rust enum.

aya/src/bpf.rs

@@ -1034,7 +1034,7 @@ impl Ebpf {
     ///
     /// let program: &mut UProbe = bpf.program_mut("SSL_read").unwrap().try_into()?;
     /// program.load()?;
-    /// program.attach(Some("SSL_read"), 0, "libssl", None)?;
+    /// program.attach("SSL_read", "libssl", None, None)?;
     /// # Ok::<(), aya::EbpfError>(())
     /// ```
     pub fn program_mut(&mut self, name: &str) -> Option<&mut Program> {

aya/src/programs/cgroup_device.rs

@@ -75,7 +75,6 @@ impl CgroupDevice {
                 prog_fd,
                 LinkTarget::Fd(cgroup_fd),
                 BPF_CGROUP_DEVICE,
-                None,
                 mode.into(),
                 None,
             )

aya/src/programs/cgroup_skb.rs

@@ -103,7 +103,6 @@ impl CgroupSkb {
                 prog_fd,
                 LinkTarget::Fd(cgroup_fd),
                 attach_type,
-                None,
                 mode.into(),
                 None,
             )

aya/src/programs/cgroup_sock.rs

@@ -81,7 +81,6 @@ impl CgroupSock {
                 prog_fd,
                 LinkTarget::Fd(cgroup_fd),
                 attach_type,
-                None,
                 mode.into(),
                 None,
             )

aya/src/programs/cgroup_sock_addr.rs

@@ -82,7 +82,6 @@ impl CgroupSockAddr {
                 prog_fd,
                 LinkTarget::Fd(cgroup_fd),
                 attach_type,
-                None,
                 mode.into(),
                 None,
             )

aya/src/programs/cgroup_sockopt.rs

@@ -79,7 +79,6 @@ impl CgroupSockopt {
                 prog_fd,
                 LinkTarget::Fd(cgroup_fd),
                 attach_type,
-                None,
                 mode.into(),
                 None,
             )

aya/src/programs/cgroup_sysctl.rs

@@ -74,7 +74,6 @@ impl CgroupSysctl {
                 prog_fd,
                 LinkTarget::Fd(cgroup_fd),
                 BPF_CGROUP_SYSCTL,
-                None,
                 mode.into(),
                 None,
             )

aya/src/programs/extension.rs

@@ -11,7 +11,7 @@ use crate::{
     programs::{
         define_link_wrapper, load_program, FdLink, FdLinkId, ProgramData, ProgramError, ProgramFd,
     },
-    sys::{self, bpf_link_create, LinkTarget, SyscallError},
+    sys::{self, bpf_link_create, BpfLinkCreateArgs, LinkTarget, SyscallError},
     Btf,
 };
 
@@ -101,9 +101,8 @@ impl Extension {
             prog_fd,
             LinkTarget::Fd(target_fd),
             BPF_CGROUP_INET_INGRESS,
-            Some(btf_id),
             0,
-            None,
+            Some(BpfLinkCreateArgs::TargetBtfId(btf_id)),
         )
         .map_err(|(_, io_error)| SyscallError {
             call: "bpf_link_create",
@@ -139,9 +138,8 @@ impl Extension {
             prog_fd,
             LinkTarget::Fd(target_fd),
             BPF_CGROUP_INET_INGRESS,
-            Some(btf_id),
             0,
-            None,
+            Some(BpfLinkCreateArgs::TargetBtfId(btf_id)),
         )
         .map_err(|(_, io_error)| SyscallError {
             call: "bpf_link_create",

aya/src/programs/iter.rs

@@ -72,11 +72,12 @@ impl Iter {
     pub fn attach(&mut self) -> Result<IterLinkId, ProgramError> {
         let prog_fd = self.fd()?;
         let prog_fd = prog_fd.as_fd();
-        let link_fd = bpf_link_create(prog_fd, LinkTarget::Iter, BPF_TRACE_ITER, None, 0, None)
-            .map_err(|(_, io_error)| SyscallError {
+        let link_fd = bpf_link_create(prog_fd, LinkTarget::Iter, BPF_TRACE_ITER, 0, None).map_err(
+            |(_, io_error)| SyscallError {
                 call: "bpf_link_create",
                 io_error,
-            })?;
+            },
+        )?;
 
         self.data
             .links

aya/src/programs/kprobe.rs

@@ -78,7 +78,14 @@ impl KProbe {
         fn_name: T,
         offset: u64,
     ) -> Result<KProbeLinkId, ProgramError> {
-        attach(&mut self.data, self.kind, fn_name.as_ref(), offset, None)
+        attach(
+            &mut self.data,
+            self.kind,
+            fn_name.as_ref(),
+            offset,
+            None, // pid
+            None, // cookie
+        )
     }
 
     /// Creates a program from a pinned entry on a bpffs.

aya/src/programs/mod.rs

@@ -219,6 +219,10 @@ pub enum ProgramError {
     /// An error occurred while working with IO.
     #[error(transparent)]
     IOError(#[from] io::Error),
+
+    /// Providing an attach cookie is not supported.
+    #[error("providing an attach cookie is not supported")]
+    AttachCookieNotSupported,
 }
 
 /// A [`Program`] file descriptor.

aya/src/programs/perf_attach.rs

@@ -7,7 +7,10 @@ use crate::{
         probe::{detach_debug_fs, ProbeEvent},
         FdLink, Link, ProgramError,
     },
-    sys::{bpf_link_create, perf_event_ioctl, LinkTarget, SysResult, SyscallError},
+    sys::{
+        bpf_link_create, is_bpf_cookie_supported, perf_event_ioctl, BpfLinkCreateArgs, LinkTarget,
+        SysResult, SyscallError,
+    },
     FEATURES, PERF_EVENT_IOC_DISABLE, PERF_EVENT_IOC_ENABLE, PERF_EVENT_IOC_SET_BPF,
 };
 
@@ -73,15 +76,18 @@ impl Link for PerfLink {
 pub(crate) fn perf_attach(
     prog_fd: BorrowedFd<'_>,
     fd: crate::MockableFd,
+    cookie: Option<u64>,
 ) -> Result<PerfLinkInner, ProgramError> {
+    if cookie.is_some() && (!is_bpf_cookie_supported() || !FEATURES.bpf_perf_link()) {
+        return Err(ProgramError::AttachCookieNotSupported);
+    }
     if FEATURES.bpf_perf_link() {
         let link_fd = bpf_link_create(
             prog_fd,
             LinkTarget::Fd(fd.as_fd()),
             BPF_PERF_EVENT,
-            None,
             0,
-            None,
+            cookie.map(|bpf_cookie| BpfLinkCreateArgs::PerfEvent { bpf_cookie }),
         )
         .map_err(|(_, io_error)| SyscallError {
             call: "bpf_link_create",

aya/src/programs/perf_event.rs

@@ -180,7 +180,7 @@ impl PerfEvent {
             io_error,
         })?;
 
-        let link = perf_attach(prog_fd, fd)?;
+        let link = perf_attach(prog_fd, fd, None /* cookie */)?;
         self.data.links.insert(PerfEventLink::new(link))
     }
 }

aya/src/programs/probe.rs

@@ -112,17 +112,21 @@ pub(crate) fn attach<T: Link + From<PerfLinkInner>>(
     fn_name: &OsStr,
     offset: u64,
     pid: Option<pid_t>,
+    cookie: Option<u64>,
 ) -> Result<T::Id, ProgramError> {
     // https://github.com/torvalds/linux/commit/e12f03d7031a977356e3d7b75a68c2185ff8d155
     // Use debugfs to create probe
     let prog_fd = program_data.fd()?;
     let prog_fd = prog_fd.as_fd();
     let link = if KernelVersion::current().unwrap() < KernelVersion::new(4, 17, 0) {
+        if cookie.is_some() {
+            return Err(ProgramError::AttachCookieNotSupported);
+        }
         let (fd, event_alias) = create_as_trace_point(kind, fn_name, offset, pid)?;
         perf_attach_debugfs(prog_fd, fd, ProbeEvent { kind, event_alias })
     } else {
         let fd = create_as_probe(kind, fn_name, offset, pid)?;
-        perf_attach(prog_fd, fd)
+        perf_attach(prog_fd, fd, cookie)
     }?;
     program_data.links.insert(T::from(link))
 }

aya/src/programs/sk_lookup.rs

@@ -65,14 +65,7 @@ impl SkLookup {
         let prog_fd = prog_fd.as_fd();
         let netns_fd = netns.as_fd();
 
-        let link_fd = bpf_link_create(
-            prog_fd,
-            LinkTarget::Fd(netns_fd),
-            BPF_SK_LOOKUP,
-            None,
-            0,
-            None,
-        )
+        let link_fd = bpf_link_create(prog_fd, LinkTarget::Fd(netns_fd), BPF_SK_LOOKUP, 0, None)
             .map_err(|(_, io_error)| SyscallError {
                 call: "bpf_link_create",
                 io_error,

aya/src/programs/sock_ops.rs

@@ -73,7 +73,6 @@ impl SockOps {
                 prog_fd,
                 LinkTarget::Fd(cgroup_fd),
                 attach_type,
-                None,
                 mode.into(),
                 None,
             )

aya/src/programs/tc.rs

@@ -23,7 +23,7 @@ use crate::{
     sys::{
         bpf_link_create, bpf_link_get_info_by_fd, bpf_link_update, bpf_prog_get_fd_by_id,
         netlink_find_filter_with_name, netlink_qdisc_add_clsact, netlink_qdisc_attach,
-        netlink_qdisc_detach, LinkTarget, ProgQueryTarget, SyscallError,
+        netlink_qdisc_detach, BpfLinkCreateArgs, LinkTarget, ProgQueryTarget, SyscallError,
     },
     util::{ifindex_from_ifname, tc_handler_make, KernelVersion},
     VerifierLogLevel,
@@ -297,9 +297,8 @@ impl SchedClassifier {
                     prog_fd,
                     LinkTarget::IfIndex(if_index),
                     attach_type.tcx_attach_type()?,
-                    None,
                     options.flags.bits(),
-                    Some(&options.link_ref),
+                    Some(BpfLinkCreateArgs::Tcx(&options.link_ref)),
                 )
                 .map_err(|(_, io_error)| SyscallError {
                     call: "bpf_mprog_attach",

aya/src/programs/trace_point.rs

@@ -78,7 +78,7 @@ impl TracePoint {
                 io_error,
             })?;
 
-        let link = perf_attach(prog_fd, fd)?;
+        let link = perf_attach(prog_fd, fd, None /* cookie */)?;
         self.data.links.insert(TracePointLink::new(link))
     }
 }

aya/src/programs/uprobe.rs

@@ -48,6 +48,30 @@ pub struct UProbe {
     pub(crate) kind: ProbeKind,
 }
 
+/// The location in the target object file to which the uprobe is to be
+/// attached.
+pub enum UProbeAttachLocation<'a> {
+    /// The location of the target function in the target object file.
+    Symbol(&'a str),
+    /// The location of the target function in the target object file, offset by
+    /// the given number of bytes.
+    SymbolOffset(&'a str, u64),
+    /// The offset in the target object file, in bytes.
+    AbsoluteOffset(u64),
+}
+
+impl<'a> From<&'a str> for UProbeAttachLocation<'a> {
+    fn from(s: &'a str) -> Self {
+        Self::Symbol(s)
+    }
+}
+
+impl From<u64> for UProbeAttachLocation<'static> {
+    fn from(offset: u64) -> Self {
+        Self::AbsoluteOffset(offset)
+    }
+}
+
 impl UProbe {
     /// Loads the program inside the kernel.
     pub fn load(&mut self) -> Result<(), ProgramError> {
@@ -64,37 +88,46 @@ impl UProbe {
     ///
     /// Attaches the uprobe to the function `fn_name` defined in the `target`.
     /// If `offset` is non-zero, it is added to the address of the target
-    /// function. If `pid` is not `None`, the program executes only when the target
-    /// function is executed by the given `pid`.
+    /// function. If `pid` is not `None`, the program executes only when the
+    /// target function is executed by the given `pid`.
     ///
     /// The `target` argument can be an absolute path to a binary or library, or
     /// a library name (eg: `"libc"`).
     ///
-    /// If the program is an `uprobe`, it is attached to the *start* address of the target
-    /// function.  Instead if the program is a `uretprobe`, it is attached to the return address of
-    /// the target function.
+    /// If the program is an `uprobe`, it is attached to the *start* address of
+    /// the target function.  Instead if the program is a `uretprobe`, it is
+    /// attached to the return address of the target function.
     ///
     /// The returned value can be used to detach, see [UProbe::detach].
-    pub fn attach<T: AsRef<Path>>(
+    ///
+    /// The cookie is supported since kernel 5.15, and it is made available to
+    /// the eBPF program via the `bpf_get_attach_cookie()` helper.
+    pub fn attach<'loc, T: AsRef<Path>, Loc: Into<UProbeAttachLocation<'loc>>>(
         &mut self,
-        fn_name: Option<&str>,
-        offset: u64,
+        location: Loc,
         target: T,
         pid: Option<pid_t>,
+        cookie: Option<u64>,
     ) -> Result<UProbeLinkId, ProgramError> {
         let path = resolve_attach_path(target.as_ref(), pid)?;
-
-        let sym_offset = if let Some(fn_name) = fn_name {
-            resolve_symbol(&path, fn_name).map_err(|error| UProbeError::SymbolError {
-                symbol: fn_name.to_string(),
+        let (symbol, offset) = match location.into() {
+            UProbeAttachLocation::Symbol(s) => (Some(s), 0),
+            UProbeAttachLocation::SymbolOffset(s, offset) => (Some(s), offset),
+            UProbeAttachLocation::AbsoluteOffset(offset) => (None, offset),
+        };
+        let offset = if let Some(symbol) = symbol {
+            let symbol_offset =
+                resolve_symbol(&path, symbol).map_err(|error| UProbeError::SymbolError {
+                    symbol: symbol.to_string(),
                     error: Box::new(error),
-            })?
+                })?;
+            symbol_offset + offset
         } else {
-            0
+            offset
         };
 
         let path = path.as_os_str();
-        attach(&mut self.data, self.kind, path, sym_offset + offset, pid)
+        attach(&mut self.data, self.kind, path, offset, pid, cookie)
     }
 
     /// Creates a program from a pinned entry on a bpffs.

aya/src/programs/xdp.rs

@@ -148,7 +148,6 @@ impl Xdp {
                 prog_fd,
                 LinkTarget::IfIndex(if_index),
                 attach_type,
-                None,
                 flags.bits(),
                 None,
             )

aya/src/sys/bpf.rs

@@ -390,14 +390,22 @@ pub(crate) enum LinkTarget<'f> {
     Iter,
 }
 
+// Models https://github.com/torvalds/linux/blob/2144da25/include/uapi/linux/bpf.h#L1724-L1782.
+pub(crate) enum BpfLinkCreateArgs<'a> {
+    TargetBtfId(u32),
+    // since kernel 5.15
+    PerfEvent { bpf_cookie: u64 },
+    // since kernel 6.6
+    Tcx(&'a LinkRef),
+}
+
 // since kernel 5.7
 pub(crate) fn bpf_link_create(
     prog_fd: BorrowedFd<'_>,
     target: LinkTarget<'_>,
     attach_type: bpf_attach_type,
-    btf_id: Option<u32>,
     flags: u32,
-    link_ref: Option<&LinkRef>,
+    args: Option<BpfLinkCreateArgs<'_>>,
 ) -> SysResult<crate::MockableFd> {
     let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
 
@@ -417,31 +425,34 @@ pub(crate) fn bpf_link_create(
         LinkTarget::Iter => {}
     };
     attr.link_create.attach_type = attach_type as u32;
+    attr.link_create.flags = flags;
 
-    if let Some(btf_id) = btf_id {
+    if let Some(args) = args {
+        match args {
+            BpfLinkCreateArgs::TargetBtfId(btf_id) => {
                 attr.link_create.__bindgen_anon_3.target_btf_id = btf_id;
             }
-
-    attr.link_create.flags = flags;
-
-    // since kernel 6.6
-    match link_ref {
-        Some(LinkRef::Fd(fd)) => {
+            BpfLinkCreateArgs::PerfEvent { bpf_cookie } => {
+                attr.link_create.__bindgen_anon_3.perf_event.bpf_cookie = bpf_cookie;
+            }
+            BpfLinkCreateArgs::Tcx(link_ref) => match link_ref {
+                LinkRef::Fd(fd) => {
                     attr.link_create
                         .__bindgen_anon_3
                         .tcx
                         .__bindgen_anon_1
                         .relative_fd = fd.to_owned() as u32;
                 }
-        Some(LinkRef::Id(id)) => {
+                LinkRef::Id(id) => {
                     attr.link_create
                         .__bindgen_anon_3
                         .tcx
                         .__bindgen_anon_1
                         .relative_id = id.to_owned();
                 }
-        None => {}
-    };
+            },
+        }
+    }
 
     // SAFETY: BPF_LINK_CREATE returns a new file descriptor.
     unsafe { fd_sys_bpf(bpf_cmd::BPF_LINK_CREATE, &mut attr) }
@@ -877,12 +888,16 @@ pub(crate) fn is_perf_link_supported() -> bool {
 
     if let Ok(fd) = bpf_prog_load(&mut attr) {
         let fd = fd.as_fd();
-        matches!(
         // Uses an invalid target FD so we get EBADF if supported.
-            bpf_link_create(fd, LinkTarget::IfIndex(u32::MAX), bpf_attach_type::BPF_PERF_EVENT, None, 0, None),
+        let link = bpf_link_create(
+            fd,
+            LinkTarget::IfIndex(u32::MAX),
+            bpf_attach_type::BPF_PERF_EVENT,
+            0,
+            None,
+        );
         // Returns EINVAL if unsupported. EBADF if supported.
-            Err((_, e)) if e.raw_os_error() == Some(libc::EBADF),
-        )
+        matches!(link, Err((_, e)) if e.raw_os_error() == Some(libc::EBADF))
     } else {
         false
     }

test/integration-ebpf/Cargo.toml

@@ -81,3 +81,7 @@ path = "src/two_progs.rs"
 [[bin]]
 name = "xdp_sec"
 path = "src/xdp_sec.rs"
+
+[[bin]]
+name = "uprobe_cookie"
+path = "src/uprobe_cookie.rs"

test/integration-ebpf/src/uprobe_cookie.rs

@@ -0,0 +1,26 @@
+#![no_std]
+#![no_main]
+
+use aya_ebpf::{
+    helpers,
+    macros::{map, uprobe},
+    maps::RingBuf,
+    programs::ProbeContext,
+    EbpfContext,
+};
+
+#[map]
+static RING_BUF: RingBuf = RingBuf::with_byte_size(0, 0);
+
+#[uprobe]
+pub fn uprobe_cookie(ctx: ProbeContext) {
+    let cookie = unsafe { helpers::bpf_get_attach_cookie(ctx.as_ptr()) };
+    let cookie_bytes = cookie.to_le_bytes();
+    let _res = RING_BUF.output(&cookie_bytes, 0);
+}
+
+#[cfg(not(test))]
+#[panic_handler]
+fn panic(_info: &core::panic::PanicInfo) -> ! {
+    loop {}
+}

test/integration-test/src/lib.rs

@@ -31,6 +31,7 @@ pub const TCX: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/tcx"));
 pub const TEST: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/test"));
 pub const TWO_PROGS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/two_progs"));
 pub const XDP_SEC: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/xdp_sec"));
+pub const UPROBE_COOKIE: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/uprobe_cookie"));
 
 #[cfg(test)]
 mod tests;

test/integration-test/src/tests.rs

@@ -12,4 +12,5 @@ mod ring_buf;
 mod smoke;
 mod strncmp;
 mod tcx;
+mod uprobe_cookie;
 mod xdp;

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

@@ -100,7 +100,7 @@ fn load_and_attach_uprobe(prog_name: &str, func_name: &str, bytes: &[u8]) -> Ebp
     let prog: &mut UProbe = bpf.program_mut(prog_name).unwrap().try_into().unwrap();
     prog.load().unwrap();
 
-    prog.attach(Some(func_name), 0, "/proc/self/exe", None)
+    prog.attach(func_name, "/proc/self/exe", None, None)
         .unwrap();
 
     bpf

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

@@ -48,10 +48,10 @@ fn relocation_tests(
     program.load().unwrap();
     program
         .attach(
-            Some("trigger_btf_relocations_program"),
-            0,
+            "trigger_btf_relocations_program",
             "/proc/self/exe",
             None,
+            None,
         )
         .unwrap();
 

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

@@ -42,7 +42,7 @@ fn multiple_btf_maps() {
 
     let prog: &mut UProbe = bpf.program_mut("bpf_prog").unwrap().try_into().unwrap();
     prog.load().unwrap();
-    prog.attach(Some("trigger_bpf_program"), 0, "/proc/self/exe", None)
+    prog.attach("trigger_bpf_program", "/proc/self/exe", None, None)
         .unwrap();
 
     trigger_bpf_program();
@@ -92,7 +92,7 @@ fn pin_lifecycle_multiple_btf_maps() {
 
     let prog: &mut UProbe = bpf.program_mut("bpf_prog").unwrap().try_into().unwrap();
     prog.load().unwrap();
-    prog.attach(Some("trigger_bpf_program"), 0, "/proc/self/exe", None)
+    prog.attach("trigger_bpf_program", "/proc/self/exe", None, None)
         .unwrap();
 
     trigger_bpf_program();
@@ -288,7 +288,7 @@ fn basic_uprobe() {
     prog.load().unwrap();
     assert_loaded("test_uprobe");
     let link = prog
-        .attach(Some("uprobe_function"), 0, "/proc/self/exe", None)
+        .attach("uprobe_function", "/proc/self/exe", None, None)
         .unwrap();
 
     {
@@ -301,7 +301,7 @@ fn basic_uprobe() {
     prog.load().unwrap();
 
     assert_loaded("test_uprobe");
-    prog.attach(Some("uprobe_function"), 0, "/proc/self/exe", None)
+    prog.attach("uprobe_function", "/proc/self/exe", None, None)
         .unwrap();
 
     assert_loaded("test_uprobe");
@@ -548,7 +548,7 @@ fn pin_lifecycle_uprobe() {
     {
         let mut prog = UProbe::from_pin(FIRST_PIN_PATH, aya::programs::ProbeKind::UProbe).unwrap();
         let link_id = prog
-            .attach(Some("uprobe_function"), 0, "/proc/self/exe", None)
+            .attach("uprobe_function", "/proc/self/exe", None, None)
             .unwrap();
         let link = prog.take_link(link_id).unwrap();
         let fd_link: FdLink = link.try_into().unwrap();

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

@@ -63,7 +63,7 @@ async fn log() {
 
     let prog: &mut UProbe = bpf.program_mut("test_log").unwrap().try_into().unwrap();
     prog.load().unwrap();
-    prog.attach(Some("trigger_ebpf_program"), 0, "/proc/self/exe", None)
+    prog.attach("trigger_ebpf_program", "/proc/self/exe", None, None)
         .unwrap();
 
     // Call the function that the uprobe is attached to, so it starts logging.

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

@@ -54,12 +54,7 @@ fn load_and_attach(name: &str, bytes: &[u8]) -> Ebpf {
     let prog: &mut UProbe = bpf.program_mut(name).unwrap().try_into().unwrap();
     prog.load().unwrap();
 
-    prog.attach(
-        Some("trigger_relocations_program"),
-        0,
-        "/proc/self/exe",
-        None,
-    )
+    prog.attach("trigger_relocations_program", "/proc/self/exe", None, None)
         .unwrap();
 
     bpf

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

@@ -57,10 +57,10 @@ impl RingBufTest {
             .unwrap();
         prog.load().unwrap();
         prog.attach(
-            Some("ring_buf_trigger_ebpf_program"),
-            0,
+            "ring_buf_trigger_ebpf_program",
             "/proc/self/exe",
             None,
+            None,
         )
         .unwrap();
 

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

@@ -22,7 +22,7 @@ fn bpf_strncmp() {
             .unwrap();
         prog.load().unwrap();
 
-        prog.attach(Some("trigger_bpf_strncmp"), 0, "/proc/self/exe", None)
+        prog.attach("trigger_bpf_strncmp", "/proc/self/exe", None, None)
             .unwrap();
     }
 

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

@@ -0,0 +1,63 @@
+use aya::{maps::ring_buf::RingBuf, programs::UProbe, EbpfLoader};
+use test_log::test;
+
+#[test]
+fn test_uprobe_cookie() {
+    const RING_BUF_BYTE_SIZE: u32 = 512; // arbitrary, but big enough
+
+    let mut bpf = EbpfLoader::new()
+        .set_max_entries("RING_BUF", RING_BUF_BYTE_SIZE)
+        .load(crate::UPROBE_COOKIE)
+        .unwrap();
+    let ring_buf = bpf.take_map("RING_BUF").unwrap();
+    let mut ring_buf = RingBuf::try_from(ring_buf).unwrap();
+    let prog: &mut UProbe = bpf
+        .program_mut("uprobe_cookie")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog.load().unwrap();
+    const PROG_A: &str = "uprobe_cookie_trigger_ebpf_program_a";
+    const PROG_B: &str = "uprobe_cookie_trigger_ebpf_program_b";
+    let attach = |prog: &mut UProbe, fn_name, cookie| {
+        prog.attach(fn_name, "/proc/self/exe", None, Some(cookie))
+            .unwrap()
+    };
+
+    // Note that the arguments we pass to the functions are meaningless, but we
+    // pass the value we expect to see in the ring buffer from the cookie for
+    // readability.
+    let a = attach(prog, PROG_A, 1);
+    let _b = attach(prog, PROG_B, 2);
+    uprobe_cookie_trigger_ebpf_program_a(1);
+    uprobe_cookie_trigger_ebpf_program_b(2);
+    uprobe_cookie_trigger_ebpf_program_a(1);
+    prog.detach(a).unwrap();
+    let _a = attach(prog, PROG_A, 3);
+    uprobe_cookie_trigger_ebpf_program_a(3);
+    const EXP: &[u64] = &[1, 2, 1, 3];
+
+    let mut seen = Vec::new();
+    while let Some(read) = ring_buf.next() {
+        let read = read.as_ref();
+        match read.try_into() {
+            Ok(read) => seen.push(u64::from_le_bytes(read)),
+            Err(std::array::TryFromSliceError { .. }) => {
+                panic!("invalid ring buffer data: {read:x?}")
+            }
+        }
+    }
+    assert_eq!(seen, EXP);
+}
+
+#[no_mangle]
+#[inline(never)]
+pub extern "C" fn uprobe_cookie_trigger_ebpf_program_a(arg: u64) {
+    std::hint::black_box(arg);
+}
+
+#[no_mangle]
+#[inline(never)]
+pub extern "C" fn uprobe_cookie_trigger_ebpf_program_b(arg: u32) {
+    std::hint::black_box(arg);
+}

xtask/public-api/aya.txt

@@ -6291,6 +6291,36 @@ pub fn aya::programs::trace_point::TracePointLinkId::borrow_mut(&mut self) -> &m
 impl<T> core::convert::From<T> for aya::programs::trace_point::TracePointLinkId
 pub fn aya::programs::trace_point::TracePointLinkId::from(t: T) -> T
 pub mod aya::programs::uprobe
+pub enum aya::programs::uprobe::UProbeAttachLocation<'a>
+pub aya::programs::uprobe::UProbeAttachLocation::AbsoluteOffset(u64)
+pub aya::programs::uprobe::UProbeAttachLocation::Symbol(&'a str)
+pub aya::programs::uprobe::UProbeAttachLocation::SymbolOffset(&'a str, u64)
+impl core::convert::From<u64> for aya::programs::uprobe::UProbeAttachLocation<'static>
+pub fn aya::programs::uprobe::UProbeAttachLocation<'static>::from(offset: u64) -> Self
+impl<'a> core::convert::From<&'a str> for aya::programs::uprobe::UProbeAttachLocation<'a>
+pub fn aya::programs::uprobe::UProbeAttachLocation<'a>::from(s: &'a str) -> Self
+impl<'a> core::marker::Freeze for aya::programs::uprobe::UProbeAttachLocation<'a>
+impl<'a> core::marker::Send for aya::programs::uprobe::UProbeAttachLocation<'a>
+impl<'a> core::marker::Sync for aya::programs::uprobe::UProbeAttachLocation<'a>
+impl<'a> core::marker::Unpin for aya::programs::uprobe::UProbeAttachLocation<'a>
+impl<'a> core::panic::unwind_safe::RefUnwindSafe for aya::programs::uprobe::UProbeAttachLocation<'a>
+impl<'a> core::panic::unwind_safe::UnwindSafe for aya::programs::uprobe::UProbeAttachLocation<'a>
+impl<T, U> core::convert::Into<U> for aya::programs::uprobe::UProbeAttachLocation<'a> where U: core::convert::From<T>
+pub fn aya::programs::uprobe::UProbeAttachLocation<'a>::into(self) -> U
+impl<T, U> core::convert::TryFrom<U> for aya::programs::uprobe::UProbeAttachLocation<'a> where U: core::convert::Into<T>
+pub type aya::programs::uprobe::UProbeAttachLocation<'a>::Error = core::convert::Infallible
+pub fn aya::programs::uprobe::UProbeAttachLocation<'a>::try_from(value: U) -> core::result::Result<T, <T as core::convert::TryFrom<U>>::Error>
+impl<T, U> core::convert::TryInto<U> for aya::programs::uprobe::UProbeAttachLocation<'a> where U: core::convert::TryFrom<T>
+pub type aya::programs::uprobe::UProbeAttachLocation<'a>::Error = <U as core::convert::TryFrom<T>>::Error
+pub fn aya::programs::uprobe::UProbeAttachLocation<'a>::try_into(self) -> core::result::Result<U, <U as core::convert::TryFrom<T>>::Error>
+impl<T> core::any::Any for aya::programs::uprobe::UProbeAttachLocation<'a> where T: 'static + ?core::marker::Sized
+pub fn aya::programs::uprobe::UProbeAttachLocation<'a>::type_id(&self) -> core::any::TypeId
+impl<T> core::borrow::Borrow<T> for aya::programs::uprobe::UProbeAttachLocation<'a> where T: ?core::marker::Sized
+pub fn aya::programs::uprobe::UProbeAttachLocation<'a>::borrow(&self) -> &T
+impl<T> core::borrow::BorrowMut<T> for aya::programs::uprobe::UProbeAttachLocation<'a> where T: ?core::marker::Sized
+pub fn aya::programs::uprobe::UProbeAttachLocation<'a>::borrow_mut(&mut self) -> &mut T
+impl<T> core::convert::From<T> for aya::programs::uprobe::UProbeAttachLocation<'a>
+pub fn aya::programs::uprobe::UProbeAttachLocation<'a>::from(t: T) -> T
 pub enum aya::programs::uprobe::UProbeError
 pub aya::programs::uprobe::UProbeError::FileError
 pub aya::programs::uprobe::UProbeError::FileError::filename: std::path::PathBuf
@@ -6336,7 +6366,7 @@ impl<T> core::convert::From<T> for aya::programs::uprobe::UProbeError
 pub fn aya::programs::uprobe::UProbeError::from(t: T) -> T
 pub struct aya::programs::uprobe::UProbe
 impl aya::programs::uprobe::UProbe
-pub fn aya::programs::uprobe::UProbe::attach<T: core::convert::AsRef<std::path::Path>>(&mut self, fn_name: core::option::Option<&str>, offset: u64, target: T, pid: core::option::Option<libc::unix::pid_t>) -> core::result::Result<aya::programs::uprobe::UProbeLinkId, aya::programs::ProgramError>
+pub fn aya::programs::uprobe::UProbe::attach<'loc, T: core::convert::AsRef<std::path::Path>, Loc: core::convert::Into<aya::programs::uprobe::UProbeAttachLocation<'loc>>>(&mut self, location: Loc, target: T, pid: core::option::Option<libc::unix::pid_t>, cookie: core::option::Option<u64>) -> core::result::Result<aya::programs::uprobe::UProbeLinkId, aya::programs::ProgramError>
 pub fn aya::programs::uprobe::UProbe::from_pin<P: core::convert::AsRef<std::path::Path>>(path: P, kind: aya::programs::ProbeKind) -> core::result::Result<Self, aya::programs::ProgramError>
 pub fn aya::programs::uprobe::UProbe::kind(&self) -> aya::programs::ProbeKind
 pub fn aya::programs::uprobe::UProbe::load(&mut self) -> core::result::Result<(), aya::programs::ProgramError>
@@ -7199,6 +7229,7 @@ pub fn aya::programs::Program::from(t: T) -> T
 pub enum aya::programs::ProgramError
 pub aya::programs::ProgramError::AlreadyAttached
 pub aya::programs::ProgramError::AlreadyLoaded
+pub aya::programs::ProgramError::AttachCookieNotSupported
 pub aya::programs::ProgramError::Btf(aya_obj::btf::btf::BtfError)
 pub aya::programs::ProgramError::ExtensionError(aya::programs::extension::ExtensionError)
 pub aya::programs::ProgramError::IOError(std::io::error::Error)
@@ -8937,7 +8968,7 @@ impl<T> core::convert::From<T> for aya::programs::trace_point::TracePoint
 pub fn aya::programs::trace_point::TracePoint::from(t: T) -> T
 pub struct aya::programs::UProbe
 impl aya::programs::uprobe::UProbe
-pub fn aya::programs::uprobe::UProbe::attach<T: core::convert::AsRef<std::path::Path>>(&mut self, fn_name: core::option::Option<&str>, offset: u64, target: T, pid: core::option::Option<libc::unix::pid_t>) -> core::result::Result<aya::programs::uprobe::UProbeLinkId, aya::programs::ProgramError>
+pub fn aya::programs::uprobe::UProbe::attach<'loc, T: core::convert::AsRef<std::path::Path>, Loc: core::convert::Into<aya::programs::uprobe::UProbeAttachLocation<'loc>>>(&mut self, location: Loc, target: T, pid: core::option::Option<libc::unix::pid_t>, cookie: core::option::Option<u64>) -> core::result::Result<aya::programs::uprobe::UProbeLinkId, aya::programs::ProgramError>
 pub fn aya::programs::uprobe::UProbe::from_pin<P: core::convert::AsRef<std::path::Path>>(path: P, kind: aya::programs::ProbeKind) -> core::result::Result<Self, aya::programs::ProgramError>
 pub fn aya::programs::uprobe::UProbe::kind(&self) -> aya::programs::ProbeKind
 pub fn aya::programs::uprobe::UProbe::load(&mut self) -> core::result::Result<(), aya::programs::ProgramError>