libbpf-rs: Plumb through support for USDT options

When attaching to USDT, users can optionally provide some option struct
object that may change the behavior or provide additional data. So far
we have not supported doing that in libbpf-rs, but with this change we
add such support.

Closes: #308
Signed-off-by: Daniel Müller <deso@posteo.net>

Author: d-e-s-o

libbpf-rs/src/lib.rs

@@ -97,7 +97,7 @@ pub use crate::map::{Map, MapFlags, MapType, OpenMap};
 pub use crate::object::{Object, ObjectBuilder, OpenObject};
 pub use crate::perf_buffer::{PerfBuffer, PerfBufferBuilder};
 pub use crate::print::{get_print, set_print, PrintCallback, PrintLevel};
-pub use crate::program::{OpenProgram, Program, ProgramAttachType, ProgramType};
+pub use crate::program::{OpenProgram, Program, ProgramAttachType, ProgramType, UsdtOpts};
 pub use crate::ringbuf::{RingBuffer, RingBufferBuilder};
 pub use crate::tc::{
     TcAttachPoint, TcHook, TcHookBuilder, TC_CUSTOM, TC_EGRESS, TC_H_CLSACT, TC_H_INGRESS,

libbpf-rs/src/program.rs

@@ -1,5 +1,6 @@
 use std::convert::TryFrom;
 use std::ffi::CStr;
+use std::mem;
 use std::path::Path;
 use std::ptr;
 
@@ -9,6 +10,28 @@ use strum_macros::Display;
 
 use crate::*;
 
+/// Options to optionally be provided when attaching to a USDT.
+#[derive(Clone, Debug, Default)]
+pub struct UsdtOpts {
+    /// Custom user-provided value accessible through `bpf_usdt_cookie`.
+    pub cookie: u64,
+    #[doc(hidden)]
+    pub _non_exhaustive: (),
+}
+
+impl From<UsdtOpts> for libbpf_sys::bpf_usdt_opts {
+    fn from(opts: UsdtOpts) -> Self {
+        let UsdtOpts {
+            cookie,
+            _non_exhaustive,
+        } = opts;
+        libbpf_sys::bpf_usdt_opts {
+            sz: mem::size_of::<Self>() as u64,
+            usdt_cookie: cookie,
+        }
+    }
+}
+
 /// Represents a parsed but not yet loaded BPF program.
 ///
 /// This object exposes operations that need to happen before the program is loaded.
@@ -488,30 +511,34 @@ impl Program {
         }
     }
 
-    /// Attach this program to a [USDT](https://lwn.net/Articles/753601/) probe
-    /// point. The entry point of the program must be defined with
-    /// `SEC("usdt")`.
-    pub fn attach_usdt(
+    fn attach_usdt_impl(
         &mut self,
         pid: i32,
-        binary_path: impl AsRef<Path>,
-        usdt_provider: impl AsRef<str>,
-        usdt_name: impl AsRef<str>,
+        binary_path: &Path,
+        usdt_provider: &str,
+        usdt_name: &str,
+        usdt_opts: Option<UsdtOpts>,
     ) -> Result<Link> {
-        let path = util::path_to_cstring(binary_path.as_ref())?;
+        let path = util::path_to_cstring(binary_path)?;
         let path_ptr = path.as_ptr();
-        let usdt_provider = util::str_to_cstring(usdt_provider.as_ref())?;
+        let usdt_provider = util::str_to_cstring(usdt_provider)?;
         let usdt_provider_ptr = usdt_provider.as_ptr();
-        let usdt_name = util::str_to_cstring(usdt_name.as_ref())?;
+        let usdt_name = util::str_to_cstring(usdt_name)?;
         let usdt_name_ptr = usdt_name.as_ptr();
+        let usdt_opts = usdt_opts.map(libbpf_sys::bpf_usdt_opts::from);
+        let usdt_opts_ptr = usdt_opts
+            .as_ref()
+            .map(|opts| opts as *const _)
+            .unwrap_or_else(ptr::null);
+
         let ptr = unsafe {
             libbpf_sys::bpf_program__attach_usdt(
                 self.ptr,
                 pid,
                 path_ptr,
                 usdt_provider_ptr,
                 usdt_name_ptr,
-                ptr::null(),
+                usdt_opts_ptr,
             )
         };
         let err = unsafe { libbpf_sys::libbpf_get_error(ptr as *const _) };
@@ -522,6 +549,45 @@ impl Program {
         }
     }
 
+    /// Attach this program to a [USDT](https://lwn.net/Articles/753601/) probe
+    /// point. The entry point of the program must be defined with
+    /// `SEC("usdt")`.
+    pub fn attach_usdt(
+        &mut self,
+        pid: i32,
+        binary_path: impl AsRef<Path>,
+        usdt_provider: impl AsRef<str>,
+        usdt_name: impl AsRef<str>,
+    ) -> Result<Link> {
+        self.attach_usdt_impl(
+            pid,
+            binary_path.as_ref(),
+            usdt_provider.as_ref(),
+            usdt_name.as_ref(),
+            None,
+        )
+    }
+
+    /// Attach this program to a [USDT](https://lwn.net/Articles/753601/) probe
+    /// point, providing additional options. The entry point of the program must
+    /// be defined with `SEC("usdt")`.
+    pub fn attach_usdt_with_opts(
+        &mut self,
+        pid: i32,
+        binary_path: impl AsRef<Path>,
+        usdt_provider: impl AsRef<str>,
+        usdt_name: impl AsRef<str>,
+        usdt_opts: UsdtOpts,
+    ) -> Result<Link> {
+        self.attach_usdt_impl(
+            pid,
+            binary_path.as_ref(),
+            usdt_provider.as_ref(),
+            usdt_name.as_ref(),
+            Some(usdt_opts),
+        )
+    }
+
     /// Returns the number of instructions that form the program.
     ///
     /// Please see note in [`OpenProgram::insn_cnt`].

libbpf-rs/tests/bin/src/usdt.bpf.c

@@ -23,4 +23,18 @@ int handle__usdt(void *ctx)
     return 0;
 }
 
+SEC("usdt")
+int handle__usdt_with_cookie(void *ctx)
+{
+    int *value;
+
+    value = bpf_ringbuf_reserve(&ringbuf, sizeof(int), 0);
+    if (value) {
+        *value = bpf_usdt_cookie(ctx);
+        bpf_ringbuf_submit(value, 0);
+    }
+
+    return 0;
+}
+
 char LICENSE[] SEC("license") = "GPL";

libbpf-rs/tests/bin/usdt.bpf.o

@@ -13,7 +13,8 @@ use probe::probe;
 use scopeguard::defer;
 
 use libbpf_rs::{
-    num_possible_cpus, Iter, Map, MapFlags, MapType, Object, ObjectBuilder, OpenObject, ProgramType,
+    num_possible_cpus, Iter, Map, MapFlags, MapType, Object, ObjectBuilder, OpenObject,
+    ProgramType, UsdtOpts,
 };
 
 fn get_test_object_path(filename: &str) -> PathBuf {
@@ -783,6 +784,56 @@ fn test_object_usdt() {
     unsafe { assert_eq!(V, 1) };
 }
 
+#[test]
+fn test_object_usdt_cookie() {
+    bump_rlimit_mlock();
+
+    let cookie_val = 1337u16;
+    let mut obj = get_test_object("usdt.bpf.o");
+    let prog = obj
+        .prog_mut("handle__usdt_with_cookie")
+        .expect("Failed to find program");
+
+    let path = std::env::current_exe().expect("Failed to find executable name");
+    let _link = prog
+        .attach_usdt_with_opts(
+            unsafe { libc::getpid() },
+            &path,
+            "test_provider",
+            "test_function",
+            UsdtOpts {
+                cookie: cookie_val.into(),
+                ..UsdtOpts::default()
+            },
+        )
+        .expect("Failed to attach prog");
+
+    let mut builder = libbpf_rs::RingBufferBuilder::new();
+    let map = obj.map("ringbuf").expect("Failed to get ringbuf map");
+
+    static mut V: i32 = 0;
+    fn callback(data: &[u8]) -> i32 {
+        let mut value: i32 = 0;
+        plain::copy_from_bytes(&mut value, data).expect("Wrong size");
+
+        unsafe {
+            V = value;
+        }
+
+        0
+    }
+
+    builder.add(map, callback).expect("Failed to add ringbuf");
+    let mgr = builder.build().expect("Failed to build");
+
+    // Define a USDT probe point and exercise it as we are attaching to self.
+    probe!(test_provider, test_function, 1);
+
+    mgr.consume().expect("Failed to consume ringbuf");
+
+    unsafe { assert_eq!(V, cookie_val.into()) };
+}
+
 #[test]
 fn test_object_map_probes() {
     bump_rlimit_mlock();