[216] Add Support For Flow Dissector Programs

aya-bpf-macros/src/flow_dissector.rs

@@ -0,0 +1,66 @@
+use proc_macro2::TokenStream;
+use proc_macro_error::abort;
+use quote::quote;
+use syn::{ItemFn, Result};
+
+pub(crate) struct FlowDissector {
+    item: ItemFn,
+}
+
+impl FlowDissector {
+    pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self> {
+        if !attrs.is_empty() {
+            abort!(attrs, "unexpected attribute")
+        }
+        let item = syn::parse2(item)?;
+        Ok(FlowDissector { item })
+    }
+
+    pub(crate) fn expand(&self) -> Result<TokenStream> {
+        let fn_name = self.item.sig.ident.clone();
+        let fn_vis = &self.item.vis;
+        let item = &self.item;
+        Ok(quote! {
+            #[no_mangle]
+            #[link_section = "flow_dissector"]
+            #fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> u32 {
+                return #fn_name(::aya_bpf::programs::FlowDissectorContext::new(ctx));
+
+                #item
+            }
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use syn::parse_quote;
+
+    use super::*;
+
+    #[test]
+    fn test_flow_dissector() {
+        let prog = FlowDissector::parse(
+            parse_quote! {},
+            parse_quote! {
+                fn prog(ctx: &mut ::aya_bpf::programs::FlowDissectorContext) -> u32 {
+                    0
+                }
+            },
+        )
+        .unwrap();
+        let expanded = prog.expand().unwrap();
+        let expected = quote! {
+            #[no_mangle]
+            #[link_section = "flow_dissector"]
+            fn prog(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> u32 {
+                return prog(::aya_bpf::programs::FlowDissectorContext::new(ctx));
+
+                fn prog(ctx: &mut ::aya_bpf::programs::FlowDissectorContext) -> u32 {
+                    0
+                }
+            }
+        };
+        assert_eq!(expected.to_string(), expanded.to_string());
+    }
+}

aya-bpf-macros/src/lib.rs

@@ -8,6 +8,7 @@ mod cgroup_sockopt;
 mod cgroup_sysctl;
 mod fentry;
 mod fexit;
+mod flow_dissector;
 mod kprobe;
 mod lsm;
 mod map;
@@ -32,6 +33,7 @@ use cgroup_sockopt::CgroupSockopt;
 use cgroup_sysctl::CgroupSysctl;
 use fentry::FEntry;
 use fexit::FExit;
+use flow_dissector::FlowDissector;
 use kprobe::{KProbe, KProbeKind};
 use lsm::Lsm;
 use map::Map;
@@ -605,6 +607,21 @@ pub fn fexit(attrs: TokenStream, item: TokenStream) -> TokenStream {
     }
 }
 
+/// Marks a function as an eBPF Flow Dissector program that can be attached to
+/// a network namespace.
+///
+#[proc_macro_error]
+#[proc_macro_attribute]
+pub fn flow_dissector(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    match FlowDissector::parse(attrs.into(), item.into()) {
+        Ok(prog) => prog
+            .expand()
+            .unwrap_or_else(|err| abort!(err.span(), "{}", err))
+            .into(),
+        Err(err) => abort!(err.span(), "{}", err),
+    }
+}
+
 /// Marks a function as an eBPF Socket Lookup program that can be attached to
 /// a network namespace.
 ///

aya-obj/src/obj.rs

@@ -264,6 +264,7 @@ pub enum ProgramSection {
     FExit {
         sleepable: bool,
     },
+    FlowDissector,
     Extension,
     SkLookup,
     CgroupSock {
@@ -419,6 +420,7 @@ impl FromStr for ProgramSection {
             "fentry.s" => FEntry { sleepable: true },
             "fexit" => FExit { sleepable: false },
             "fexit.s" => FExit { sleepable: true },
+            "flow_dissector" => FlowDissector,
             "freplace" => Extension,
             "sk_lookup" => SkLookup,
             _ => {

aya/src/bpf.rs

@@ -31,9 +31,10 @@ use crate::{
     },
     programs::{
         BtfTracePoint, CgroupDevice, 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,
+        CgroupSockopt, CgroupSysctl, Extension, FEntry, FExit, FlowDissector, KProbe, LircMode2,
+        Lsm, PerfEvent, ProbeKind, Program, ProgramData, ProgramError, RawTracePoint,
+        SchedClassifier, SkLookup, SkMsg, SkSkb, SkSkbKind, SockOps, SocketFilter, TracePoint,
+        UProbe, Xdp,
     },
     sys::{
         bpf_load_btf, is_bpf_cookie_supported, is_bpf_global_data_supported,
@@ -434,6 +435,7 @@ impl<'a> BpfLoader<'a> {
                                 | ProgramSection::PerfEvent
                                 | ProgramSection::RawTracePoint
                                 | ProgramSection::SkLookup
+                                | ProgramSection::FlowDissector
                                 | ProgramSection::CgroupSock { attach_type: _ }
                                 | ProgramSection::CgroupDevice => {}
                             }
@@ -666,6 +668,9 @@ impl<'a> BpfLoader<'a> {
                             }
                             Program::FExit(FExit { data })
                         }
+                        ProgramSection::FlowDissector => Program::FlowDissector(FlowDissector {
+                            data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
+                        }),
                         ProgramSection::Extension => Program::Extension(Extension {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
                         }),

aya/src/programs/flow_dissector.rs

@@ -0,0 +1,109 @@
+//! Flow dissector programs.
+
+use std::os::fd::AsFd;
+
+use crate::{
+    generated::{bpf_attach_type::BPF_FLOW_DISSECTOR, bpf_prog_type::BPF_PROG_TYPE_FLOW_DISSECTOR},
+    programs::{define_link_wrapper, load_program, FdLink, FdLinkId, ProgramData, ProgramError},
+    sys::{bpf_link_create, LinkTarget, SyscallError},
+};
+
+/// A program that can be attached as a Flow Dissector routine
+///
+/// ['FlowDissector'] programs operate on an __sk_buff.
+/// However, only the limited set of fields is allowed: data, data_end and flow_keys.
+/// flow_keys is struct bpf_flow_keys and contains flow dissector input and output arguments.
+///
+/// # Minimum kernel version
+///
+/// The minimum kernel version required to use this feature is 4.2.
+///
+/// # Examples
+///
+/// ```no_run
+/// # #[derive(Debug, thiserror::Error)]
+/// # 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 = Bpf::load_file("ebpf_programs.o")?;
+/// use aya::{Bpf, programs::FlowDissector};
+/// use std::fs::File;
+///
+/// let program: &mut FlowDissector = bpf.program_mut("filename_lookup").unwrap().try_into()?;
+/// program.load()?;
+///
+/// let net_ns = File::open("/proc/self/ns/net")?;
+/// program.attach(net_ns)?;
+/// # Ok::<(), Error>(())
+/// ```
+#[derive(Debug)]
+#[doc(alias = "BPF_PROG_TYPE_FLOW_DISSECTOR")]
+pub struct FlowDissector {
+    pub(crate) data: ProgramData<FlowDissectorLink>,
+}
+
+impl FlowDissector {
+    /// Loads the program inside the kernel.
+    pub fn load(&mut self) -> Result<(), ProgramError> {
+        self.data.expected_attach_type = Some(BPF_FLOW_DISSECTOR);
+        load_program(BPF_PROG_TYPE_FLOW_DISSECTOR, &mut self.data)
+    }
+
+    /// Attaches the program to the given network namespace.
+    ///
+    /// The returned value can be used to detach, see [FlowDissector::detach].
+    pub fn attach<T: AsFd>(&mut self, netns: T) -> Result<FlowDissectorLinkId, ProgramError> {
+        let prog_fd = self.fd()?;
+        let prog_fd = prog_fd.as_fd();
+        let netns_fd = netns.as_fd();
+
+        let link_fd = bpf_link_create(
+            prog_fd,
+            LinkTarget::Fd(netns_fd),
+            BPF_FLOW_DISSECTOR,
+            None,
+            0,
+        )
+        .map_err(|(_, io_error)| SyscallError {
+            call: "bpf_link_create",
+            io_error,
+        })?;
+        self.data
+            .links
+            .insert(FlowDissectorLink::new(FdLink::new(link_fd)))
+    }
+
+    /// Detaches the program.
+    ///
+    /// See [FlowDissector::attach].
+    pub fn detach(&mut self, link_id: FlowDissectorLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
+
+    /// 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: FlowDissectorLinkId,
+    ) -> Result<FlowDissectorLink, ProgramError> {
+        self.data.take_link(link_id)
+    }
+}
+
+define_link_wrapper!(
+    /// The link used by [FlowDissector] programs.
+    FlowDissectorLink,
+    /// The type returned by [FlowDissector::attach]. Can be passed to [FlowDissector::detach].
+    FlowDissectorLinkId,
+    FdLink,
+    FdLinkId
+);

aya/src/programs/mod.rs

@@ -44,6 +44,7 @@ pub mod cgroup_sysctl;
 pub mod extension;
 pub mod fentry;
 pub mod fexit;
+pub mod flow_dissector;
 pub mod kprobe;
 pub mod links;
 pub mod lirc_mode2;
@@ -83,6 +84,7 @@ pub use cgroup_sysctl::CgroupSysctl;
 pub use extension::{Extension, ExtensionError};
 pub use fentry::FEntry;
 pub use fexit::FExit;
+pub use flow_dissector::FlowDissector;
 pub use kprobe::{KProbe, KProbeError};
 use libc::ENOSPC;
 pub use links::Link;
@@ -276,6 +278,8 @@ pub enum Program {
     FEntry(FEntry),
     /// A [`FExit`] program
     FExit(FExit),
+    /// A [`FlowDissector`] program
+    FlowDissector(FlowDissector),
     /// A [`Extension`] program
     Extension(Extension),
     /// A [`SkLookup`] program
@@ -310,6 +314,7 @@ impl Program {
             Self::BtfTracePoint(_) => BPF_PROG_TYPE_TRACING,
             Self::FEntry(_) => BPF_PROG_TYPE_TRACING,
             Self::FExit(_) => BPF_PROG_TYPE_TRACING,
+            Self::FlowDissector(_) => BPF_PROG_TYPE_FLOW_DISSECTOR,
             Self::Extension(_) => BPF_PROG_TYPE_EXT,
             Self::CgroupSockAddr(_) => BPF_PROG_TYPE_CGROUP_SOCK_ADDR,
             Self::SkLookup(_) => BPF_PROG_TYPE_SK_LOOKUP,
@@ -340,6 +345,7 @@ impl Program {
             Self::BtfTracePoint(p) => p.pin(path),
             Self::FEntry(p) => p.pin(path),
             Self::FExit(p) => p.pin(path),
+            Self::FlowDissector(p) => p.pin(path),
             Self::Extension(p) => p.pin(path),
             Self::CgroupSockAddr(p) => p.pin(path),
             Self::SkLookup(p) => p.pin(path),
@@ -370,6 +376,7 @@ impl Program {
             Self::BtfTracePoint(mut p) => p.unload(),
             Self::FEntry(mut p) => p.unload(),
             Self::FExit(mut p) => p.unload(),
+            Self::FlowDissector(mut p) => p.unload(),
             Self::Extension(mut p) => p.unload(),
             Self::CgroupSockAddr(mut p) => p.unload(),
             Self::SkLookup(mut p) => p.unload(),
@@ -402,6 +409,7 @@ impl Program {
             Self::BtfTracePoint(p) => p.fd(),
             Self::FEntry(p) => p.fd(),
             Self::FExit(p) => p.fd(),
+            Self::FlowDissector(p) => p.fd(),
             Self::Extension(p) => p.fd(),
             Self::CgroupSockAddr(p) => p.fd(),
             Self::SkLookup(p) => p.fd(),
@@ -435,6 +443,7 @@ impl Program {
             Self::BtfTracePoint(p) => p.info(),
             Self::FEntry(p) => p.info(),
             Self::FExit(p) => p.info(),
+            Self::FlowDissector(p) => p.info(),
             Self::Extension(p) => p.info(),
             Self::CgroupSockAddr(p) => p.info(),
             Self::SkLookup(p) => p.info(),
@@ -748,6 +757,7 @@ impl_program_unload!(
     BtfTracePoint,
     FEntry,
     FExit,
+    FlowDissector,
     Extension,
     CgroupSockAddr,
     SkLookup,
@@ -788,6 +798,7 @@ impl_fd!(
     BtfTracePoint,
     FEntry,
     FExit,
+    FlowDissector,
     Extension,
     CgroupSockAddr,
     SkLookup,
@@ -842,6 +853,7 @@ impl_program_pin!(
     BtfTracePoint,
     FEntry,
     FExit,
+    FlowDissector,
     Extension,
     CgroupSockAddr,
     SkLookup,
@@ -882,6 +894,7 @@ impl_from_pin!(
     BtfTracePoint,
     FEntry,
     FExit,
+    FlowDissector,
     Extension,
     SkLookup,
     SockOps,
@@ -936,6 +949,7 @@ impl_try_from_program!(
     BtfTracePoint,
     FEntry,
     FExit,
+    FlowDissector,
     Extension,
     CgroupSockAddr,
     SkLookup,
@@ -982,6 +996,7 @@ impl_info!(
     BtfTracePoint,
     FEntry,
     FExit,
+    FlowDissector,
     Extension,
     CgroupSockAddr,
     SkLookup,