Merge pull request #531 from dave-tucker/probe-cookie

aya: Make FEATURES public

aya-obj/src/btf/btf.rs

@@ -166,12 +166,68 @@ pub enum BtfError {
 #[derive(Default, Debug)]
 #[allow(missing_docs)]
 pub struct BtfFeatures {
-    pub btf_func: bool,
-    pub btf_func_global: bool,
-    pub btf_datasec: bool,
-    pub btf_float: bool,
-    pub btf_decl_tag: bool,
-    pub btf_type_tag: bool,
+    btf_func: bool,
+    btf_func_global: bool,
+    btf_datasec: bool,
+    btf_float: bool,
+    btf_decl_tag: bool,
+    btf_type_tag: bool,
+}
+
+impl BtfFeatures {
+    #[doc(hidden)]
+    pub fn new(
+        btf_func: bool,
+        btf_func_global: bool,
+        btf_datasec: bool,
+        btf_float: bool,
+        btf_decl_tag: bool,
+        btf_type_tag: bool,
+    ) -> Self {
+        BtfFeatures {
+            btf_func,
+            btf_func_global,
+            btf_datasec,
+            btf_float,
+            btf_decl_tag,
+            btf_type_tag,
+        }
+    }
+
+    /// Returns true if the BTF_TYPE_FUNC is supported.
+    pub fn btf_func(&self) -> bool {
+        self.btf_func
+    }
+
+    /// Returns true if the BTF_TYPE_FUNC_GLOBAL is supported.
+    pub fn btf_func_global(&self) -> bool {
+        self.btf_func_global
+    }
+
+    /// Returns true if the BTF_TYPE_DATASEC is supported.
+    pub fn btf_datasec(&self) -> bool {
+        self.btf_datasec
+    }
+
+    /// Returns true if the BTF_FLOAT is supported.
+    pub fn btf_float(&self) -> bool {
+        self.btf_float
+    }
+
+    /// Returns true if the BTF_DECL_TAG is supported.
+    pub fn btf_decl_tag(&self) -> bool {
+        self.btf_decl_tag
+    }
+
+    /// Returns true if the BTF_TYPE_TAG is supported.
+    pub fn btf_type_tag(&self) -> bool {
+        self.btf_type_tag
+    }
+
+    /// Returns true if the BTF_KIND_FUNC_PROTO is supported.
+    pub fn btf_kind_func_proto(&self) -> bool {
+        self.btf_func && self.btf_decl_tag
+    }
 }
 
 /// Bpf Type Format metadata.

aya-obj/src/obj.rs

@@ -42,11 +42,63 @@ const KERNEL_VERSION_ANY: u32 = 0xFFFF_FFFE;
 #[derive(Default, Debug)]
 #[allow(missing_docs)]
 pub struct Features {
-    pub bpf_name: bool,
-    pub bpf_probe_read_kernel: bool,
-    pub bpf_perf_link: bool,
-    pub bpf_global_data: bool,
-    pub btf: Option<BtfFeatures>,
+    bpf_name: bool,
+    bpf_probe_read_kernel: bool,
+    bpf_perf_link: bool,
+    bpf_global_data: bool,
+    bpf_cookie: bool,
+    btf: Option<BtfFeatures>,
+}
+
+impl Features {
+    #[doc(hidden)]
+    pub fn new(
+        bpf_name: bool,
+        bpf_probe_read_kernel: bool,
+        bpf_perf_link: bool,
+        bpf_global_data: bool,
+        bpf_cookie: bool,
+        btf: Option<BtfFeatures>,
+    ) -> Self {
+        Self {
+            bpf_name,
+            bpf_probe_read_kernel,
+            bpf_perf_link,
+            bpf_global_data,
+            bpf_cookie,
+            btf,
+        }
+    }
+
+    /// Returns whether BPF program names are supported.
+    pub fn bpf_name(&self) -> bool {
+        self.bpf_name
+    }
+
+    /// Returns whether the bpf_probe_read_kernel helper is supported.
+    pub fn bpf_probe_read_kernel(&self) -> bool {
+        self.bpf_probe_read_kernel
+    }
+
+    /// Returns whether bpf_links are supported for Kprobes/Uprobes/Tracepoints.
+    pub fn bpf_perf_link(&self) -> bool {
+        self.bpf_perf_link
+    }
+
+    /// Returns whether BPF program global data is supported.
+    pub fn bpf_global_data(&self) -> bool {
+        self.bpf_global_data
+    }
+
+    /// Returns whether BPF program cookie is supported.
+    pub fn bpf_cookie(&self) -> bool {
+        self.bpf_cookie
+    }
+
+    /// If BTF is supported, returns which BTF features are supported.
+    pub fn btf(&self) -> Option<&BtfFeatures> {
+        self.btf.as_ref()
+    }
 }
 
 /// The loaded object file representation

aya-obj/src/relocation.rs

@@ -242,10 +242,7 @@ fn relocate_maps<'a, I: Iterator<Item = &'a Relocation>>(
             m
         } else {
             let Some(m) = maps_by_section.get(&section_index) else {
-                debug!(
-                    "failed relocating map by section index {}",
-                    section_index
-                );
+                debug!("failed relocating map by section index {}", section_index);
                 return Err(RelocationError::SectionNotFound {
                     symbol_index: rel.symbol_index,
                     symbol_name: sym.name.clone(),

aya/src/bpf.rs

@@ -33,11 +33,11 @@ use crate::{
         SkMsg, SkSkb, SkSkbKind, SockOps, SocketFilter, TracePoint, UProbe, Xdp,
     },
     sys::{
-        bpf_load_btf, bpf_map_freeze, bpf_map_update_elem_ptr, is_bpf_global_data_supported,
-        is_btf_datasec_supported, is_btf_decl_tag_supported, is_btf_float_supported,
-        is_btf_func_global_supported, is_btf_func_supported, is_btf_supported,
-        is_btf_type_tag_supported, is_perf_link_supported, is_probe_read_kernel_supported,
-        is_prog_name_supported, retry_with_verifier_logs,
+        bpf_load_btf, bpf_map_freeze, bpf_map_update_elem_ptr, is_bpf_cookie_supported,
+        is_bpf_global_data_supported, is_btf_datasec_supported, is_btf_decl_tag_supported,
+        is_btf_float_supported, is_btf_func_global_supported, is_btf_func_supported,
+        is_btf_supported, is_btf_type_tag_supported, is_perf_link_supported,
+        is_probe_read_kernel_supported, is_prog_name_supported, retry_with_verifier_logs,
     },
     util::{bytes_of, bytes_of_slice, possible_cpus, VerifierLog, POSSIBLE_CPUS},
 };
@@ -72,28 +72,34 @@ lazy_static! {
 
 fn detect_features() -> Features {
     let btf = if is_btf_supported() {
-        Some(BtfFeatures {
-            btf_func: is_btf_func_supported(),
-            btf_func_global: is_btf_func_global_supported(),
-            btf_datasec: is_btf_datasec_supported(),
-            btf_float: is_btf_float_supported(),
-            btf_decl_tag: is_btf_decl_tag_supported(),
-            btf_type_tag: is_btf_type_tag_supported(),
-        })
+        Some(BtfFeatures::new(
+            is_btf_func_supported(),
+            is_btf_func_global_supported(),
+            is_btf_datasec_supported(),
+            is_btf_float_supported(),
+            is_btf_decl_tag_supported(),
+            is_btf_type_tag_supported(),
+        ))
     } else {
         None
     };
-    let f = Features {
-        bpf_name: is_prog_name_supported(),
-        bpf_probe_read_kernel: is_probe_read_kernel_supported(),
-        bpf_perf_link: is_perf_link_supported(),
-        bpf_global_data: is_bpf_global_data_supported(),
+    let f = Features::new(
+        is_prog_name_supported(),
+        is_probe_read_kernel_supported(),
+        is_perf_link_supported(),
+        is_bpf_global_data_supported(),
+        is_bpf_cookie_supported(),
         btf,
-    };
+    );
     debug!("BPF Feature Detection: {:#?}", f);
     f
 }
 
+/// Returns a reference to the detected BPF features.
+pub fn features() -> &'static Features {
+    &FEATURES
+}
+
 /// Builder style API for advanced loading of eBPF programs.
 ///
 /// Loading eBPF code involves a few steps, including loading maps and applying
@@ -347,7 +353,7 @@ impl<'a> BpfLoader<'a> {
         let mut obj = Object::parse(data)?;
         obj.patch_map_data(self.globals.clone())?;
 
-        let btf_fd = if let Some(features) = &FEATURES.btf {
+        let btf_fd = if let Some(features) = &FEATURES.btf() {
             if let Some(btf) = obj.fixup_and_sanitize_btf(features)? {
                 // load btf to the kernel
                 Some(load_btf(btf.to_bytes())?)
@@ -364,7 +370,7 @@ impl<'a> BpfLoader<'a> {
         let mut maps = HashMap::new();
         for (name, mut obj) in obj.maps.drain() {
             if let (false, BpfSectionKind::Bss | BpfSectionKind::Data | BpfSectionKind::Rodata) =
-                (FEATURES.bpf_global_data, obj.section_kind())
+                (FEATURES.bpf_global_data(), obj.section_kind())
             {
                 continue;
             }
@@ -452,7 +458,7 @@ impl<'a> BpfLoader<'a> {
             .map(|(name, prog_obj)| {
                 let function_obj = obj.functions.get(&prog_obj.function_key()).unwrap().clone();
 
-                let prog_name = if FEATURES.bpf_name {
+                let prog_name = if FEATURES.bpf_name() {
                     Some(name.clone())
                 } else {
                     None

aya/src/programs/perf_attach.rs

@@ -73,7 +73,7 @@ impl Link for PerfLink {
 }
 
 pub(crate) fn perf_attach(prog_fd: RawFd, fd: RawFd) -> Result<PerfLinkInner, ProgramError> {
-    if FEATURES.bpf_perf_link {
+    if FEATURES.bpf_perf_link() {
         let link_fd =
             bpf_link_create(prog_fd, fd, BPF_PERF_EVENT, None, 0).map_err(|(_, io_error)| {
                 ProgramError::SyscallError {

aya/src/sys/bpf.rs

@@ -719,6 +719,33 @@ pub(crate) fn is_bpf_global_data_supported() -> bool {
     false
 }
 
+pub(crate) fn is_bpf_cookie_supported() -> bool {
+    let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
+    let u = unsafe { &mut attr.__bindgen_anon_3 };
+
+    let prog: &[u8] = &[
+        0x85, 0x00, 0x00, 0x00, 0xae, 0x00, 0x00, 0x00, // call bpf_get_attach_cookie
+        0x95, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // exit
+    ];
+
+    let gpl = b"GPL\0";
+    u.license = gpl.as_ptr() as u64;
+
+    let insns = copy_instructions(prog).unwrap();
+    u.insn_cnt = insns.len() as u32;
+    u.insns = insns.as_ptr() as u64;
+    u.prog_type = bpf_prog_type::BPF_PROG_TYPE_KPROBE as u32;
+
+    match sys_bpf(bpf_cmd::BPF_PROG_LOAD, &attr) {
+        Ok(v) => {
+            let fd = v as RawFd;
+            unsafe { close(fd) };
+            true
+        }
+        Err(_) => false,
+    }
+}
+
 pub(crate) fn is_btf_supported() -> bool {
     let mut btf = Btf::new();
     let name_offset = btf.add_string("int".to_string());