Merge pull request #261 from dave-tucker/cgroup_sock

Add Support for BPF_PROG_TYPE_CGROUP_SOCK_ADDR
pull/267/head
Alessandro Decina 2 years ago committed by GitHub
commit 8fd8816dfd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -22,10 +22,10 @@ use crate::{
MapKind, Object, ParseError, ProgramSection, MapKind, Object, ParseError, ProgramSection,
}, },
programs::{ programs::{
BtfTracePoint, CgroupSkb, CgroupSkbAttachType, CgroupSysctl, Extension, FEntry, FExit, BtfTracePoint, CgroupSkb, CgroupSkbAttachType, CgroupSockAddr, CgroupSysctl, Extension,
KProbe, LircMode2, Lsm, PerfEvent, ProbeKind, Program, ProgramData, ProgramError, FEntry, FExit, KProbe, LircMode2, Lsm, PerfEvent, ProbeKind, Program, ProgramData,
RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind, SockOps, SocketFilter, TracePoint, ProgramError, RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind, SockOps,
UProbe, Xdp, SocketFilter, TracePoint, UProbe, Xdp,
}, },
sys::{ sys::{
bpf_load_btf, bpf_map_freeze, bpf_map_update_elem_ptr, is_btf_datasec_supported, bpf_load_btf, bpf_map_freeze, bpf_map_update_elem_ptr, is_btf_datasec_supported,
@ -481,6 +481,12 @@ impl<'a> BpfLoader<'a> {
data: ProgramData::new(prog_name, obj, btf_fd), data: ProgramData::new(prog_name, obj, btf_fd),
expected_attach_type: Some(CgroupSkbAttachType::Egress), expected_attach_type: Some(CgroupSkbAttachType::Egress),
}), }),
ProgramSection::CgroupSockAddr { attach_type, .. } => {
Program::CgroupSockAddr(CgroupSockAddr {
data: ProgramData::new(prog_name, obj, btf_fd),
attach_type: *attach_type,
})
}
ProgramSection::LircMode2 { .. } => Program::LircMode2(LircMode2 { ProgramSection::LircMode2 { .. } => Program::LircMode2(LircMode2 {
data: ProgramData::new(prog_name, obj, btf_fd), data: ProgramData::new(prog_name, obj, btf_fd),
}), }),

@ -21,6 +21,7 @@ use crate::{
bpf_map_def, bpf_map_def,
generated::{bpf_insn, bpf_map_type::BPF_MAP_TYPE_ARRAY, BPF_F_RDONLY_PROG}, generated::{bpf_insn, bpf_map_type::BPF_MAP_TYPE_ARRAY, BPF_F_RDONLY_PROG},
obj::btf::{Btf, BtfError, BtfExt}, obj::btf::{Btf, BtfError, BtfExt},
programs::CgroupSockAddrAttachType,
BpfError, BpfError,
}; };
use std::slice::from_raw_parts_mut; use std::slice::from_raw_parts_mut;
@ -104,30 +105,82 @@ pub(crate) struct Function {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub enum ProgramSection { pub enum ProgramSection {
KRetProbe { name: String }, KRetProbe {
KProbe { name: String }, name: String,
UProbe { name: String }, },
URetProbe { name: String }, KProbe {
TracePoint { name: String }, name: String,
SocketFilter { name: String }, },
Xdp { name: String }, UProbe {
SkMsg { name: String }, name: String,
SkSkbStreamParser { name: String }, },
SkSkbStreamVerdict { name: String }, URetProbe {
SockOps { name: String }, name: String,
SchedClassifier { name: String }, },
CgroupSkb { name: String }, TracePoint {
CgroupSkbIngress { name: String }, name: String,
CgroupSkbEgress { name: String }, },
CgroupSysctl { name: String }, SocketFilter {
LircMode2 { name: String }, name: String,
PerfEvent { name: String }, },
RawTracePoint { name: String }, Xdp {
Lsm { name: String }, name: String,
BtfTracePoint { name: String }, },
FEntry { name: String }, SkMsg {
FExit { name: String }, name: String,
Extension { name: String }, },
SkSkbStreamParser {
name: String,
},
SkSkbStreamVerdict {
name: String,
},
SockOps {
name: String,
},
SchedClassifier {
name: String,
},
CgroupSkb {
name: String,
},
CgroupSkbIngress {
name: String,
},
CgroupSkbEgress {
name: String,
},
CgroupSockAddr {
name: String,
attach_type: CgroupSockAddrAttachType,
},
CgroupSysctl {
name: String,
},
LircMode2 {
name: String,
},
PerfEvent {
name: String,
},
RawTracePoint {
name: String,
},
Lsm {
name: String,
},
BtfTracePoint {
name: String,
},
FEntry {
name: String,
},
FExit {
name: String,
},
Extension {
name: String,
},
} }
impl ProgramSection { impl ProgramSection {
@ -145,9 +198,10 @@ impl ProgramSection {
ProgramSection::SkSkbStreamVerdict { name } => name, ProgramSection::SkSkbStreamVerdict { name } => name,
ProgramSection::SockOps { name } => name, ProgramSection::SockOps { name } => name,
ProgramSection::SchedClassifier { name } => name, ProgramSection::SchedClassifier { name } => name,
ProgramSection::CgroupSkb { name } => name, ProgramSection::CgroupSkb { name, .. } => name,
ProgramSection::CgroupSkbIngress { name } => name, ProgramSection::CgroupSkbIngress { name, .. } => name,
ProgramSection::CgroupSkbEgress { name } => name, ProgramSection::CgroupSkbEgress { name, .. } => name,
ProgramSection::CgroupSockAddr { name, .. } => name,
ProgramSection::CgroupSysctl { name } => name, ProgramSection::CgroupSysctl { name } => name,
ProgramSection::LircMode2 { name } => name, ProgramSection::LircMode2 { name } => name,
ProgramSection::PerfEvent { name } => name, ProgramSection::PerfEvent { name } => name,
@ -221,11 +275,63 @@ impl FromStr for ProgramSection {
"skb" => CgroupSkb { name }, "skb" => CgroupSkb { name },
"sysctl" => CgroupSysctl { name }, "sysctl" => CgroupSysctl { name },
_ => { _ => {
return Err(ParseError::InvalidProgramSection { if let Ok(attach_type) = CgroupSockAddrAttachType::try_from(name.as_str()) {
section: section.to_owned(), CgroupSockAddr { name, attach_type }
}) } else {
return Err(ParseError::InvalidProgramSection {
section: section.to_owned(),
});
}
} }
}, },
"cgroup/bind4" => CgroupSockAddr {
name,
attach_type: CgroupSockAddrAttachType::Bind4,
},
"cgroup/bind6" => CgroupSockAddr {
name,
attach_type: CgroupSockAddrAttachType::Bind6,
},
"cgroup/connect4" => CgroupSockAddr {
name,
attach_type: CgroupSockAddrAttachType::Connect4,
},
"cgroup/connect6" => CgroupSockAddr {
name,
attach_type: CgroupSockAddrAttachType::Connect6,
},
"cgroup/getpeername4" => CgroupSockAddr {
name,
attach_type: CgroupSockAddrAttachType::GetPeerName4,
},
"cgroup/getpeername6" => CgroupSockAddr {
name,
attach_type: CgroupSockAddrAttachType::GetPeerName6,
},
"cgroup/getsockname4" => CgroupSockAddr {
name,
attach_type: CgroupSockAddrAttachType::GetSockName4,
},
"cgroup/getsockname6" => CgroupSockAddr {
name,
attach_type: CgroupSockAddrAttachType::GetSockName6,
},
"cgroup/sendmsg4" => CgroupSockAddr {
name,
attach_type: CgroupSockAddrAttachType::UDPSendMsg4,
},
"cgroup/sendmsg6" => CgroupSockAddr {
name,
attach_type: CgroupSockAddrAttachType::UDPSendMsg6,
},
"cgroup/recvmsg4" => CgroupSockAddr {
name,
attach_type: CgroupSockAddrAttachType::UDPRecvMsg4,
},
"cgroup/recvmsg6" => CgroupSockAddr {
name,
attach_type: CgroupSockAddrAttachType::UDPRecvMsg6,
},
"lirc_mode2" => LircMode2 { name }, "lirc_mode2" => LircMode2 { name },
"perf_event" => PerfEvent { name }, "perf_event" => PerfEvent { name },
"raw_tp" | "raw_tracepoint" => RawTracePoint { name }, "raw_tp" | "raw_tracepoint" => RawTracePoint { name },
@ -1641,6 +1747,54 @@ mod tests {
); );
} }
#[test]
fn test_parse_section_sock_addr_named() {
let mut obj = fake_obj();
assert_matches!(
obj.parse_section(fake_section(
BpfSectionKind::Program,
"cgroup/connect4/foo",
bytes_of(&fake_ins())
)),
Ok(())
);
assert_matches!(
obj.programs.get("foo"),
Some(Program {
section: ProgramSection::CgroupSockAddr {
attach_type: CgroupSockAddrAttachType::Connect4,
..
},
..
})
);
}
#[test]
fn test_parse_section_sock_addr_unnamed() {
let mut obj = fake_obj();
assert_matches!(
obj.parse_section(fake_section(
BpfSectionKind::Program,
"cgroup/connect4",
bytes_of(&fake_ins())
)),
Ok(())
);
assert_matches!(
obj.programs.get("connect4"),
Some(Program {
section: ProgramSection::CgroupSockAddr {
attach_type: CgroupSockAddrAttachType::Connect4,
..
},
..
})
);
}
#[test] #[test]
fn test_patch_map_data() { fn test_patch_map_data() {
let mut obj = fake_obj(); let mut obj = fake_obj();

@ -0,0 +1,233 @@
use thiserror::Error;
use crate::generated::bpf_attach_type;
use std::{
hash::Hash,
os::unix::prelude::{AsRawFd, RawFd},
};
use crate::{
generated::bpf_prog_type::BPF_PROG_TYPE_CGROUP_SOCK_ADDR,
programs::{
define_link_wrapper, load_program, FdLink, Link, OwnedLink, ProgAttachLink, ProgramData,
ProgramError,
},
sys::{bpf_link_create, bpf_prog_attach, kernel_version},
};
/// A program that can be used to inspect or modify socket addresses (`struct sockaddr`).
///
/// [`SockAddr`] programs can be used to inspect or modify socket addresses passed to
/// various syscalls within a [cgroup]. They can be attached to a number of different
/// places as described in [`SockAddrAttachType`].
///
/// [cgroup]: https://man7.org/linux/man-pages/man7/cgroups.7.html
///
/// # Minimum kernel version
///
/// The minimum kernel version required to use this feature is 4.17.
///
/// # Examples
///
/// ```no_run
/// # #[derive(thiserror::Error, Debug)]
/// # enum Error {
/// # #[error(transparent)]
/// # IO(#[from] std::io::Error),
/// # #[error(transparent)]
/// # Map(#[from] aya::maps::MapError),
/// # #[error(transparent)]
/// # Program(#[from] aya::programs::ProgramError),
/// # #[error(transparent)]
/// # Bpf(#[from] aya::BpfError)
/// # }
/// # let mut bpf = aya::Bpf::load(&[])?;
/// use std::fs::File;
/// use std::convert::TryInto;
/// use aya::programs::{CgroupSockAddr, CgroupSockAddrAttachType};
///
/// let file = File::open("/sys/fs/cgroup/unified")?;
/// let egress: &mut CgroupSockAddr = bpf.program_mut("connect4").unwrap().try_into()?;
/// egress.load()?;
/// egress.attach(file)?;
/// # Ok::<(), Error>(())
/// ```
#[derive(Debug)]
#[doc(alias = "BPF_PROG_TYPE_CGROUP_SOCK_ADDR")]
pub struct CgroupSockAddr {
pub(crate) data: ProgramData<CgroupSockAddrLink>,
pub(crate) attach_type: CgroupSockAddrAttachType,
}
impl CgroupSockAddr {
/// Loads the program inside the kernel.
pub fn load(&mut self) -> Result<(), ProgramError> {
self.data.expected_attach_type = Some(self.attach_type.into());
load_program(BPF_PROG_TYPE_CGROUP_SOCK_ADDR, &mut self.data)
}
/// Attaches the program to the given cgroup.
///
/// The returned value can be used to detach, see [SockAddr::detach].
pub fn attach<T: AsRawFd>(&mut self, cgroup: T) -> Result<CgroupSockAddrLinkId, ProgramError> {
let prog_fd = self.data.fd_or_err()?;
let cgroup_fd = cgroup.as_raw_fd();
let attach_type = self.data.expected_attach_type.unwrap();
let k_ver = kernel_version().unwrap();
if k_ver >= (5, 7, 0) {
let link_fd = bpf_link_create(prog_fd, cgroup_fd, attach_type, None, 0).map_err(
|(_, io_error)| ProgramError::SyscallError {
call: "bpf_link_create".to_owned(),
io_error,
},
)? as RawFd;
self.data
.links
.insert(CgroupSockAddrLink(CgroupSockAddrLinkInner::Fd(
FdLink::new(link_fd),
)))
} else {
bpf_prog_attach(prog_fd, cgroup_fd, attach_type).map_err(|(_, io_error)| {
ProgramError::SyscallError {
call: "bpf_prog_attach".to_owned(),
io_error,
}
})?;
self.data
.links
.insert(CgroupSockAddrLink(CgroupSockAddrLinkInner::ProgAttach(
ProgAttachLink::new(prog_fd, cgroup_fd, attach_type),
)))
}
}
/// Takes ownership of the link referenced by the provided link_id.
///
/// The link will be detached on `Drop` and the caller is now responsible
/// for managing its lifetime.
pub fn forget_link(
&mut self,
link_id: CgroupSockAddrLinkId,
) -> Result<OwnedLink<CgroupSockAddrLink>, ProgramError> {
Ok(OwnedLink::new(self.data.forget_link(link_id)?))
}
/// Detaches the program.
///
/// See [SockAddr::attach].
pub fn detach(&mut self, link_id: CgroupSockAddrLinkId) -> Result<(), ProgramError> {
self.data.links.remove(link_id)
}
}
#[derive(Debug, Hash, Eq, PartialEq)]
enum CgroupSockAddrLinkIdInner {
Fd(<FdLink as Link>::Id),
ProgAttach(<ProgAttachLink as Link>::Id),
}
#[derive(Debug)]
enum CgroupSockAddrLinkInner {
Fd(FdLink),
ProgAttach(ProgAttachLink),
}
impl Link for CgroupSockAddrLinkInner {
type Id = CgroupSockAddrLinkIdInner;
fn id(&self) -> Self::Id {
match self {
CgroupSockAddrLinkInner::Fd(fd) => CgroupSockAddrLinkIdInner::Fd(fd.id()),
CgroupSockAddrLinkInner::ProgAttach(p) => CgroupSockAddrLinkIdInner::ProgAttach(p.id()),
}
}
fn detach(self) -> Result<(), ProgramError> {
match self {
CgroupSockAddrLinkInner::Fd(fd) => fd.detach(),
CgroupSockAddrLinkInner::ProgAttach(p) => p.detach(),
}
}
}
define_link_wrapper!(
/// The link used by [SockAddr] programs.
CgroupSockAddrLink,
/// The type returned by [SockAddr::attach]. Can be passed to [SockAddr::detach].
CgroupSockAddrLinkId,
CgroupSockAddrLinkInner,
CgroupSockAddrLinkIdInner
);
/// Defines where to attach a [`SockAddr`] program.
#[derive(Copy, Clone, Debug)]
pub enum CgroupSockAddrAttachType {
/// Attach to IPv4 bind events.
Bind4,
/// Attach to IPv6 bind events.
Bind6,
/// Attach to IPv4 connect events.
Connect4,
/// Attach to IPv6 connect events.
Connect6,
/// Attach to IPv4 getpeername events.
GetPeerName4,
/// Attach to IPv6 getpeername events.
GetPeerName6,
/// Attach to IPv4 getsockname events.
GetSockName4,
/// Attach to IPv6 getsockname events.
GetSockName6,
/// Attach to IPv4 udp_sendmsg events.
UDPSendMsg4,
/// Attach to IPv6 udp_sendmsg events.
UDPSendMsg6,
/// Attach to IPv4 udp_recvmsg events.
UDPRecvMsg4,
/// Attach to IPv6 udp_recvmsg events.
UDPRecvMsg6,
}
impl From<CgroupSockAddrAttachType> for bpf_attach_type {
fn from(s: CgroupSockAddrAttachType) -> bpf_attach_type {
match s {
CgroupSockAddrAttachType::Bind4 => bpf_attach_type::BPF_CGROUP_INET4_BIND,
CgroupSockAddrAttachType::Bind6 => bpf_attach_type::BPF_CGROUP_INET6_BIND,
CgroupSockAddrAttachType::Connect4 => bpf_attach_type::BPF_CGROUP_INET4_CONNECT,
CgroupSockAddrAttachType::Connect6 => bpf_attach_type::BPF_CGROUP_INET6_CONNECT,
CgroupSockAddrAttachType::GetPeerName4 => bpf_attach_type::BPF_CGROUP_INET4_GETPEERNAME,
CgroupSockAddrAttachType::GetPeerName6 => bpf_attach_type::BPF_CGROUP_INET6_GETPEERNAME,
CgroupSockAddrAttachType::GetSockName4 => bpf_attach_type::BPF_CGROUP_INET4_GETSOCKNAME,
CgroupSockAddrAttachType::GetSockName6 => bpf_attach_type::BPF_CGROUP_INET6_GETSOCKNAME,
CgroupSockAddrAttachType::UDPSendMsg4 => bpf_attach_type::BPF_CGROUP_UDP4_SENDMSG,
CgroupSockAddrAttachType::UDPSendMsg6 => bpf_attach_type::BPF_CGROUP_UDP6_SENDMSG,
CgroupSockAddrAttachType::UDPRecvMsg4 => bpf_attach_type::BPF_CGROUP_UDP4_RECVMSG,
CgroupSockAddrAttachType::UDPRecvMsg6 => bpf_attach_type::BPF_CGROUP_UDP6_RECVMSG,
}
}
}
#[derive(Debug, Error)]
#[error("{0} is not a valid attach type for a CGROUP_SOCK_ADDR program")]
pub(crate) struct InvalidAttachType(String);
impl CgroupSockAddrAttachType {
pub(crate) fn try_from(value: &str) -> Result<CgroupSockAddrAttachType, InvalidAttachType> {
match value {
"bind4" => Ok(CgroupSockAddrAttachType::Bind4),
"bind6" => Ok(CgroupSockAddrAttachType::Bind6),
"connect4" => Ok(CgroupSockAddrAttachType::Connect4),
"connect6" => Ok(CgroupSockAddrAttachType::Connect6),
"getpeername4" => Ok(CgroupSockAddrAttachType::GetPeerName4),
"getpeername6" => Ok(CgroupSockAddrAttachType::GetPeerName6),
"getsockname4" => Ok(CgroupSockAddrAttachType::GetSockName4),
"getsockname6" => Ok(CgroupSockAddrAttachType::GetSockName6),
"sendmsg4" => Ok(CgroupSockAddrAttachType::UDPSendMsg4),
"sendmsg6" => Ok(CgroupSockAddrAttachType::UDPSendMsg6),
"recvmsg4" => Ok(CgroupSockAddrAttachType::UDPRecvMsg4),
"recvmsg6" => Ok(CgroupSockAddrAttachType::UDPRecvMsg6),
_ => Err(InvalidAttachType(value.to_owned())),
}
}
}

@ -37,6 +37,7 @@
//! [`Bpf::program_mut`]: crate::Bpf::program_mut //! [`Bpf::program_mut`]: crate::Bpf::program_mut
//! [`maps`]: crate::maps //! [`maps`]: crate::maps
mod cgroup_skb; mod cgroup_skb;
mod cgroup_sock_addr;
mod cgroup_sysctl; mod cgroup_sysctl;
mod extension; mod extension;
mod fentry; mod fentry;
@ -71,6 +72,7 @@ use std::{
use thiserror::Error; use thiserror::Error;
pub use cgroup_skb::{CgroupSkb, CgroupSkbAttachType}; pub use cgroup_skb::{CgroupSkb, CgroupSkbAttachType};
pub use cgroup_sock_addr::{CgroupSockAddr, CgroupSockAddrAttachType};
pub use cgroup_sysctl::CgroupSysctl; pub use cgroup_sysctl::CgroupSysctl;
pub use extension::{Extension, ExtensionError}; pub use extension::{Extension, ExtensionError};
pub use fentry::FEntry; pub use fentry::FEntry;
@ -229,6 +231,8 @@ pub enum Program {
SkMsg(SkMsg), SkMsg(SkMsg),
/// A [`SkSkb`] program /// A [`SkSkb`] program
SkSkb(SkSkb), SkSkb(SkSkb),
/// A [`CgroupSockAddr`] program
CgroupSockAddr(CgroupSockAddr),
/// A [`SockOps`] program /// A [`SockOps`] program
SockOps(SockOps), SockOps(SockOps),
/// A [`SchedClassifier`] program /// A [`SchedClassifier`] program
@ -279,6 +283,7 @@ impl Program {
Program::FEntry(_) => BPF_PROG_TYPE_TRACING, Program::FEntry(_) => BPF_PROG_TYPE_TRACING,
Program::FExit(_) => BPF_PROG_TYPE_TRACING, Program::FExit(_) => BPF_PROG_TYPE_TRACING,
Program::Extension(_) => BPF_PROG_TYPE_EXT, Program::Extension(_) => BPF_PROG_TYPE_EXT,
Program::CgroupSockAddr(_) => BPF_PROG_TYPE_CGROUP_SOCK_ADDR,
} }
} }
@ -304,6 +309,7 @@ impl Program {
Program::FEntry(p) => p.data.pin(path), Program::FEntry(p) => p.data.pin(path),
Program::FExit(p) => p.data.pin(path), Program::FExit(p) => p.data.pin(path),
Program::Extension(p) => p.data.pin(path), Program::Extension(p) => p.data.pin(path),
Program::CgroupSockAddr(p) => p.data.pin(path),
} }
} }
} }
@ -509,6 +515,7 @@ impl ProgramFd for Program {
Program::FEntry(p) => p.data.fd, Program::FEntry(p) => p.data.fd,
Program::FExit(p) => p.data.fd, Program::FExit(p) => p.data.fd,
Program::Extension(p) => p.data.fd, Program::Extension(p) => p.data.fd,
Program::CgroupSockAddr(p) => p.data.fd,
} }
} }
} }
@ -556,6 +563,7 @@ impl_program_fd!(
FEntry, FEntry,
FExit, FExit,
Extension, Extension,
CgroupSockAddr,
); );
macro_rules! impl_try_from_program { macro_rules! impl_try_from_program {
@ -606,6 +614,7 @@ impl_try_from_program!(
FEntry, FEntry,
FExit, FExit,
Extension, Extension,
CgroupSockAddr,
); );
/// Provides information about a loaded program, like name, id and statistics /// Provides information about a loaded program, like name, id and statistics

@ -37,6 +37,29 @@ impl Parse for Args {
} }
} }
pub struct SockAddrArgs {
pub(crate) attach_type: Ident,
pub(crate) args: Args,
}
impl Parse for SockAddrArgs {
fn parse(input: ParseStream) -> Result<SockAddrArgs> {
let attach_type: Ident = input.parse()?;
match attach_type.to_string().as_str() {
"connect4" | "connect6" | "bind4" | "bind6" | "getpeername4" | "getpeername6"
| "getsockname4" | "getsockname6" | "sendmsg4" | "sendmsg6" | "recvmsg4"
| "recvmsg6" => (),
_ => return Err(input.error("invalid attach type")),
}
let args = if input.parse::<Token![,]>().is_ok() {
Args::parse(input)?
} else {
Args { args: vec![] }
};
Ok(SockAddrArgs { attach_type, args })
}
}
pub struct Map { pub struct Map {
item: ItemStatic, item: ItemStatic,
name: String, name: String,
@ -257,6 +280,7 @@ impl CgroupSkb {
pub fn from_syn(mut args: Args, item: ItemFn) -> Result<CgroupSkb> { pub fn from_syn(mut args: Args, item: ItemFn) -> Result<CgroupSkb> {
let name = pop_arg(&mut args, "name"); let name = pop_arg(&mut args, "name");
let expected_attach_type = pop_arg(&mut args, "attach"); let expected_attach_type = pop_arg(&mut args, "attach");
err_on_unknown_args(&args)?;
Ok(CgroupSkb { Ok(CgroupSkb {
item, item,
@ -291,6 +315,44 @@ impl CgroupSkb {
} }
} }
pub struct CgroupSockAddr {
item: ItemFn,
attach_type: String,
name: Option<String>,
}
impl CgroupSockAddr {
pub fn from_syn(mut args: Args, item: ItemFn, attach_type: String) -> Result<CgroupSockAddr> {
let name = pop_arg(&mut args, "name");
err_on_unknown_args(&args)?;
Ok(CgroupSockAddr {
item,
attach_type,
name,
})
}
pub fn expand(&self) -> Result<TokenStream> {
let section_name = if let Some(name) = &self.name {
format!("cgroup/{}/{}", self.attach_type, name)
} else {
format!("cgroup/{}", self.attach_type)
};
let fn_name = &self.item.sig.ident;
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return #fn_name(::aya_bpf::programs::SockAddrContext::new(ctx));
#item
}
})
}
}
fn pop_arg(args: &mut Args, name: &str) -> Option<String> { fn pop_arg(args: &mut Args, name: &str) -> Option<String> {
match args.args.iter().position(|arg| arg.name == name) { match args.args.iter().position(|arg| arg.name == name) {
Some(index) => Some(args.args.remove(index).value.value()), Some(index) => Some(args.args.remove(index).value.value()),
@ -507,6 +569,8 @@ pub struct SkSkb {
impl SkSkb { impl SkSkb {
pub fn from_syn(kind: SkSkbKind, mut args: Args, item: ItemFn) -> Result<SkSkb> { pub fn from_syn(kind: SkSkbKind, mut args: Args, item: ItemFn) -> Result<SkSkb> {
let name = pop_arg(&mut args, "name"); let name = pop_arg(&mut args, "name");
err_on_unknown_args(&args)?;
Ok(SkSkb { item, kind, name }) Ok(SkSkb { item, kind, name })
} }
@ -539,6 +603,7 @@ pub struct SocketFilter {
impl SocketFilter { impl SocketFilter {
pub fn from_syn(mut args: Args, item: ItemFn) -> Result<SocketFilter> { pub fn from_syn(mut args: Args, item: ItemFn) -> Result<SocketFilter> {
let name = name_arg(&mut args)?; let name = name_arg(&mut args)?;
err_on_unknown_args(&args)?;
Ok(SocketFilter { item, name }) Ok(SocketFilter { item, name })
} }

@ -1,9 +1,9 @@
mod expand; mod expand;
use expand::{ use expand::{
Args, BtfTracePoint, CgroupSkb, CgroupSysctl, FEntry, FExit, Lsm, Map, PerfEvent, Probe, Args, BtfTracePoint, CgroupSkb, CgroupSockAddr, CgroupSysctl, FEntry, FExit, Lsm, Map,
ProbeKind, RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind, SockOps, SocketFilter, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind,
TracePoint, Xdp, SockAddrArgs, SockOps, SocketFilter, TracePoint, Xdp,
}; };
use proc_macro::TokenStream; use proc_macro::TokenStream;
use syn::{parse_macro_input, ItemFn, ItemStatic}; use syn::{parse_macro_input, ItemFn, ItemStatic};
@ -105,6 +105,18 @@ pub fn cgroup_skb(attrs: TokenStream, item: TokenStream) -> TokenStream {
.into() .into()
} }
#[proc_macro_attribute]
pub fn cgroup_sock_addr(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as SockAddrArgs);
let attach_type = args.attach_type.to_string();
let item = parse_macro_input!(item as ItemFn);
CgroupSockAddr::from_syn(args.args, item, attach_type)
.and_then(|u| u.expand())
.unwrap_or_else(|err| err.to_compile_error())
.into()
}
fn probe(kind: ProbeKind, attrs: TokenStream, item: TokenStream) -> TokenStream { fn probe(kind: ProbeKind, attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); let args = parse_macro_input!(attrs as Args);
let item = parse_macro_input!(item as ItemFn); let item = parse_macro_input!(item as ItemFn);

@ -6,6 +6,7 @@ pub mod probe;
pub mod raw_tracepoint; pub mod raw_tracepoint;
pub mod sk_buff; pub mod sk_buff;
pub mod sk_msg; pub mod sk_msg;
pub mod sock_addr;
pub mod sock_ops; pub mod sock_ops;
pub mod sysctl; pub mod sysctl;
pub mod tp_btf; pub mod tp_btf;
@ -20,6 +21,7 @@ pub use probe::ProbeContext;
pub use raw_tracepoint::RawTracePointContext; pub use raw_tracepoint::RawTracePointContext;
pub use sk_buff::SkBuffContext; pub use sk_buff::SkBuffContext;
pub use sk_msg::SkMsgContext; pub use sk_msg::SkMsgContext;
pub use sock_addr::SockAddrContext;
pub use sock_ops::SockOpsContext; pub use sock_ops::SockOpsContext;
pub use sysctl::SysctlContext; pub use sysctl::SysctlContext;
pub use tp_btf::BtfTracePointContext; pub use tp_btf::BtfTracePointContext;

@ -0,0 +1,19 @@
use core::ffi::c_void;
use crate::{bindings::bpf_sock_addr, BpfContext};
pub struct SockAddrContext {
pub sock_addr: *mut bpf_sock_addr,
}
impl SockAddrContext {
pub fn new(sock_addr: *mut bpf_sock_addr) -> SockAddrContext {
SockAddrContext { sock_addr }
}
}
impl BpfContext for SockAddrContext {
fn as_ptr(&self) -> *mut c_void {
self.sock_addr as *mut _
}
}
Loading…
Cancel
Save