From 628b7fb0221321aa0cc649cd738c449a7dbb0108 Mon Sep 17 00:00:00 2001 From: Andrew Werner Date: Thu, 9 Jan 2025 19:48:21 -0500 Subject: [PATCH] aya::programs::uprobe: add support for cookies 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 | 2 +- aya/src/programs/cgroup_device.rs | 1 - aya/src/programs/cgroup_skb.rs | 1 - aya/src/programs/cgroup_sock.rs | 1 - aya/src/programs/cgroup_sock_addr.rs | 1 - aya/src/programs/cgroup_sockopt.rs | 1 - aya/src/programs/cgroup_sysctl.rs | 1 - aya/src/programs/extension.rs | 8 +- aya/src/programs/iter.rs | 7 +- aya/src/programs/kprobe.rs | 9 ++- aya/src/programs/mod.rs | 4 + aya/src/programs/perf_attach.rs | 12 ++- aya/src/programs/perf_event.rs | 2 +- aya/src/programs/probe.rs | 6 +- aya/src/programs/sk_lookup.rs | 17 ++--- aya/src/programs/sock_ops.rs | 1 - aya/src/programs/tc.rs | 5 +- aya/src/programs/trace_point.rs | 2 +- aya/src/programs/uprobe.rs | 65 ++++++++++++---- aya/src/programs/xdp.rs | 1 - aya/src/sys/bpf.rs | 75 +++++++++++-------- test/integration-ebpf/Cargo.toml | 4 + test/integration-ebpf/src/uprobe_cookie.rs | 26 +++++++ test/integration-test/src/lib.rs | 1 + test/integration-test/src/tests.rs | 1 + .../src/tests/bpf_probe_read.rs | 2 +- .../src/tests/btf_relocations.rs | 4 +- test/integration-test/src/tests/load.rs | 10 +-- test/integration-test/src/tests/log.rs | 2 +- .../integration-test/src/tests/relocations.rs | 9 +-- test/integration-test/src/tests/ring_buf.rs | 4 +- test/integration-test/src/tests/strncmp.rs | 2 +- .../src/tests/uprobe_cookie.rs | 63 ++++++++++++++++ xtask/public-api/aya.txt | 35 ++++++++- 34 files changed, 279 insertions(+), 106 deletions(-) create mode 100644 test/integration-ebpf/src/uprobe_cookie.rs create mode 100644 test/integration-test/src/tests/uprobe_cookie.rs diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs index 940c8f07..acce6c16 100644 --- a/aya/src/bpf.rs +++ b/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> { diff --git a/aya/src/programs/cgroup_device.rs b/aya/src/programs/cgroup_device.rs index 716deae0..6dec42f0 100644 --- a/aya/src/programs/cgroup_device.rs +++ b/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, ) diff --git a/aya/src/programs/cgroup_skb.rs b/aya/src/programs/cgroup_skb.rs index 3b9aed22..71ed4380 100644 --- a/aya/src/programs/cgroup_skb.rs +++ b/aya/src/programs/cgroup_skb.rs @@ -103,7 +103,6 @@ impl CgroupSkb { prog_fd, LinkTarget::Fd(cgroup_fd), attach_type, - None, mode.into(), None, ) diff --git a/aya/src/programs/cgroup_sock.rs b/aya/src/programs/cgroup_sock.rs index 79e118b6..77ccd48d 100644 --- a/aya/src/programs/cgroup_sock.rs +++ b/aya/src/programs/cgroup_sock.rs @@ -81,7 +81,6 @@ impl CgroupSock { prog_fd, LinkTarget::Fd(cgroup_fd), attach_type, - None, mode.into(), None, ) diff --git a/aya/src/programs/cgroup_sock_addr.rs b/aya/src/programs/cgroup_sock_addr.rs index 2a10fce6..8f361a3e 100644 --- a/aya/src/programs/cgroup_sock_addr.rs +++ b/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, ) diff --git a/aya/src/programs/cgroup_sockopt.rs b/aya/src/programs/cgroup_sockopt.rs index 6832d9fa..6ba6b2be 100644 --- a/aya/src/programs/cgroup_sockopt.rs +++ b/aya/src/programs/cgroup_sockopt.rs @@ -79,7 +79,6 @@ impl CgroupSockopt { prog_fd, LinkTarget::Fd(cgroup_fd), attach_type, - None, mode.into(), None, ) diff --git a/aya/src/programs/cgroup_sysctl.rs b/aya/src/programs/cgroup_sysctl.rs index ad4f595b..be8048c4 100644 --- a/aya/src/programs/cgroup_sysctl.rs +++ b/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, ) diff --git a/aya/src/programs/extension.rs b/aya/src/programs/extension.rs index 703bed5d..20a12be2 100644 --- a/aya/src/programs/extension.rs +++ b/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", diff --git a/aya/src/programs/iter.rs b/aya/src/programs/iter.rs index ab0bd4dc..1fb5d256 100644 --- a/aya/src/programs/iter.rs +++ b/aya/src/programs/iter.rs @@ -72,11 +72,12 @@ impl Iter { pub fn attach(&mut self) -> Result { 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 diff --git a/aya/src/programs/kprobe.rs b/aya/src/programs/kprobe.rs index 9e43983a..0b4ce599 100644 --- a/aya/src/programs/kprobe.rs +++ b/aya/src/programs/kprobe.rs @@ -78,7 +78,14 @@ impl KProbe { fn_name: T, offset: u64, ) -> Result { - 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. diff --git a/aya/src/programs/mod.rs b/aya/src/programs/mod.rs index f9fbeec8..fc0fc7dc 100644 --- a/aya/src/programs/mod.rs +++ b/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. diff --git a/aya/src/programs/perf_attach.rs b/aya/src/programs/perf_attach.rs index 89f68ace..81fef70d 100644 --- a/aya/src/programs/perf_attach.rs +++ b/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, ) -> Result { + 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", diff --git a/aya/src/programs/perf_event.rs b/aya/src/programs/perf_event.rs index 73d44865..73d07526 100644 --- a/aya/src/programs/perf_event.rs +++ b/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)) } } diff --git a/aya/src/programs/probe.rs b/aya/src/programs/probe.rs index 4d737e34..f5b796ca 100644 --- a/aya/src/programs/probe.rs +++ b/aya/src/programs/probe.rs @@ -112,17 +112,21 @@ pub(crate) fn attach>( fn_name: &OsStr, offset: u64, pid: Option, + cookie: Option, ) -> Result { // 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)) } diff --git a/aya/src/programs/sk_lookup.rs b/aya/src/programs/sk_lookup.rs index 8d70ee34..add1b04d 100644 --- a/aya/src/programs/sk_lookup.rs +++ b/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))) diff --git a/aya/src/programs/sock_ops.rs b/aya/src/programs/sock_ops.rs index 9f2028bd..2a01e53b 100644 --- a/aya/src/programs/sock_ops.rs +++ b/aya/src/programs/sock_ops.rs @@ -73,7 +73,6 @@ impl SockOps { prog_fd, LinkTarget::Fd(cgroup_fd), attach_type, - None, mode.into(), None, ) diff --git a/aya/src/programs/tc.rs b/aya/src/programs/tc.rs index 69d1bfd1..b76b0f50 100644 --- a/aya/src/programs/tc.rs +++ b/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", diff --git a/aya/src/programs/trace_point.rs b/aya/src/programs/trace_point.rs index 075c6007..3038752d 100644 --- a/aya/src/programs/trace_point.rs +++ b/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)) } } diff --git a/aya/src/programs/uprobe.rs b/aya/src/programs/uprobe.rs index 07a27a38..21f721ac 100644 --- a/aya/src/programs/uprobe.rs +++ b/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 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>( + /// + /// 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, Loc: Into>>( &mut self, - fn_name: Option<&str>, - offset: u64, + location: Loc, target: T, pid: Option, + cookie: Option, ) -> Result { 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. diff --git a/aya/src/programs/xdp.rs b/aya/src/programs/xdp.rs index 6c087f12..3471e318 100644 --- a/aya/src/programs/xdp.rs +++ b/aya/src/programs/xdp.rs @@ -148,7 +148,6 @@ impl Xdp { prog_fd, LinkTarget::IfIndex(if_index), attach_type, - None, flags.bits(), None, ) diff --git a/aya/src/sys/bpf.rs b/aya/src/sys/bpf.rs index 2502fde5..1c2fadef 100644 --- a/aya/src/sys/bpf.rs +++ b/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, flags: u32, - link_ref: Option<&LinkRef>, + args: Option>, ) -> SysResult { let mut attr = unsafe { mem::zeroed::() }; @@ -417,31 +425,34 @@ pub(crate) fn bpf_link_create( LinkTarget::Iter => {} }; attr.link_create.attach_type = attach_type as u32; - - if let Some(btf_id) = 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)) => { - attr.link_create - .__bindgen_anon_3 - .tcx - .__bindgen_anon_1 - .relative_fd = fd.to_owned() as u32; - } - Some(LinkRef::Id(id)) => { - attr.link_create - .__bindgen_anon_3 - .tcx - .__bindgen_anon_1 - .relative_id = id.to_owned(); + if let Some(args) = args { + match args { + BpfLinkCreateArgs::TargetBtfId(btf_id) => { + attr.link_create.__bindgen_anon_3.target_btf_id = btf_id; + } + 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; + } + 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), - // Returns EINVAL if unsupported. EBADF if supported. - Err((_, e)) if e.raw_os_error() == Some(libc::EBADF), - ) + // Uses an invalid target FD so we get EBADF if supported. + 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. + matches!(link, Err((_, e)) if e.raw_os_error() == Some(libc::EBADF)) } else { false } diff --git a/test/integration-ebpf/Cargo.toml b/test/integration-ebpf/Cargo.toml index bee1b0f6..24780717 100644 --- a/test/integration-ebpf/Cargo.toml +++ b/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" diff --git a/test/integration-ebpf/src/uprobe_cookie.rs b/test/integration-ebpf/src/uprobe_cookie.rs new file mode 100644 index 00000000..c937aee9 --- /dev/null +++ b/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 {} +} diff --git a/test/integration-test/src/lib.rs b/test/integration-test/src/lib.rs index 4f54f036..5dcef22a 100644 --- a/test/integration-test/src/lib.rs +++ b/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; diff --git a/test/integration-test/src/tests.rs b/test/integration-test/src/tests.rs index c278e3fd..9ca83669 100644 --- a/test/integration-test/src/tests.rs +++ b/test/integration-test/src/tests.rs @@ -12,4 +12,5 @@ mod ring_buf; mod smoke; mod strncmp; mod tcx; +mod uprobe_cookie; mod xdp; diff --git a/test/integration-test/src/tests/bpf_probe_read.rs b/test/integration-test/src/tests/bpf_probe_read.rs index 9218a5b1..a20f11dc 100644 --- a/test/integration-test/src/tests/bpf_probe_read.rs +++ b/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 diff --git a/test/integration-test/src/tests/btf_relocations.rs b/test/integration-test/src/tests/btf_relocations.rs index 29616850..b6074ef9 100644 --- a/test/integration-test/src/tests/btf_relocations.rs +++ b/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(); diff --git a/test/integration-test/src/tests/load.rs b/test/integration-test/src/tests/load.rs index f9ba4660..0b4f176e 100644 --- a/test/integration-test/src/tests/load.rs +++ b/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(); diff --git a/test/integration-test/src/tests/log.rs b/test/integration-test/src/tests/log.rs index 9c1b1ac9..3ef1891e 100644 --- a/test/integration-test/src/tests/log.rs +++ b/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. diff --git a/test/integration-test/src/tests/relocations.rs b/test/integration-test/src/tests/relocations.rs index 2dce6a60..2328ae72 100644 --- a/test/integration-test/src/tests/relocations.rs +++ b/test/integration-test/src/tests/relocations.rs @@ -54,13 +54,8 @@ 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, - ) - .unwrap(); + prog.attach("trigger_relocations_program", "/proc/self/exe", None, None) + .unwrap(); bpf } diff --git a/test/integration-test/src/tests/ring_buf.rs b/test/integration-test/src/tests/ring_buf.rs index 2493bdf6..515ed421 100644 --- a/test/integration-test/src/tests/ring_buf.rs +++ b/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(); diff --git a/test/integration-test/src/tests/strncmp.rs b/test/integration-test/src/tests/strncmp.rs index 9ee66152..56cf5ef3 100644 --- a/test/integration-test/src/tests/strncmp.rs +++ b/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(); } diff --git a/test/integration-test/src/tests/uprobe_cookie.rs b/test/integration-test/src/tests/uprobe_cookie.rs new file mode 100644 index 00000000..7f9f70e3 --- /dev/null +++ b/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); +} diff --git a/xtask/public-api/aya.txt b/xtask/public-api/aya.txt index 4a87a3a4..b24048c1 100644 --- a/xtask/public-api/aya.txt +++ b/xtask/public-api/aya.txt @@ -6291,6 +6291,36 @@ pub fn aya::programs::trace_point::TracePointLinkId::borrow_mut(&mut self) -> &m impl core::convert::From 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 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 core::convert::Into for aya::programs::uprobe::UProbeAttachLocation<'a> where U: core::convert::From +pub fn aya::programs::uprobe::UProbeAttachLocation<'a>::into(self) -> U +impl core::convert::TryFrom for aya::programs::uprobe::UProbeAttachLocation<'a> where U: core::convert::Into +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>::Error> +impl core::convert::TryInto for aya::programs::uprobe::UProbeAttachLocation<'a> where U: core::convert::TryFrom +pub type aya::programs::uprobe::UProbeAttachLocation<'a>::Error = >::Error +pub fn aya::programs::uprobe::UProbeAttachLocation<'a>::try_into(self) -> core::result::Result>::Error> +impl 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 core::borrow::Borrow for aya::programs::uprobe::UProbeAttachLocation<'a> where T: ?core::marker::Sized +pub fn aya::programs::uprobe::UProbeAttachLocation<'a>::borrow(&self) -> &T +impl core::borrow::BorrowMut 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 core::convert::From 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 core::convert::From 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>(&mut self, fn_name: core::option::Option<&str>, offset: u64, target: T, pid: core::option::Option) -> core::result::Result +pub fn aya::programs::uprobe::UProbe::attach<'loc, T: core::convert::AsRef, Loc: core::convert::Into>>(&mut self, location: Loc, target: T, pid: core::option::Option, cookie: core::option::Option) -> core::result::Result pub fn aya::programs::uprobe::UProbe::from_pin>(path: P, kind: aya::programs::ProbeKind) -> core::result::Result 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 core::convert::From 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>(&mut self, fn_name: core::option::Option<&str>, offset: u64, target: T, pid: core::option::Option) -> core::result::Result +pub fn aya::programs::uprobe::UProbe::attach<'loc, T: core::convert::AsRef, Loc: core::convert::Into>>(&mut self, location: Loc, target: T, pid: core::option::Option, cookie: core::option::Option) -> core::result::Result pub fn aya::programs::uprobe::UProbe::from_pin>(path: P, kind: aya::programs::ProbeKind) -> core::result::Result 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>