Make set_print return the previous callback

Author: benfogle

libbpf-rs/Cargo.toml

@@ -31,3 +31,4 @@ vsprintf = "2.0"
 libc = "0.2"
 plain = "0.2.3"
 scopeguard = "1.1"
+serial_test = "0.5"

libbpf-rs/src/object.rs

@@ -29,30 +29,14 @@ impl ObjectBuilder {
     }
 
     /// Option to print debug output to stderr.
+    ///
+    /// Note: This function interacts poorly with [`set_print`].
     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(|_, s| print!("{}", s));
         } else {
-            unsafe { libbpf_sys::libbpf_set_print(None) };
+            set_print(|_, _| ());
         }
-
         self
     }
 

libbpf-rs/src/print.rs

@@ -1,4 +1,5 @@
 use crate::*;
+use std::io::{self, Write};
 use std::os::raw::c_char;
 use std::sync::atomic::{AtomicPtr, Ordering};
 
@@ -24,8 +25,29 @@ impl From<libbpf_sys::libbpf_print_level> for PrintLevel {
 
 pub type PrintCallback = fn(PrintLevel, &str);
 
+/// Mimic the default print functionality of libbpf. This way if the user calls `set_print` when no
+/// previous callback had been set and then tries to restore it, it will appear to work correctly.
+fn default_callback(lvl: PrintLevel, msg: &str) {
+    if lvl == PrintLevel::Debug {
+        return;
+    }
+
+    let _ = io::stderr().write(msg.as_bytes());
+}
+
+// not allowed to use function pointers in const functions, so it needs to be a macro
+macro_rules! to_ptr {
+    ($x:expr) => {
+        unsafe { std::mem::transmute::<PrintCallback, *mut ()>($x) }
+    };
+}
+
+fn from_ptr(ptr: *mut ()) -> PrintCallback {
+    unsafe { std::mem::transmute::<*mut (), PrintCallback>(ptr) }
+}
+
 // There is no AtomicFnPtr. This is a workaround.
-static PRINT_CB: AtomicPtr<()> = AtomicPtr::new(std::ptr::null_mut());
+static PRINT_CB: AtomicPtr<()> = AtomicPtr::new(to_ptr!(default_callback));
 
 /// libbpf's default cb uses vfprintf's return code...which is ignored everywhere. Mimic that for
 /// completeness
@@ -34,27 +56,33 @@ extern "C" fn outer_print_cb(
     fmtstr: *const c_char,
     va_list: *mut libbpf_sys::__va_list_tag,
 ) -> i32 {
-    // can't be null: we cant't get here until print_cb has been initialized with a function
-    // pointer.
-    let cb = PRINT_CB.load(Ordering::Relaxed);
-    let cb = unsafe { std::mem::transmute::<*const (), PrintCallback>(cb) };
+    let cb = from_ptr(PRINT_CB.load(Ordering::Relaxed));
     match unsafe { vsprintf::vsprintf(fmtstr, va_list) } {
         Ok(s) => {
             cb(level.into(), &s);
             s.len() as i32
         }
-        Err(_) => {
-            cb(PrintLevel::Warn, "Could not format libbpf log string");
+        Err(e) => {
+            cb(
+                PrintLevel::Warn,
+                &format!("Failed to parse libbpf output: {}", e),
+            );
             -1
         }
     }
 }
 
-/// Set a callback to receive log messages from libbpf, instead of printing them to stderr
+/// Set a callback to receive log messages from libbpf, instead of printing them to stderr. Returns
+/// the previous callback.
+///
+/// # Arguments
+///
+/// * `func` - The callback
 ///
 /// This overrides (and is overridden by) [`ObjectBuilder::debug`]
-pub fn set_print(func: PrintCallback) {
-    let cb = unsafe { std::mem::transmute::<PrintCallback, *mut ()>(func) };
-    PRINT_CB.store(cb, Ordering::Relaxed);
+pub fn set_print(func: PrintCallback) -> PrintCallback {
+    let cb = to_ptr!(func);
+    let prev = PRINT_CB.swap(cb, Ordering::Relaxed);
     unsafe { libbpf_sys::libbpf_set_print(Some(outer_print_cb)) };
+    from_ptr(prev)
 }

libbpf-rs/tests/test_print.rs

@@ -2,11 +2,15 @@
 //! 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::{set_print, ObjectBuilder, 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);
@@ -31,3 +35,18 @@ fn test_set_print() {
     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, _: &str) {
+        println!("one");
+    }
+    fn callback2(_: PrintLevel, _: &str) {
+        println!("two");
+    }
+
+    set_print(callback1);
+    assert_eq!(callback1 as usize, set_print(callback2) as usize);
+    assert_eq!(callback2 as usize, set_print(callback1) as usize);
+}