Merge pull request #315 from dave-tucker/sock

Add support for BPF_PROG_TYPE_CGROUP_SOCK

aya/src/bpf.rs

@@ -22,9 +22,9 @@ use crate::{
         MapKind, Object, ParseError, ProgramSection,
     },
     programs::{
-        BtfTracePoint, CgroupSkb, CgroupSkbAttachType, CgroupSockAddr, CgroupSockopt, CgroupSysctl,
-        Extension, FEntry, FExit, KProbe, LircMode2, Lsm, PerfEvent, ProbeKind, Program,
-        ProgramData, ProgramError, RawTracePoint, SchedClassifier, SkLookup, SkMsg, SkSkb,
+        BtfTracePoint, CgroupSkb, CgroupSkbAttachType, CgroupSock, CgroupSockAddr, CgroupSockopt,
+        CgroupSysctl, Extension, FEntry, FExit, KProbe, LircMode2, Lsm, PerfEvent, ProbeKind,
+        Program, ProgramData, ProgramError, RawTracePoint, SchedClassifier, SkLookup, SkMsg, SkSkb,
         SkSkbKind, SockOps, SocketFilter, TracePoint, UProbe, Xdp,
     },
     sys::{
@@ -524,6 +524,12 @@ impl<'a> BpfLoader<'a> {
                         ProgramSection::SkLookup { .. } => Program::SkLookup(SkLookup {
                             data: ProgramData::new(prog_name, obj, btf_fd),
                         }),
+                        ProgramSection::CgroupSock { attach_type, .. } => {
+                            Program::CgroupSock(CgroupSock {
+                                data: ProgramData::new(prog_name, obj, btf_fd),
+                                attach_type: *attach_type,
+                            })
+                        }
                     }
                 };
                 (name, program)

aya/src/obj/mod.rs

@@ -21,7 +21,7 @@ use crate::{
     bpf_map_def,
     generated::{bpf_insn, bpf_map_type::BPF_MAP_TYPE_ARRAY, BPF_F_RDONLY_PROG},
     obj::btf::{Btf, BtfError, BtfExt},
-    programs::{CgroupSockAddrAttachType, CgroupSockoptAttachType},
+    programs::{CgroupSockAddrAttachType, CgroupSockAttachType, CgroupSockoptAttachType},
     BpfError,
 };
 use std::slice::from_raw_parts_mut;
@@ -188,6 +188,10 @@ pub enum ProgramSection {
     SkLookup {
         name: String,
     },
+    CgroupSock {
+        name: String,
+        attach_type: CgroupSockAttachType,
+    },
 }
 
 impl ProgramSection {
@@ -220,6 +224,7 @@ impl ProgramSection {
             ProgramSection::FExit { name } => name,
             ProgramSection::Extension { name } => name,
             ProgramSection::SkLookup { name } => name,
+            ProgramSection::CgroupSock { name, .. } => name,
         }
     }
 }
@@ -279,6 +284,10 @@ impl FromStr for ProgramSection {
             "cgroup_skb/ingress" => CgroupSkbIngress { name },
             "cgroup_skb/egress" => CgroupSkbEgress { name },
             "cgroup/skb" => CgroupSkb { name },
+            "cgroup/sock" => CgroupSock {
+                name,
+                attach_type: CgroupSockAttachType::default(),
+            },
             "cgroup/sysctl" => CgroupSysctl { name },
             "cgroup/getsockopt" => CgroupSockopt {
                 name,
@@ -300,6 +309,19 @@ impl FromStr for ProgramSection {
                         });
                     }
                 }
+                "sock" => CgroupSock {
+                    name,
+                    attach_type: CgroupSockAttachType::default(),
+                },
+                "post_bind4" | "post_bind6" | "sock_create" | "sock_release" => {
+                    if let Ok(attach_type) = CgroupSockAttachType::try_from(name.as_str()) {
+                        CgroupSock { name, attach_type }
+                    } else {
+                        return Err(ParseError::InvalidProgramSection {
+                            section: section.to_owned(),
+                        });
+                    }
+                }
                 _ => {
                     if let Ok(attach_type) = CgroupSockAddrAttachType::try_from(name.as_str()) {
                         CgroupSockAddr { name, attach_type }
@@ -310,6 +332,22 @@ impl FromStr for ProgramSection {
                     }
                 }
             },
+            "cgroup/post_bind4" => CgroupSock {
+                name,
+                attach_type: CgroupSockAttachType::PostBind4,
+            },
+            "cgroup/post_bind6" => CgroupSock {
+                name,
+                attach_type: CgroupSockAttachType::PostBind6,
+            },
+            "cgroup/sock_create" => CgroupSock {
+                name,
+                attach_type: CgroupSockAttachType::SockCreate,
+            },
+            "cgroup/sock_release" => CgroupSock {
+                name,
+                attach_type: CgroupSockAttachType::SockRelease,
+            },
             "cgroup/bind4" => CgroupSockAddr {
                 name,
                 attach_type: CgroupSockAddrAttachType::Bind4,

aya/src/programs/cgroup_sock.rs

@@ -0,0 +1,209 @@
+//! Cgroup socket programs.
+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,
+    programs::{
+        define_link_wrapper, load_program, FdLink, Link, OwnedLink, ProgAttachLink, ProgramData,
+        ProgramError,
+    },
+    sys::{bpf_link_create, bpf_prog_attach, kernel_version},
+};
+
+/// A program that is called on socket creation, bind and release.
+///
+/// [`CgroupSock`] programs can be used to allow or deny socket creation from
+/// within a [cgroup], or they can be used to monitor and gather statistics.
+///
+/// [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.10.
+///
+/// # 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::{CgroupSock, CgroupSockAttachType};
+///
+/// let file = File::open("/sys/fs/cgroup/unified")?;
+/// let bind: &mut CgroupSock = bpf.program_mut("bind").unwrap().try_into()?;
+/// bind.load()?;
+/// bind.attach(file)?;
+/// # Ok::<(), Error>(())
+/// ```
+#[derive(Debug)]
+#[doc(alias = "BPF_PROG_TYPE_CGROUP_SOCK")]
+pub struct CgroupSock {
+    pub(crate) data: ProgramData<CgroupSockLink>,
+    pub(crate) attach_type: CgroupSockAttachType,
+}
+
+impl CgroupSock {
+    /// 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, &mut self.data)
+    }
+
+    /// Attaches the program to the given cgroup.
+    ///
+    /// The returned value can be used to detach, see [CgroupSock::detach].
+    pub fn attach<T: AsRawFd>(&mut self, cgroup: T) -> Result<CgroupSockLinkId, 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(CgroupSockLink(CgroupSockLinkInner::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(CgroupSockLink(CgroupSockLinkInner::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 take_link(
+        &mut self,
+        link_id: CgroupSockLinkId,
+    ) -> Result<OwnedLink<CgroupSockLink>, ProgramError> {
+        Ok(OwnedLink::new(self.data.take_link(link_id)?))
+    }
+
+    /// Detaches the program.
+    ///
+    /// See [CgroupSock::attach].
+    pub fn detach(&mut self, link_id: CgroupSockLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
+}
+
+#[derive(Debug, Hash, Eq, PartialEq)]
+enum CgroupSockLinkIdInner {
+    Fd(<FdLink as Link>::Id),
+    ProgAttach(<ProgAttachLink as Link>::Id),
+}
+
+#[derive(Debug)]
+enum CgroupSockLinkInner {
+    Fd(FdLink),
+    ProgAttach(ProgAttachLink),
+}
+
+impl Link for CgroupSockLinkInner {
+    type Id = CgroupSockLinkIdInner;
+
+    fn id(&self) -> Self::Id {
+        match self {
+            CgroupSockLinkInner::Fd(fd) => CgroupSockLinkIdInner::Fd(fd.id()),
+            CgroupSockLinkInner::ProgAttach(p) => CgroupSockLinkIdInner::ProgAttach(p.id()),
+        }
+    }
+
+    fn detach(self) -> Result<(), ProgramError> {
+        match self {
+            CgroupSockLinkInner::Fd(fd) => fd.detach(),
+            CgroupSockLinkInner::ProgAttach(p) => p.detach(),
+        }
+    }
+}
+
+define_link_wrapper!(
+    /// The link used by [CgroupSock] programs.
+    CgroupSockLink,
+    /// The type returned by [CgroupSock::attach]. Can be passed to [CgroupSock::detach].
+    CgroupSockLinkId,
+    CgroupSockLinkInner,
+    CgroupSockLinkIdInner
+);
+
+/// Defines where to attach a [`CgroupSock`] program.
+#[derive(Copy, Clone, Debug)]
+pub enum CgroupSockAttachType {
+    /// Called after the IPv4 bind events.
+    PostBind4,
+    /// Called after the IPv6 bind events.
+    PostBind6,
+    /// Attach to IPv4 connect events.
+    SockCreate,
+    /// Attach to IPv6 connect events.
+    SockRelease,
+}
+
+impl Default for CgroupSockAttachType {
+    // The kernel checks for a 0 attach_type and sets it to sock_create
+    // We may as well do that here also
+    fn default() -> Self {
+        CgroupSockAttachType::SockCreate
+    }
+}
+
+impl From<CgroupSockAttachType> for bpf_attach_type {
+    fn from(s: CgroupSockAttachType) -> bpf_attach_type {
+        match s {
+            CgroupSockAttachType::PostBind4 => bpf_attach_type::BPF_CGROUP_INET4_POST_BIND,
+            CgroupSockAttachType::PostBind6 => bpf_attach_type::BPF_CGROUP_INET6_POST_BIND,
+            CgroupSockAttachType::SockCreate => bpf_attach_type::BPF_CGROUP_INET_SOCK_CREATE,
+            CgroupSockAttachType::SockRelease => bpf_attach_type::BPF_CGROUP_INET_SOCK_RELEASE,
+        }
+    }
+}
+
+#[derive(Debug, Error)]
+#[error("{0} is not a valid attach type for a CGROUP_SOCK program")]
+pub(crate) struct InvalidAttachType(String);
+
+impl CgroupSockAttachType {
+    pub(crate) fn try_from(value: &str) -> Result<CgroupSockAttachType, InvalidAttachType> {
+        match value {
+            "post_bind4" => Ok(CgroupSockAttachType::PostBind4),
+            "post_bind6" => Ok(CgroupSockAttachType::PostBind6),
+            "sock_create" => Ok(CgroupSockAttachType::SockCreate),
+            "sock_release" => Ok(CgroupSockAttachType::SockRelease),
+            _ => Err(InvalidAttachType(value.to_owned())),
+        }
+    }
+}

aya/src/programs/mod.rs

@@ -37,6 +37,7 @@
 //! [`Bpf::program_mut`]: crate::Bpf::program_mut
 //! [`maps`]: crate::maps
 pub mod cgroup_skb;
+pub mod cgroup_sock;
 pub mod cgroup_sock_addr;
 pub mod cgroup_sockopt;
 pub mod cgroup_sysctl;
@@ -74,6 +75,7 @@ use std::{
 use thiserror::Error;
 
 pub use cgroup_skb::{CgroupSkb, CgroupSkbAttachType};
+pub use cgroup_sock::{CgroupSock, CgroupSockAttachType};
 pub use cgroup_sock_addr::{CgroupSockAddr, CgroupSockAddrAttachType};
 pub use cgroup_sockopt::{CgroupSockopt, CgroupSockoptAttachType};
 pub use cgroup_sysctl::CgroupSysctl;
@@ -265,6 +267,8 @@ pub enum Program {
     Extension(Extension),
     /// A [`SkLookup`] program
     SkLookup(SkLookup),
+    /// A [`CgroupSock`] program
+    CgroupSock(CgroupSock),
 }
 
 impl Program {
@@ -294,6 +298,7 @@ impl Program {
             Program::Extension(_) => BPF_PROG_TYPE_EXT,
             Program::CgroupSockAddr(_) => BPF_PROG_TYPE_CGROUP_SOCK_ADDR,
             Program::SkLookup(_) => BPF_PROG_TYPE_SK_LOOKUP,
+            Program::CgroupSock(_) => BPF_PROG_TYPE_CGROUP_SOCK,
         }
     }
 
@@ -322,6 +327,7 @@ impl Program {
             Program::Extension(p) => p.data.pin(path),
             Program::CgroupSockAddr(p) => p.data.pin(path),
             Program::SkLookup(p) => p.data.pin(path),
+            Program::CgroupSock(p) => p.data.pin(path),
         }
     }
 
@@ -350,6 +356,7 @@ impl Program {
             Program::Extension(p) => p.unload(),
             Program::CgroupSockAddr(p) => p.unload(),
             Program::SkLookup(p) => p.unload(),
+            Program::CgroupSock(p) => p.unload(),
         }
     }
 }
@@ -573,6 +580,7 @@ impl ProgramFd for Program {
             Program::Extension(p) => p.data.fd,
             Program::CgroupSockAddr(p) => p.data.fd,
             Program::SkLookup(p) => p.data.fd,
+            Program::CgroupSock(p) => p.data.fd,
         }
     }
 }
@@ -622,7 +630,8 @@ impl_program_unload!(
     Extension,
     CgroupSockAddr,
     SkLookup,
-    SockOps
+    SockOps,
+    CgroupSock,
 );
 
 macro_rules! impl_program_fd {
@@ -665,6 +674,7 @@ impl_program_fd!(
     Extension,
     CgroupSockAddr,
     SkLookup,
+    CgroupSock,
 );
 
 macro_rules! impl_try_from_program {
@@ -718,6 +728,7 @@ impl_try_from_program!(
     Extension,
     CgroupSockAddr,
     SkLookup,
+    CgroupSock,
 );
 
 /// Provides information about a loaded program, like name, id and statistics

bpf/aya-bpf-macros/src/expand.rs

@@ -410,6 +410,51 @@ impl CgroupSockAddr {
     }
 }
 
+pub struct CgroupSock {
+    item: ItemFn,
+    attach_type: Option<String>,
+    name: Option<String>,
+}
+
+impl CgroupSock {
+    pub fn from_syn(mut args: Args, item: ItemFn) -> Result<CgroupSock> {
+        let name = pop_arg(&mut args, "name");
+        let attach_type = pop_arg(&mut args, "attach");
+        err_on_unknown_args(&args)?;
+
+        Ok(CgroupSock {
+            item,
+            attach_type,
+            name,
+        })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        let section_name = if let Some(name) = &self.name {
+            if let Some(attach_type) = &self.attach_type {
+                format!("cgroup/{}/{}", attach_type, name)
+            } else {
+                format!("cgroup/sock/{}", name)
+            }
+        } else if let Some(attach_type) = &self.attach_type {
+            format!("cgroup/{}", attach_type)
+        } else {
+            "cgroup/sock".to_string()
+        };
+        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) -> i32 {
+                return #fn_name(::aya_bpf::programs::SockContext::new(ctx));
+
+                #item
+            }
+        })
+    }
+}
+
 fn pop_arg(args: &mut Args, name: &str) -> Option<String> {
     match args.args.iter().position(|arg| arg.name == name) {
         Some(index) => Some(args.args.remove(index).value.value()),

bpf/aya-bpf-macros/src/lib.rs

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

bpf/aya-bpf/src/programs/mod.rs

@@ -7,6 +7,7 @@ pub mod raw_tracepoint;
 pub mod sk_buff;
 pub mod sk_lookup;
 pub mod sk_msg;
+pub mod sock;
 pub mod sock_addr;
 pub mod sock_ops;
 pub mod sockopt;
@@ -24,6 +25,7 @@ pub use raw_tracepoint::RawTracePointContext;
 pub use sk_buff::SkBuffContext;
 pub use sk_lookup::SkLookupContext;
 pub use sk_msg::SkMsgContext;
+pub use sock::SockContext;
 pub use sock_addr::SockAddrContext;
 pub use sock_ops::SockOpsContext;
 pub use sockopt::SockoptContext;

bpf/aya-bpf/src/programs/sock.rs

@@ -0,0 +1,19 @@
+use core::ffi::c_void;
+
+use crate::{bindings::bpf_sock, BpfContext};
+
+pub struct SockContext {
+    pub sock: *mut bpf_sock,
+}
+
+impl SockContext {
+    pub fn new(sock: *mut bpf_sock) -> SockContext {
+        SockContext { sock }
+    }
+}
+
+impl BpfContext for SockContext {
+    fn as_ptr(&self) -> *mut c_void {
+        self.sock as *mut _
+    }
+}