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
dave-tucker committed Mar 6, 2025
1 parent 73a34e1 commit dc47958
Show file tree
Hide file tree
Showing 7 changed files with 287 additions and 15 deletions.
6 changes: 3 additions & 3 deletions aya/src/bpf.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -613,15 +613,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
View file
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
60 changes: 59 additions & 1 deletion aya/src/programs/kprobe.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,12 @@ use std::{
use aya_obj::generated::{bpf_link_type, bpf_prog_type::BPF_PROG_TYPE_KPROBE};
use thiserror::Error;

use super::ProgramType;
use crate::{
VerifierLogLevel,
programs::{
FdLink, LinkError, ProgramData, ProgramError, define_link_wrapper, load_program,
FdLink, LinkError, ProgramData, ProgramError, ProgramInfo, define_link_wrapper,
load_program,
perf_attach::{PerfLinkIdInner, PerfLinkInner},
probe::{ProbeKind, attach},
},
Expand Down Expand Up @@ -97,6 +99,62 @@ impl KProbe {
let data = ProgramData::from_pinned_path(path, VerifierLogLevel::default())?;
Ok(Self { data, kind })
}

/// Constructs an instance of [KProbe] from a [ProgramInfo].
///
/// This allows you to get a handle to an already loaded program
/// from the kernel without having to load it again.
///
/// # Safety
///
/// The type of program in the [ProgramInfo] may not be of
/// [`ProgramType::Kprobe`]. This function is unable to
/// check since the kernel doesn't expose this information.
/// The caller is responsible for ensuring that the program
/// type is correct (i.e they previously loaded it) otherwise
/// the behavior is undefined.
///
/// # Errors
///
/// - If the program type is not [`ProgramType::Kprobe`]
/// - If the file descriptor of the program cannot be cloned
///
/// # Example
/// ```no_run
/// # use aya::programs::loaded_programs;
/// # use aya::programs::{KProbe, ProbeKind, ProgramInfo, ProgramType};
///
/// for program in loaded_programs() {
/// let program = program?;
/// if program.program_type()? == ProgramType::KProbe {
/// let p = unsafe { KProbe::from_program_info(None, program, ProbeKind::KProbe)? };
/// // do something with the program
/// }
/// }
/// # Ok::<(), aya::programs::ProgramError>(())
/// ```
pub unsafe fn from_program_info(
name: Option<String>,
info: ProgramInfo,
kind: ProbeKind,
) -> Result<Self, ProgramError> {
if info.program_type()? != ProgramType::KProbe {
return Err(ProgramError::UnexpectedProgramType {});
}
if kind != ProbeKind::KProbe && kind != ProbeKind::KRetProbe {
return Err(ProgramError::UnexpectedProgramType {});
}
Ok(Self {
data: ProgramData::from_bpf_prog_info(
name,
crate::MockableFd::from_fd(info.fd()?.as_fd().try_clone_to_owned()?),
Path::new(""),
info.0,
VerifierLogLevel::default(),
)?,
kind,
})
}
}

define_link_wrapper!(
Expand Down
153 changes: 153 additions & 0 deletions aya/src/programs/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -990,6 +990,159 @@ impl_from_pin!(
Iter,
);

macro_rules! impl_from_prog_info {
($($struct_name:ident),+ $(,)?) => {
$(
impl $struct_name {
/// Consructs an program from a [ProgramInfo].
///
/// This allows you to get a handle to an already loaded program
/// from the kernel without having to load it again.
///
/// # Errors
///
/// - If the program type is not a match
/// - If the file descriptor of the program cannot be cloned
pub fn from_program_info(
name: Option<String>,
info: ProgramInfo,
) -> Result<$struct_name, ProgramError> {
if info.program_type()? != ProgramType::$struct_name {
return Err(ProgramError::UnexpectedProgramType {});
}
Ok($struct_name {
data: ProgramData::from_bpf_prog_info(
name,
crate::MockableFd::from_fd(info.fd()?.as_fd().try_clone_to_owned()?),
Path::new(""),
info.0,
VerifierLogLevel::default(),
)?,
})
}
}
)+
}
}

macro_rules! impl_from_prog_info_attach_type {
($($struct_name:ident, $attach_type:ty),+ $(,)?) => {
$(
impl $struct_name {

/// Consructs an program from a [ProgramInfo].
///
/// This allows you to get a handle to an already loaded program
/// from the kernel without having to load it again.
///
/// # Errors
///
/// - If the program type is not a match
/// - If the file descriptor of the program cannot be cloned
pub fn from_program_info(
name: Option<String>,
info: ProgramInfo,
attach_type: $attach_type,
) -> Result<$struct_name, ProgramError> {
if info.program_type()? != ProgramType::$struct_name {
return Err(ProgramError::UnexpectedProgramType {});
}
Ok($struct_name {
data: ProgramData::from_bpf_prog_info(
name,
crate::MockableFd::from_fd(info.fd()?.as_fd().try_clone_to_owned()?),
Path::new(""),
info.0,
VerifierLogLevel::default(),
)?,
attach_type,
})
}
}
)+
}
}

macro_rules! impl_from_prog_info_unsafe {
($($struct_name:ident, $prog_type:expr),+ $(,)?) => {
$(
impl $struct_name {

/// Consructs an program from a [ProgramInfo].
///
/// This allows you to get a handle to an already loaded program
/// from the kernel without having to load it again.
///
/// # Errors
///
/// - If the program type is not a match
/// - If the file descriptor of the program cannot be cloned
///
/// # Safety
///
/// The type of program in the [ProgramInfo] may not be of
/// [`$prog_type`]. This function is unable to
/// check since the kernel doesn't expose this information.
/// The caller is responsible for ensuring that the program
/// type is correct (i.e they previously loaded it) otherwise
/// the behavior is undefined.
pub unsafe fn from_program_info(
name: Option<String>,
info: ProgramInfo,
) -> Result<$struct_name, ProgramError> {
if info.program_type()? != $prog_type {
return Err(ProgramError::UnexpectedProgramType {});
}
Ok($struct_name {
data: ProgramData::from_bpf_prog_info(
name,
crate::MockableFd::from_fd(info.fd()?.as_fd().try_clone_to_owned()?),
Path::new(""),
info.0,
VerifierLogLevel::default(),
)?,
})
}
}
)+
}
}

impl_from_prog_info_unsafe!(
BtfTracePoint,
ProgramType::Tracing,
FEntry,
ProgramType::Tracing,
FExit,
ProgramType::Tracing,
);

impl_from_prog_info!(
CgroupDevice,
CgroupSysctl,
Extension,
LircMode2,
Lsm,
PerfEvent,
RawTracePoint,
SkLookup,
SkMsg,
SockOps,
SocketFilter,
SchedClassifier,
);

impl_from_prog_info_attach_type!(
CgroupSkb,
Option<CgroupSkbAttachType>,
CgroupSockAddr,
CgroupSockAddrAttachType,
CgroupSock,
CgroupSockAttachType,
CgroupSockopt,
CgroupSockoptAttachType,
);

macro_rules! impl_try_from_program {
($($ty:ident),+ $(,)?) => {
$(
Expand Down
2 changes: 1 addition & 1 deletion aya/src/programs/probe.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ use crate::{
static PROBE_NAME_INDEX: AtomicUsize = AtomicUsize::new(0);

/// Kind of probe program
#[derive(Debug, Copy, Clone)]
#[derive(Debug, Copy, Clone, PartialEq)]
pub enum ProbeKind {
/// Kernel probe
KProbe,
Expand Down
59 changes: 58 additions & 1 deletion aya/src/programs/uprobe.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,8 @@ use thiserror::Error;
use crate::{
VerifierLogLevel,
programs::{
FdLink, LinkError, ProgramData, ProgramError, define_link_wrapper, load_program,
FdLink, LinkError, ProgramData, ProgramError, ProgramInfo, ProgramType,
define_link_wrapper, load_program,
perf_attach::{PerfLinkIdInner, PerfLinkInner},
probe::{OsStringExt as _, ProbeKind, attach},
},
Expand Down Expand Up @@ -139,6 +140,62 @@ impl UProbe {
let data = ProgramData::from_pinned_path(path, VerifierLogLevel::default())?;
Ok(Self { data, kind })
}

/// Constructs an instance of [UProbe] from a [ProgramInfo].
///
/// This allows you to get a handle to an already loaded program
/// from the kernel without having to load it again.
///
/// # Safety
///
/// The type of program in the [ProgramInfo] may not be of
/// [`ProgramType::Kprobe`]. This function is unable to
/// check since the kernel doesn't expose this information.
/// The caller is responsible for ensuring that the program
/// type is correct (i.e they previously loaded it) otherwise
/// the behavior is undefined.
///
/// # Errors
///
/// - If the program type is not [`ProgramType::Kprobe`]
/// - If the file descriptor of the program cannot be cloned
///
/// # Example
/// ```no_run
/// # use aya::programs::loaded_programs;
/// # use aya::programs::{UProbe, ProbeKind, ProgramInfo, ProgramType};
///
/// for program in loaded_programs() {
/// let program = program?;
/// if program.program_type()? == ProgramType::KProbe {
/// let p = unsafe { UProbe::from_program_info(None, program, ProbeKind::UProbe)? };
/// // do something with the program
/// }
/// }
/// # Ok::<(), aya::programs::ProgramError>(())
/// ```
pub unsafe fn from_program_info(
name: Option<String>,
info: ProgramInfo,
kind: ProbeKind,
) -> Result<Self, ProgramError> {
if info.program_type()? != ProgramType::KProbe {
return Err(ProgramError::UnexpectedProgramType {});
}
if kind != ProbeKind::UProbe && kind != ProbeKind::URetProbe {
return Err(ProgramError::UnexpectedProgramType {});
}
Ok(Self {
data: ProgramData::from_bpf_prog_info(
name,
crate::MockableFd::from_fd(info.fd()?.as_fd().try_clone_to_owned()?),
Path::new(""),
info.0,
VerifierLogLevel::default(),
)?,
kind,
})
}
}

fn resolve_attach_path<'a, 'b, 'c, T>(
Expand Down
6 changes: 6 additions & 0 deletions test/integration-test/src/tests/info.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,12 @@ fn test_loaded_programs() {
programs.any(|prog| prog.id() == test_prog.id()),
KernelVersion::new(4, 13, 0)
);

for program in programs {
let mut p: SocketFilter =
SocketFilter::from_program_info(Some("simple_prog".to_string()), program).unwrap();
p.unload().unwrap();
}
}

#[test]
Expand Down

0 comments on commit dc47958

Please sign in to comment.