Skip to content

Commit

Permalink
feat: Allow conversions to Program from ProgramInfo
Browse files Browse the repository at this point in the history
Allow for a ProgramInfo to be converted into one of the program types
that we support. This allows for a user of Aya access to reattach,
pin or unload a program that was either, previously loaded, or was
loaded by another process.

Signed-off-by: Dave Tucker <[email protected]>
  • Loading branch information
dave-tucker committed Mar 6, 2025
1 parent 65489e1 commit 7e54054
Show file tree
Hide file tree
Showing 5 changed files with 329 additions and 15 deletions.
6 changes: 3 additions & 3 deletions aya/src/bpf.rs
Original file line number Diff line number Diff line change
Expand Up @@ -609,15 +609,15 @@ impl<'a> EbpfLoader<'a> {
}
ProgramSection::CgroupSkb => Program::CgroupSkb(CgroupSkb {
data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
expected_attach_type: None,
attach_type: None,
}),
ProgramSection::CgroupSkbIngress => Program::CgroupSkb(CgroupSkb {
data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
expected_attach_type: Some(CgroupSkbAttachType::Ingress),
attach_type: Some(CgroupSkbAttachType::Ingress),
}),
ProgramSection::CgroupSkbEgress => Program::CgroupSkb(CgroupSkb {
data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
expected_attach_type: Some(CgroupSkbAttachType::Egress),
attach_type: Some(CgroupSkbAttachType::Egress),
}),
ProgramSection::CgroupSockAddr { attach_type, .. } => {
Program::CgroupSockAddr(CgroupSockAddr {
Expand Down
16 changes: 7 additions & 9 deletions aya/src/programs/cgroup_skb.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,18 +57,16 @@ use crate::{
#[doc(alias = "BPF_PROG_TYPE_CGROUP_SKB")]
pub struct CgroupSkb {
pub(crate) data: ProgramData<CgroupSkbLink>,
pub(crate) expected_attach_type: Option<CgroupSkbAttachType>,
pub(crate) attach_type: Option<CgroupSkbAttachType>,
}

impl CgroupSkb {
/// Loads the program inside the kernel.
pub fn load(&mut self) -> Result<(), ProgramError> {
self.data.expected_attach_type =
self.expected_attach_type
.map(|attach_type| match attach_type {
CgroupSkbAttachType::Ingress => BPF_CGROUP_INET_INGRESS,
CgroupSkbAttachType::Egress => BPF_CGROUP_INET_EGRESS,
});
self.data.expected_attach_type = self.attach_type.map(|attach_type| match attach_type {
CgroupSkbAttachType::Ingress => BPF_CGROUP_INET_INGRESS,
CgroupSkbAttachType::Egress => BPF_CGROUP_INET_EGRESS,
});
load_program(BPF_PROG_TYPE_CGROUP_SKB, &mut self.data)
}

Expand All @@ -79,7 +77,7 @@ impl CgroupSkb {
/// method returns `None` for programs defined with the generic section
/// `cgroup/skb`.
pub fn expected_attach_type(&self) -> &Option<CgroupSkbAttachType> {
&self.expected_attach_type
&self.attach_type
}

/// Attaches the program to the given cgroup.
Expand Down Expand Up @@ -138,7 +136,7 @@ impl CgroupSkb {
let data = ProgramData::from_pinned_path(path, VerifierLogLevel::default())?;
Ok(Self {
data,
expected_attach_type: Some(expected_attach_type),
attach_type: Some(expected_attach_type),
})
}
}
Expand Down
207 changes: 207 additions & 0 deletions aya/src/programs/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ use aya_obj::{
VerifierLog,
btf::BtfError,
generated::{bpf_attach_type, bpf_link_info, bpf_prog_info, bpf_prog_type},
programs::XdpAttachType,
};
use info::impl_info;
pub use info::{ProgramInfo, ProgramType, loaded_programs};
Expand Down Expand Up @@ -991,6 +992,212 @@ impl_from_pin!(
Iter,
);

macro_rules! impl_from_prog_info {
($(($struct_name:ident, $prog_type:expr, $unsafe:ident $(, $var:ident, $var_ty:ty)?)),* $(,)?) => {
$(
impl_from_prog_info!($struct_name, $prog_type, $unsafe $(, $var, $var_ty)?);
)*
};
($struct_name:ident, $prog_type:expr, false) => {
impl $struct_name {
/// Constructs an instance of a Program from a [`ProgramInfo`].
///
/// This allows the caller to get a handle to an already loaded
/// program from the kernel without having to load it again.
///
/// # Errors
///
/// - If the program type reported by the kernel does not match
/// that of the type you're converting to.
/// - If the file descriptor of the program cannot be cloned.
pub fn from_program_info(
name: Option<&'static str>,
info: ProgramInfo,
) -> Result<Self, ProgramError> {
if info.program_type()? != $prog_type {
return Err(ProgramError::UnexpectedProgramType {});
}
let ProgramInfo { 0: bpf_progam_info} = info;
let fd = info.fd()?;
let fd = fd.as_fd().try_clone_to_owned()?;

Ok(Self {
data: ProgramData::from_bpf_prog_info(
name.map(|n| Cow::Borrowed(n)),
crate::MockableFd::from_fd(fd),
Path::new(""),
bpf_progam_info,
VerifierLogLevel::default(),
)?,
})
}
}
};
($struct_name:ident, $prog_type:expr, true) => {
impl $struct_name {
/// Constructs an instance of a Program from a [`ProgramInfo`].
///
/// This allows the caller to get a handle to an already loaded
/// program from the kernel without having to load it again.
///
/// # Errors
///
/// - If the program type reported by the kernel does not match
/// that of the type you're converting to.
/// - If the file descriptor of the program cannot be cloned.
///
/// # Safety
///
/// We can't safely cast to `ProgramInfo` since we don't know the
/// concrete type of the program. It's up to the caller to ensure
/// that the program type matches the type you're converting to.
/// Otherwise, the behavior is undefined.
pub unsafe fn from_program_info(
name: Option<&'static str>,
info: ProgramInfo,
) -> Result<Self, ProgramError> {
if info.program_type()? != $prog_type {
return Err(ProgramError::UnexpectedProgramType {});
}
let ProgramInfo { 0: bpf_progam_info} = info;
let fd = info.fd()?;
let fd = fd.as_fd().try_clone_to_owned()?;

Ok(Self {
data: ProgramData::from_bpf_prog_info(
name.map(|n| Cow::Borrowed(n)),
crate::MockableFd::from_fd(fd),
Path::new(""),
bpf_progam_info,
VerifierLogLevel::default(),
)?,
})
}
}
};
($struct_name:ident, $prog_type:expr, false, $var:ident, $var_ty:ty) => {
impl $struct_name {
/// Constructs an instance of a Program from a [`ProgramInfo`].
///
/// This allows the caller to get a handle to an already loaded
/// program from the kernel without having to load it again.
///
/// # Errors
///
/// - If the program type reported by the kernel does not match
/// that of the type you're converting to.
/// - If the file descriptor of the program cannot be cloned.
pub fn from_program_info(
name: Option<&'static str>,
info: ProgramInfo,
$var: $var_ty,
) -> Result<Self, ProgramError> {
if info.program_type()? != $prog_type {
return Err(ProgramError::UnexpectedProgramType {});
}
let ProgramInfo { 0: bpf_progam_info} = info;
let fd = info.fd()?;
let fd = fd.as_fd().try_clone_to_owned()?;

Ok(Self {
data: ProgramData::from_bpf_prog_info(
name.map(|n| Cow::Borrowed(n)),
crate::MockableFd::from_fd(fd),
Path::new(""),
bpf_progam_info,
VerifierLogLevel::default(),
)?,
$var,
})
}
}
};
($struct_name:ident, $prog_type:expr, true, $var:ident, $var_ty:ty) => {
impl $struct_name {
/// Constructs an instance of a Program from a [`ProgramInfo`].
///
/// This allows the caller to get a handle to an already loaded
/// program from the kernel without having to load it again.
///
/// # Errors
///
/// - If the program type reported by the kernel does not match
/// that of the type you're converting to.
/// - If the file descriptor of the program cannot be cloned.
///
/// # Safety
///
/// We can't safely cast to `ProgramInfo` since we don't know the
/// concrete type of the program. It's up to the caller to ensure
/// that the program type matches the type you're converting to.
/// Otherwise, the behavior is undefined.
pub unsafe fn from_program_info(
name: Option<&'static str>,
info: ProgramInfo,
$var: $var_ty,
) -> Result<Self, ProgramError> {
if info.program_type()? != $prog_type {
return Err(ProgramError::UnexpectedProgramType {});
}
let ProgramInfo { 0: bpf_progam_info} = info;
let fd = info.fd()?;
let fd = fd.as_fd().try_clone_to_owned()?;

Ok(Self {
data: ProgramData::from_bpf_prog_info(
name.map(|n| Cow::Borrowed(n)),
crate::MockableFd::from_fd(fd),
Path::new(""),
bpf_progam_info,
VerifierLogLevel::default(),
)?,
$var,
})
}
}
};
}

// Order of arguments is as follows:
// - Program, bpf_prog_type, unsafe
// - Program, bpf_prog_type, unsafe, additional variable, variable type
impl_from_prog_info!(
(KProbe, ProgramType::KProbe, true, kind, ProbeKind),
(UProbe, ProgramType::KProbe, true, kind, ProbeKind),
(TracePoint, ProgramType::TracePoint, false),
(SocketFilter, ProgramType::SocketFilter, false),
(Xdp, ProgramType::Xdp, false, attach_type, XdpAttachType),
(SkMsg, ProgramType::SkMsg, false),
(SkSkb, ProgramType::SkSkb, false, kind, SkSkbKind),
(SockOps, ProgramType::SockOps, false),
(SchedClassifier, ProgramType::SchedClassifier, false),
(
CgroupSkb,
ProgramType::CgroupSkb,
false,
attach_type,
Option<CgroupSkbAttachType>
),
(CgroupSysctl, ProgramType::CgroupSysctl, false),
(
CgroupSockopt,
ProgramType::CgroupSockopt,
false,
attach_type,
CgroupSockoptAttachType
),
(LircMode2, ProgramType::LircMode2, false),
(PerfEvent, ProgramType::PerfEvent, false),
(Lsm, ProgramType::Lsm, false),
(RawTracePoint, ProgramType::RawTracePoint, false),
(BtfTracePoint, ProgramType::Tracing, true),
(FEntry, ProgramType::Tracing, true),
(FExit, ProgramType::Tracing, true),
(Extension, ProgramType::Extension, false),
(SkLookup, ProgramType::SkLookup, false),
(CgroupDevice, ProgramType::CgroupDevice, false),
);

macro_rules! impl_try_from_program {
($($ty:ident),+ $(,)?) => {
$(
Expand Down
27 changes: 24 additions & 3 deletions test/integration-test/src/tests/info.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use std::{fs, panic, path::Path, time::SystemTime};
use aya::{
Ebpf,
maps::{Array, HashMap, IterableMap as _, MapError, MapType, loaded_maps},
programs::{ProgramError, ProgramType, SocketFilter, TracePoint, loaded_programs},
programs::{ProgramError, ProgramType, SocketFilter, TracePoint, UProbe, loaded_programs},
sys::enable_stats,
util::KernelVersion,
};
Expand All @@ -25,8 +25,8 @@ const BPF_STATS_ENABLED: &str = "/proc/sys/kernel/bpf_stats_enabled";
fn test_loaded_programs() {
// Load a program.
// Since we are only testing the programs for their metadata, there is no need to "attach" them.
let mut bpf = Ebpf::load(crate::SIMPLE_PROG).unwrap();
let prog: &mut SocketFilter = bpf.program_mut("simple_prog").unwrap().try_into().unwrap();
let mut bpf = Ebpf::load(crate::TEST).unwrap();
let prog: &mut UProbe = bpf.program_mut("test_uprobe").unwrap().try_into().unwrap();
prog.load().unwrap();
let test_prog = prog.info().unwrap();

Expand Down Expand Up @@ -55,6 +55,27 @@ fn test_loaded_programs() {
programs.any(|prog| prog.id() == test_prog.id()),
KernelVersion::new(4, 13, 0)
);

// Iterate through loaded programs to exercise `from_program_info()`.
for program in loaded_programs() {
let program = program.unwrap();
let mut p: UProbe = unsafe {
UProbe::from_program_info(
Some("simple_prog"),
program,
aya::programs::ProbeKind::UProbe,
)
.unwrap()
};

// Ensure we can perform basic operations on the re-created program.
let res = p
.attach("uprobe_function", "/proc/self/exe", None, None)
.unwrap();

// Ensure the program can be detached
p.detach(res).unwrap();
}
}

#[test]
Expand Down
Loading

0 comments on commit 7e54054

Please sign in to comment.