aya::programs::uprobe: add support for cookies

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,18 +65,11 @@ 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,
-        )
-        .map_err(|(_, io_error)| SyscallError {
-            call: "bpf_link_create",
-            io_error,
-        })?;
+        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,
+            })?;
         self.data
             .links
             .insert(SkLookupLink::new(FdLink::new(link_fd)))

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(),
-                error: Box::new(error),
-            })?
+        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,
             )