libbpf-rs: Add set_print/get_print, to send libbpf output to a callback

This commit exposes libbpf_set_print to direct libbpf output to a
user-defined callback.

Signed-off-by: Benjamin Fogle <benfogle@gmail.com>

libbpf-rs/Cargo.toml

@@ -26,8 +26,11 @@ nix = "0.22"
 num_enum = "0.5"
 strum_macros = "0.21"
 vsprintf = "2.0"
+lazy_static = "1.4"
 
 [dev-dependencies]
 libc = "0.2"
 plain = "0.2.3"
 scopeguard = "1.1"
+serial_test = "0.5"
+log = "0.4"

libbpf-rs/src/lib.rs

@@ -72,6 +72,7 @@ mod link;
 mod map;
 mod object;
 mod perf_buffer;
+mod print;
 mod program;
 pub mod query;
 mod ringbuf;
@@ -88,6 +89,7 @@ pub use crate::link::Link;
 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::ringbuf::{RingBuffer, RingBufferBuilder};
 pub use crate::util::num_possible_cpus;

libbpf-rs/src/object.rs

@@ -29,30 +29,15 @@ impl ObjectBuilder {
     }
 
     /// Option to print debug output to stderr.
+    ///
+    /// Note: This function uses [`set_print`] internally and will overwrite any callbacks
+    /// currently in use.
     pub fn debug(&mut self, dbg: bool) -> &mut Self {
-        extern "C" fn cb(
-            _level: libbpf_sys::libbpf_print_level,
-            fmtstr: *const c_char,
-            va_list: *mut libbpf_sys::__va_list_tag,
-        ) -> i32 {
-            match unsafe { vsprintf::vsprintf(fmtstr, va_list) } {
-                Ok(s) => {
-                    print!("{}", s);
-                    0
-                }
-                Err(e) => {
-                    eprintln!("Failed to parse libbpf output: {}", e);
-                    1
-                }
-            }
-        }
-
         if dbg {
-            unsafe { libbpf_sys::libbpf_set_print(Some(cb)) };
+            set_print(Some((PrintLevel::Debug, |_, s| print!("{}", s))));
         } else {
-            unsafe { libbpf_sys::libbpf_set_print(None) };
+            set_print(None);
         }
-
         self
     }
 

libbpf-rs/src/print.rs

@@ -0,0 +1,137 @@
+use crate::*;
+use lazy_static::lazy_static;
+use std::io::{self, Write};
+use std::os::raw::c_char;
+use std::sync::Mutex;
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
+#[repr(u32)]
+pub enum PrintLevel {
+    Warn = libbpf_sys::LIBBPF_WARN,
+    Info = libbpf_sys::LIBBPF_INFO,
+    Debug = libbpf_sys::LIBBPF_DEBUG,
+}
+
+impl From<libbpf_sys::libbpf_print_level> for PrintLevel {
+    fn from(level: libbpf_sys::libbpf_print_level) -> Self {
+        match level {
+            libbpf_sys::LIBBPF_WARN => Self::Warn,
+            libbpf_sys::LIBBPF_INFO => Self::Info,
+            libbpf_sys::LIBBPF_DEBUG => Self::Debug,
+            // shouldn't happen, but anything unknown becomes the highest level
+            _ => Self::Warn,
+        }
+    }
+}
+
+pub type PrintCallback = fn(PrintLevel, String);
+
+/// Mimic the default print functionality of libbpf. This way if the user calls `get_print` when no
+/// previous callback had been set, with the intention of restoring it, everything will behave as
+/// expected.
+fn default_callback(_lvl: PrintLevel, msg: String) {
+    let _ = io::stderr().write(msg.as_bytes());
+}
+
+// While we can't say that set_print is thread-safe, because we shouldn't assume that of
+// libbpf_set_print, we should still make sure that things are sane on the rust side of things.
+// Therefore we are using a lock to keep the log level and the callback in sync.
+//
+// We don't do anything that can panic with the lock held, so we'll unconditionally unwrap() when
+// locking the mutex.
+//
+// Note that default print behavior ignores debug messages.
+lazy_static! {
+    static ref PRINT_CB: Mutex<Option<(PrintLevel, PrintCallback)>> =
+        Mutex::new(Some((PrintLevel::Info, default_callback)));
+}
+
+extern "C" fn outer_print_cb(
+    level: libbpf_sys::libbpf_print_level,
+    fmtstr: *const c_char,
+    va_list: *mut libbpf_sys::__va_list_tag,
+) -> i32 {
+    let level = level.into();
+    if let Some((min_level, func)) = { *PRINT_CB.lock().unwrap() } {
+        if level <= min_level {
+            let msg = match unsafe { vsprintf::vsprintf(fmtstr, va_list) } {
+                Ok(s) => s,
+                Err(e) => format!("Failed to parse libbpf output: {}", e),
+            };
+            func(level, msg);
+        }
+    }
+    0 // return value is ignored by libbpf
+}
+
+/// Set a callback to receive log messages from libbpf, instead of printing them to stderr.
+///
+/// # Arguments
+///
+/// * `callback` - Either a tuple `(min_level, function)` where `min_level` is the lowest priority
+///   log message to handle, or `None` to disable all printing.
+///
+/// This overrides (and is overridden by) [`ObjectBuilder::debug`]
+///
+/// # Examples
+///
+/// To pass all messages to the `log` crate:
+///
+/// ```
+/// use log;
+/// use libbpf_rs::{PrintLevel, set_print};
+///
+/// fn print_to_log(level: PrintLevel, msg: String) {
+///     match level {
+///         PrintLevel::Debug => log::debug!("{}", msg),
+///         PrintLevel::Info => log::info!("{}", msg),
+///         PrintLevel::Warn => log::warn!("{}", msg),
+///     }
+/// }
+///
+/// set_print(Some((PrintLevel::Debug, print_to_log)));
+/// ```
+///
+/// To disable printing completely:
+///
+/// ```
+/// use libbpf_rs::set_print;
+/// set_print(None);
+/// ```
+///
+/// To temporarliy suppress output:
+///
+/// ```
+/// use libbpf_rs::set_print;
+///
+/// let prev = set_print(None);
+/// // do things quietly
+/// set_print(prev);
+/// ```
+pub fn set_print(
+    mut callback: Option<(PrintLevel, PrintCallback)>,
+) -> Option<(PrintLevel, PrintCallback)> {
+    let real_cb: libbpf_sys::libbpf_print_fn_t;
+    real_cb = callback.as_ref().and(Some(outer_print_cb));
+    std::mem::swap(&mut callback, &mut *PRINT_CB.lock().unwrap());
+    unsafe { libbpf_sys::libbpf_set_print(real_cb) };
+    callback
+}
+
+/// Return the current print callback and level.
+///
+/// # Examples
+///
+/// To temporarliy suppress output:
+///
+/// ```
+/// use libbpf_rs::{get_print, set_print};
+///
+/// let prev = get_print();
+/// set_print(None);
+/// // do things quietly
+/// set_print(prev);
+/// ```
+pub fn get_print() -> Option<(PrintLevel, PrintCallback)> {
+    *PRINT_CB.lock().unwrap()
+}

libbpf-rs/tests/test_print.rs

@@ -0,0 +1,74 @@
+//! This test is in its own file because the underlying libbpf_set_print function used by
+//! set_print() and ObjectBuilder::debug() sets global state. The default is to run multiple tests
+//! in different threads, so this test will always race with the others unless its isolated to a
+//! different process.
+//!
+//! For the same reason, all tests here must run serially.
+
+use libbpf_rs::{get_print, set_print, ObjectBuilder, PrintCallback, PrintLevel};
+use serial_test::serial;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+#[test]
+#[serial]
+fn test_set_print() {
+    static CORRECT_LEVEL: AtomicBool = AtomicBool::new(false);
+    static CORRECT_MESSAGE: AtomicBool = AtomicBool::new(false);
+
+    fn callback(level: PrintLevel, msg: String) {
+        if level == PrintLevel::Warn {
+            CORRECT_LEVEL.store(true, Ordering::Relaxed);
+        }
+
+        if msg.starts_with("libbpf: ") {
+            CORRECT_MESSAGE.store(true, Ordering::Relaxed);
+        }
+    }
+
+    set_print(Some((PrintLevel::Debug, callback)));
+    // expect_err requires that OpenObject implement Debug, which it does not.
+    let obj = ObjectBuilder::default().open_file("/dev/null");
+    assert!(obj.is_err(), "Successfully loaded /dev/null?");
+
+    let correct_level = CORRECT_LEVEL.load(Ordering::Relaxed);
+    let correct_message = CORRECT_MESSAGE.load(Ordering::Relaxed);
+    assert!(correct_level, "Did not capture a warning");
+    assert!(correct_message, "Did not capture the correct message");
+}
+
+#[test]
+#[serial]
+fn test_set_restore_print() {
+    fn callback1(_: PrintLevel, _: String) {
+        println!("one");
+    }
+    fn callback2(_: PrintLevel, _: String) {
+        println!("two");
+    }
+
+    set_print(Some((PrintLevel::Warn, callback1)));
+    let prev = get_print();
+    assert_eq!(prev, Some((PrintLevel::Warn, callback1 as PrintCallback)));
+
+    set_print(Some((PrintLevel::Debug, callback2)));
+    let prev = get_print();
+    assert_eq!(prev, Some((PrintLevel::Debug, callback2 as PrintCallback)));
+}
+
+#[test]
+#[serial]
+fn test_set_and_save_print() {
+    fn callback1(_: PrintLevel, _: String) {
+        println!("one");
+    }
+    fn callback2(_: PrintLevel, _: String) {
+        println!("two");
+    }
+
+    set_print(Some((PrintLevel::Warn, callback1)));
+    let prev = set_print(Some((PrintLevel::Debug, callback2)));
+    assert_eq!(prev, Some((PrintLevel::Warn, callback1 as PrintCallback)));
+
+    let prev = set_print(None);
+    assert_eq!(prev, Some((PrintLevel::Debug, callback2 as PrintCallback)));
+}