aya/aya-bpf: implement btf tracepoint programs

aya/src/bpf.rs

@@ -20,9 +20,9 @@ use crate::{
         Object, ParseError, ProgramSection,
     },
     programs::{
-        CgroupSkb, CgroupSkbAttachType, KProbe, LircMode2, Lsm, PerfEvent, ProbeKind, Program,
-        ProgramData, ProgramError, RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind,
-        SockOps, SocketFilter, TracePoint, UProbe, Xdp,
+        BtfTracePoint, CgroupSkb, CgroupSkbAttachType, KProbe, LircMode2, Lsm, PerfEvent,
+        ProbeKind, Program, ProgramData, ProgramError, RawTracePoint, SchedClassifier, SkMsg,
+        SkSkb, SkSkbKind, SockOps, SocketFilter, TracePoint, UProbe, Xdp,
     },
     sys::bpf_map_update_elem_ptr,
     util::{possible_cpus, POSSIBLE_CPUS},
@@ -306,6 +306,9 @@ impl<'a> BpfLoader<'a> {
                         Program::RawTracePoint(RawTracePoint { data })
                     }
                     ProgramSection::Lsm { .. } => Program::Lsm(Lsm { data }),
+                    ProgramSection::BtfTracePoint { .. } => {
+                        Program::BtfTracePoint(BtfTracePoint { data })
+                    }
                 };
 
                 (name, program)

aya/src/obj/mod.rs

@@ -87,6 +87,7 @@ pub enum ProgramSection {
     PerfEvent { name: String },
     RawTracePoint { name: String },
     Lsm { name: String },
+    BtfTracePoint { name: String },
 }
 
 impl ProgramSection {
@@ -110,6 +111,7 @@ impl ProgramSection {
             ProgramSection::PerfEvent { name } => name,
             ProgramSection::RawTracePoint { name } => name,
             ProgramSection::Lsm { name } => name,
+            ProgramSection::BtfTracePoint { name } => name,
         }
     }
 }
@@ -135,6 +137,7 @@ impl FromStr for ProgramSection {
             "uprobe" => UProbe { name },
             "uretprobe" => URetProbe { name },
             "xdp" => Xdp { name },
+            "tp_btf" => BtfTracePoint { name },
             _ if kind.starts_with("tracepoint") || kind.starts_with("tp") => {
                 // tracepoint sections are named `tracepoint/category/event_name`,
                 // and we want to parse the name as "category/event_name"
@@ -579,6 +582,7 @@ fn is_program_section(name: &str) -> bool {
         "raw_tp",
         "raw_tracepoint",
         "lsm",
+        "tp_btf",
     ] {
         if name.starts_with(prefix) {
             return true;
@@ -1046,4 +1050,21 @@ mod tests {
             })
         );
     }
+
+    #[test]
+    fn test_parse_section_btf_tracepoint() {
+        let mut obj = fake_obj();
+
+        assert_matches!(
+            obj.parse_section(fake_section("tp_btf/foo", bytes_of(&fake_ins()))),
+            Ok(())
+        );
+        assert_matches!(
+            obj.programs.get("foo"),
+            Some(Program {
+                section: ProgramSection::BtfTracePoint { .. },
+                ..
+            })
+        );
+    }
 }

aya/src/programs/mod.rs

@@ -49,6 +49,7 @@ mod sk_skb;
 mod sock_ops;
 mod socket_filter;
 pub mod tc;
+mod tp_btf;
 mod trace_point;
 mod uprobe;
 mod xdp;
@@ -79,6 +80,7 @@ pub use sk_skb::{SkSkb, SkSkbKind};
 pub use sock_ops::SockOps;
 pub use socket_filter::{SocketFilter, SocketFilterError};
 pub use tc::{SchedClassifier, TcAttachType, TcError};
+pub use tp_btf::{BtfTracePoint, BtfTracePointError};
 pub use trace_point::{TracePoint, TracePointError};
 pub use uprobe::{UProbe, UProbeError};
 pub use xdp::{Xdp, XdpError, XdpFlags};
@@ -171,6 +173,10 @@ pub enum ProgramError {
     /// An error occurred while working with a TC program.
     #[error(transparent)]
     TcError(#[from] TcError),
+
+    /// An error occurred while working with a BTF raw tracepoint program.
+    #[error(transparent)]
+    BtfTracePointError(#[from] BtfTracePointError),
 }
 
 pub trait ProgramFd {
@@ -194,6 +200,7 @@ pub enum Program {
     PerfEvent(PerfEvent),
     RawTracePoint(RawTracePoint),
     Lsm(Lsm),
+    BtfTracePoint(BtfTracePoint),
 }
 
 impl Program {
@@ -229,6 +236,7 @@ impl Program {
             Program::PerfEvent(_) => BPF_PROG_TYPE_PERF_EVENT,
             Program::RawTracePoint(_) => BPF_PROG_TYPE_RAW_TRACEPOINT,
             Program::Lsm(_) => BPF_PROG_TYPE_LSM,
+            Program::BtfTracePoint(_) => BPF_PROG_TYPE_TRACING,
         }
     }
 
@@ -258,6 +266,7 @@ impl Program {
             Program::PerfEvent(p) => &p.data,
             Program::RawTracePoint(p) => &p.data,
             Program::Lsm(p) => &p.data,
+            Program::BtfTracePoint(p) => &p.data,
         }
     }
 
@@ -277,6 +286,7 @@ impl Program {
             Program::PerfEvent(p) => &mut p.data,
             Program::RawTracePoint(p) => &mut p.data,
             Program::Lsm(p) => &mut p.data,
+            Program::BtfTracePoint(p) => &mut p.data,
         }
     }
 }
@@ -590,6 +600,7 @@ impl_program_fd!(
     PerfEvent,
     Lsm,
     RawTracePoint,
+    BtfTracePoint,
 );
 
 macro_rules! impl_try_from_program {
@@ -635,6 +646,7 @@ impl_try_from_program!(
     PerfEvent,
     Lsm,
     RawTracePoint,
+    BtfTracePoint,
 );
 
 /// Provides information about a loaded program, like name, id and statistics

aya/src/programs/tp_btf.rs

@@ -0,0 +1,103 @@
+//! BTF-enabled raw tracepoints.
+use std::os::unix::io::RawFd;
+
+use thiserror::Error;
+
+use crate::{
+    generated::{bpf_attach_type::BPF_TRACE_RAW_TP, bpf_prog_type::BPF_PROG_TYPE_TRACING},
+    obj::btf::{Btf, BtfError, BtfKind},
+    programs::{load_program, FdLink, LinkRef, ProgramData, ProgramError},
+    sys::bpf_raw_tracepoint_open,
+};
+
+/// Marks a function as a [BTF-enabled raw tracepoint][1] eBPF program that can be attached at
+/// a pre-defined kernel trace point.
+///
+/// The kernel provides a set of pre-defined trace points that eBPF programs can
+/// be attached to. See `/sys/kernel/debug/tracing/events` for a list of which
+/// events can be traced.
+///
+/// # Minimum kernel version
+///
+/// The minimum kernel version required to use this feature is 5.5.
+///
+/// # Examples
+///
+/// ```no_run
+/// # #[derive(thiserror::Error, Debug)]
+/// # enum Error {
+/// #     #[error(transparent)]
+/// #     BtfTracePointError(#[from] aya::programs::BtfTracePointError),
+/// #     #[error(transparent)]
+/// #     BtfError(#[from] aya::BtfError),
+/// #     #[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::BtfTracePoint, BtfError, Btf};
+/// use std::convert::TryInto;
+///
+/// let btf = Btf::from_sys_fs()?;
+/// let program: &mut BtfTracePoint = bpf.program_mut("sched_process_fork")?.try_into()?;
+/// program.load("sched_process_fork", &btf)?;
+/// program.attach()?;
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// [1]: https://github.com/torvalds/linux/commit/9e15db66136a14cde3f35691f1d839d950118826
+#[derive(Debug)]
+#[doc(alias = "BPF_TRACE_RAW_TP")]
+#[doc(alias = "BPF_PROG_TYPE_TRACING")]
+pub struct BtfTracePoint {
+    pub(crate) data: ProgramData,
+}
+
+/// Error type returned when loading LSM programs.
+#[derive(Debug, Error)]
+pub enum BtfTracePointError {
+    /// An error occured while working with BTF.
+    #[error(transparent)]
+    Btf(#[from] BtfError),
+}
+
+impl BtfTracePoint {
+    /// Loads the program inside the kernel.
+    ///
+    /// See also [`Program::load`](crate::programs::Program::load).
+    ///
+    /// # Arguments
+    ///
+    /// * `tracepoint` - full name of the tracepoint that we should attach to
+    /// * `btf` - btf information for the target system
+    pub fn load(&mut self, tracepoint: &str, btf: &Btf) -> Result<(), ProgramError> {
+        self.data.expected_attach_type = Some(BPF_TRACE_RAW_TP);
+        let type_name = format!("btf_trace_{}", tracepoint);
+        self.data.attach_btf_id = Some(
+            btf.id_by_type_name_kind(type_name.as_str(), BtfKind::Typedef)
+                .map_err(BtfTracePointError::from)?,
+        );
+        load_program(BPF_PROG_TYPE_TRACING, &mut self.data)
+    }
+
+    /// Returns the name of the program.
+    pub fn name(&self) -> String {
+        self.data.name.to_string()
+    }
+
+    /// Attaches the program.
+    pub fn attach(&mut self) -> Result<LinkRef, ProgramError> {
+        let prog_fd = self.data.fd_or_err()?;
+
+        // BTF programs specify their attach name at program load time
+        let pfd = bpf_raw_tracepoint_open(None, prog_fd).map_err(|(_code, io_error)| {
+            ProgramError::SyscallError {
+                call: "bpf_raw_tracepoint_open".to_owned(),
+                io_error,
+            }
+        })? as RawFd;
+
+        Ok(self.data.link(FdLink { fd: Some(pfd) }))
+    }
+}

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

@@ -412,3 +412,32 @@ impl Lsm {
         })
     }
 }
+
+pub struct BtfTracePoint {
+    item: ItemFn,
+    name: String,
+}
+
+impl BtfTracePoint {
+    pub fn from_syn(mut args: Args, item: ItemFn) -> Result<BtfTracePoint> {
+        let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string());
+
+        Ok(BtfTracePoint { item, name })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        let section_name = format!("tp_btf/{}", self.name);
+        let fn_name = &self.item.sig.ident;
+        let item = &self.item;
+        Ok(quote! {
+            #[no_mangle]
+            #[link_section = #section_name]
+            fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 {
+                let _ = #fn_name(::aya_bpf::programs::BtfTracePointContext::new(ctx));
+                return 0;
+
+                #item
+            }
+        })
+    }
+}

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

@@ -1,8 +1,8 @@
 mod expand;
 
 use expand::{
-    Args, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkMsg, SockOps,
-    TracePoint, Xdp,
+    Args, BtfTracePoint, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier,
+    SkMsg, SockOps, TracePoint, Xdp,
 };
 use proc_macro::TokenStream;
 use syn::{parse_macro_input, ItemFn, ItemStatic};
@@ -207,3 +207,44 @@ pub fn lsm(attrs: TokenStream, item: TokenStream) -> TokenStream {
         .unwrap_or_else(|err| err.to_compile_error())
         .into()
 }
+
+/// Marks a function as a [BTF-enabled raw tracepoint][1] eBPF program that can be attached at
+/// a pre-defined kernel trace point.
+///
+/// The kernel provides a set of pre-defined trace points that eBPF programs can
+/// be attached to. See `/sys/kernel/debug/tracing/events` for a list of which
+/// events can be traced.
+///
+/// # Minimum kernel version
+///
+/// The minimum kernel version required to use this feature is 5.5.
+///
+/// # Examples
+///
+/// ```no_run
+/// use aya_bpf::{macros::btf_tracepoint, programs::BtfTracePointContext};
+///
+/// #[btf_tracepoint(name = "sched_process_fork")]
+/// pub fn sched_process_fork(ctx: BtfTracePointContext) -> u32 {
+///     match unsafe { try_sched_process_fork(ctx) } {
+///         Ok(ret) => ret,
+///         Err(ret) => ret,
+///     }
+/// }
+///
+/// unsafe fn try_sched_process_fork(_ctx: BtfTracePointContext) -> Result<u32, u32> {
+///     Ok(0)
+/// }
+/// ```
+///
+/// [1]: https://github.com/torvalds/linux/commit/9e15db66136a14cde3f35691f1d839d950118826
+#[proc_macro_attribute]
+pub fn btf_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    let args = parse_macro_input!(attrs as Args);
+    let item = parse_macro_input!(item as ItemFn);
+
+    BtfTracePoint::from_syn(args, item)
+        .and_then(|u| u.expand())
+        .unwrap_or_else(|err| err.to_compile_error())
+        .into()
+}

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

@@ -5,6 +5,7 @@ pub mod raw_tracepoint;
 pub mod sk_msg;
 pub mod sk_skb;
 pub mod sock_ops;
+pub mod tp_btf;
 pub mod tracepoint;
 pub mod xdp;
 
@@ -15,5 +16,6 @@ pub use raw_tracepoint::RawTracePointContext;
 pub use sk_msg::SkMsgContext;
 pub use sk_skb::SkSkbContext;
 pub use sock_ops::SockOpsContext;
+pub use tp_btf::BtfTracePointContext;
 pub use tracepoint::TracePointContext;
 pub use xdp::XdpContext;

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

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