Skip to content

aya: add multi-uprobe attach support #1417

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
swananan wants to merge 1 commit into
base: main
Choose a base branch
from
Open

aya: add multi-uprobe attach support #1417

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions aya-obj/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,3 +25,4 @@ thiserror = { workspace = true, features = ["std"] }
[dev-dependencies]
assert_matches = { workspace = true }
rbpf = { workspace = true }
test-case = { workspace = true }
150 changes: 63 additions & 87 deletions aya-obj/src/obj.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -231,9 +231,11 @@ pub enum ProgramSection {
KProbe,
UProbe {
sleepable: bool,
multi: bool,
},
URetProbe {
sleepable: bool,
multi: bool,
},
TracePoint,
SocketFilter,
Expand Down Expand Up @@ -305,10 +307,38 @@ impl FromStr for ProgramSection {
Ok(match kind {
"kprobe" => Self::KProbe,
"kretprobe" => Self::KRetProbe,
"uprobe" => Self::UProbe { sleepable: false },
"uprobe.s" => Self::UProbe { sleepable: true },
"uretprobe" => Self::URetProbe { sleepable: false },
"uretprobe.s" => Self::URetProbe { sleepable: true },
"uprobe" => Self::UProbe {
sleepable: false,
multi: false,
},
"uprobe.s" => Self::UProbe {
sleepable: true,
multi: false,
},
"uprobe.multi" => Self::UProbe {
sleepable: false,
multi: true,
},
"uprobe.multi.s" => Self::UProbe {
sleepable: true,
multi: true,
},
"uretprobe" => Self::URetProbe {
sleepable: false,
multi: false,
},
"uretprobe.s" => Self::URetProbe {
sleepable: true,
multi: false,
},
"uretprobe.multi" => Self::URetProbe {
sleepable: false,
multi: true,
},
"uretprobe.multi.s" => Self::URetProbe {
sleepable: true,
multi: true,
},
"xdp" | "xdp.frags" => Self::Xdp {
frags: kind == "xdp.frags",
attach_type: match pieces.next() {
Expand Down Expand Up @@ -1457,6 +1487,7 @@ fn get_func_and_line_info(
#[cfg(test)]
mod tests {
use assert_matches::assert_matches;
use test_case::test_case;

use super::*;
use crate::generated::{bpf_map_type::BPF_MAP_TYPE_BLOOM_FILTER, btf_ext_header};
Expand Down Expand Up @@ -1971,102 +2002,47 @@ mod tests {
);
}

#[test]
fn test_parse_section_uprobe() {
let mut obj = fake_obj();
fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN);

assert_matches!(
obj.parse_section(fake_section(
EbpfSectionKind::Program,
"uprobe/foo",
bytes_of(&fake_ins()),
None
)),
Ok(())
);
assert_matches!(
obj.programs.get("foo"),
Some(Program {
section: ProgramSection::UProbe { .. },
..
})
);
}

#[test]
fn test_parse_section_uprobe_sleepable() {
#[test_case("uprobe/foo", ProgramSection::UProbe { sleepable: false, multi: false }; "uprobe_plain")]
#[test_case("uprobe.s/foo", ProgramSection::UProbe { sleepable: true, multi: false }; "uprobe_sleepable")]
#[test_case("uprobe.multi/foo", ProgramSection::UProbe { sleepable: false, multi: true }; "uprobe_multi")]
#[test_case("uprobe.multi.s/foo", ProgramSection::UProbe { sleepable: true, multi: true }; "uprobe_multi_sleepable")]
#[test_case("uretprobe/foo", ProgramSection::URetProbe { sleepable: false, multi: false }; "uretprobe_plain")]
#[test_case("uretprobe.s/foo", ProgramSection::URetProbe { sleepable: true, multi: false }; "uretprobe_sleepable")]
#[test_case("uretprobe.multi/foo", ProgramSection::URetProbe { sleepable: false, multi: true }; "uretprobe_multi")]
#[test_case("uretprobe.multi.s/foo", ProgramSection::URetProbe { sleepable: true, multi: true }; "uretprobe_multi_sleepable")]
fn test_parse_section_user_probe(section: &str, expected_section: ProgramSection) {
let mut obj = fake_obj();
fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN);

assert_matches!(
obj.parse_section(fake_section(
EbpfSectionKind::Program,
"uprobe.s/foo",
section,
bytes_of(&fake_ins()),
None
)),
Ok(())
);
assert_matches!(
obj.programs.get("foo"),
Some(Program {
section: ProgramSection::UProbe {
sleepable: true,
..
let program = obj.programs.remove("foo").unwrap();
match (program.section, expected_section) {
(
ProgramSection::UProbe {
sleepable: actual_sleepable,
multi: actual_multi,
},
..
})
);
}

#[test]
fn test_parse_section_uretprobe() {
let mut obj = fake_obj();
fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN);

assert_matches!(
obj.parse_section(fake_section(
EbpfSectionKind::Program,
"uretprobe/foo",
bytes_of(&fake_ins()),
None
)),
Ok(())
);
assert_matches!(
obj.programs.get("foo"),
Some(Program {
section: ProgramSection::URetProbe { .. },
..
})
);
}

#[test]
fn test_parse_section_uretprobe_sleepable() {
let mut obj = fake_obj();
fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN);

assert_matches!(
obj.parse_section(fake_section(
EbpfSectionKind::Program,
"uretprobe.s/foo",
bytes_of(&fake_ins()),
None
)),
Ok(())
);
assert_matches!(
obj.programs.get("foo"),
Some(Program {
section: ProgramSection::URetProbe {
sleepable: true,
..
ProgramSection::UProbe { sleepable, multi },
)
| (
ProgramSection::URetProbe {
sleepable: actual_sleepable,
multi: actual_multi,
},
..
})
);
ProgramSection::URetProbe { sleepable, multi },
) => assert_eq!((actual_sleepable, actual_multi), (sleepable, multi)),
(section, expected_section) => {
panic!("unexpected section: {section:?}, expected: {expected_section:?}")
}
}
}

#[test]
Expand Down
28 changes: 22 additions & 6 deletions aya/src/bpf.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use crate::{
CgroupSysctl, Extension, FEntry, FExit, FlowDissector, Iter, KProbe, LircMode2, Lsm,
LsmCgroup, PerfEvent, ProbeKind, Program, ProgramData, ProgramError, RawTracePoint,
SchedClassifier, SkLookup, SkMsg, SkReuseport, SkSkb, SockOps, SocketFilter, TracePoint,
UProbe, Xdp,
UProbe, Xdp, uprobe::AttachMode,
},
sys::{
bpf_load_btf, is_bpf_cookie_supported, is_bpf_global_data_supported,
Expand Down Expand Up @@ -449,8 +449,14 @@ impl<'a> EbpfLoader<'a> {
}
ProgramSection::KRetProbe
| ProgramSection::KProbe
| ProgramSection::UProbe { sleepable: _ }
| ProgramSection::URetProbe { sleepable: _ }
| ProgramSection::UProbe {
sleepable: _,
multi: _,
}
| ProgramSection::URetProbe {
sleepable: _,
multi: _,
}
| ProgramSection::TracePoint
| ProgramSection::SocketFilter
| ProgramSection::Xdp {
Expand Down Expand Up @@ -582,7 +588,7 @@ impl<'a> EbpfLoader<'a> {
data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
kind: ProbeKind::Return,
}),
ProgramSection::UProbe { sleepable } => {
ProgramSection::UProbe { sleepable, multi } => {
let mut data =
ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level);
if *sleepable {
Expand All @@ -591,9 +597,14 @@ impl<'a> EbpfLoader<'a> {
Program::UProbe(UProbe {
data,
kind: ProbeKind::Entry,
attach_mode: if *multi {
AttachMode::Multi
} else {
AttachMode::Single
},
})
}
ProgramSection::URetProbe { sleepable } => {
ProgramSection::URetProbe { sleepable, multi } => {
let mut data =
ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level);
if *sleepable {
Expand All @@ -602,6 +613,11 @@ impl<'a> EbpfLoader<'a> {
Program::UProbe(UProbe {
data,
kind: ProbeKind::Return,
attach_mode: if *multi {
AttachMode::Multi
} else {
AttachMode::Single
},
})
}
ProgramSection::TracePoint => Program::TracePoint(TracePoint {
Expand Down Expand Up @@ -1103,7 +1119,7 @@ impl Ebpf {
///
/// let program: &mut UProbe = bpf.program_mut("SSL_read").unwrap().try_into()?;
/// program.load()?;
/// program.attach("SSL_read", "libssl", UProbeScope::AllProcesses)?;
/// program.attach(["SSL_read"], "libssl", UProbeScope::AllProcesses)?;
/// # Ok::<(), aya::EbpfError>(())
/// ```
pub fn program_mut(&mut self, name: &str) -> Option<&mut Program> {
Expand Down
28 changes: 14 additions & 14 deletions aya/src/programs/iter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ impl Iter {

self.data
.links
.insert(IterLink::new(PerfLinkInner::Fd(FdLink::new(link_fd))))
.insert(IterLink::new(FdLink::new(link_fd).into()))
}
}

Expand All @@ -104,25 +104,25 @@ impl AsFd for IterFd {
}
}

impl_try_into_fdlink!(IterLink, PerfLinkInner);
impl_try_from_fdlink!(IterLink, PerfLinkInner, BPF_LINK_TYPE_ITER);
impl_try_into_fdlink!(IterLink, method = into_fd_link);
impl_try_from_fdlink!(IterLink, BPF_LINK_TYPE_ITER);

define_link_wrapper!(IterLink, IterLinkId, PerfLinkInner, PerfLinkIdInner, Iter);

impl IterLink {
/// Converts [`IterLink`] into a [`File`] that can be used to retrieve the
/// outputs of the iterator program.
pub fn into_file(self) -> Result<File, LinkError> {
if let PerfLinkInner::Fd(fd) = self.into_inner() {
let fd = bpf_create_iter(fd.fd.as_fd()).map_err(|io_error| {
LinkError::SyscallError(SyscallError {
call: "bpf_iter_create",
io_error,
})
})?;
Ok(fd.into_inner().into())
} else {
Err(LinkError::InvalidLink)
}
let fd = self
.into_inner()
.into_fd_link()
.map_err(|_inner| LinkError::InvalidLink)?;
let fd = bpf_create_iter(fd.fd.as_fd()).map_err(|io_error| {
LinkError::SyscallError(SyscallError {
call: "bpf_iter_create",
io_error,
})
})?;
Ok(fd.into_inner().into())
}
}
14 changes: 5 additions & 9 deletions aya/src/programs/kprobe.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,9 @@ use thiserror::Error;
use crate::{
VerifierLogLevel,
programs::{
ProgramData, ProgramError, ProgramType, define_link_wrapper, impl_try_from_fdlink,
impl_try_into_fdlink, load_program_without_attach_type,
perf_attach::{PerfLinkIdInner, PerfLinkInner},
PerfLinkIdInner, PerfLinkInner, ProgramData, ProgramError, ProgramType,
define_link_wrapper, impl_try_from_fdlink, impl_try_into_fdlink,
load_program_without_attach_type,
probe::{Probe, ProbeKind, attach},
},
};
Expand Down Expand Up @@ -143,9 +143,5 @@ pub enum KProbeError {
},
}

impl_try_into_fdlink!(KProbeLink, PerfLinkInner);
impl_try_from_fdlink!(
KProbeLink,
PerfLinkInner,
bpf_link_type::BPF_LINK_TYPE_PERF_EVENT
);
impl_try_into_fdlink!(KProbeLink, method = into_fd_link);
impl_try_from_fdlink!(KProbeLink, bpf_link_type::BPF_LINK_TYPE_PERF_EVENT);
31 changes: 26 additions & 5 deletions aya/src/programs/links.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -579,26 +579,47 @@ macro_rules! impl_try_into_fdlink {
}
}
};
($wrapper:ident, method = $method:ident) => {
impl TryFrom<$wrapper> for $crate::programs::FdLink {
type Error = $crate::programs::LinkError;

fn try_from(value: $wrapper) -> Result<Self, Self::Error> {
match value.into_inner().$method() {
Ok(fd) => Ok(fd),
Err(_) => Err($crate::programs::LinkError::InvalidLink),
}
}
}
};
}

pub(crate) use impl_try_into_fdlink;

macro_rules! impl_try_from_fdlink {
($wrapper:ident, $inner:ident, $link_type:expr) => {
(@impl $wrapper:ident, [$($link_type:expr),+ $(,)?], $fd_link:ident => $inner:expr) => {
impl TryFrom<$crate::programs::FdLink> for $wrapper {
type Error = $crate::programs::LinkError;

fn try_from(fd_link: $crate::programs::FdLink) -> Result<Self, Self::Error> {
fn try_from($fd_link: $crate::programs::FdLink) -> Result<Self, Self::Error> {
use std::os::fd::AsFd as _;

let info = $crate::sys::bpf_link_get_info_by_fd(fd_link.fd.as_fd())?;
if info.type_ == ($link_type as u32) {
return Ok(Self::new($inner::Fd(fd_link)));
let info = $crate::sys::bpf_link_get_info_by_fd($fd_link.fd.as_fd())?;
if [$(($link_type as u32)),+].contains(&info.type_) {
return Ok(Self::new($inner));
}
Err($crate::programs::LinkError::InvalidLink)
}
}
};
($wrapper:ident, $inner:ident, $link_type:expr) => {
impl_try_from_fdlink!(@impl $wrapper, [$link_type], fd_link => $inner::Fd(fd_link));
};
($wrapper:ident, link_types = [$($link_type:expr),+ $(,)?]) => {
impl_try_from_fdlink!(@impl $wrapper, [$($link_type),+], fd_link => fd_link.into());
};
($wrapper:ident, $link_type:expr) => {
impl_try_from_fdlink!(@impl $wrapper, [$link_type], fd_link => fd_link.into());
};
}

pub(crate) use impl_try_from_fdlink;
Expand Down
Loading