From e32ccb7e35ead7ae2d00dff593d4940b1c8b7f6f Mon Sep 17 00:00:00 2001
From: Benjamin Fogle <benfogle@gmail.com>
Date: Thu, 4 Nov 2021 08:48:56 -0400
Subject: [PATCH] Using a Display argument for lazy evaluation

---
 libbpf-rs/src/print.rs        | 40 +++++++++++++++++++----------------
 libbpf-rs/tests/test_print.rs |  9 ++++----
 2 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/libbpf-rs/src/print.rs b/libbpf-rs/src/print.rs
index 8e44a556..839bb90c 100644
--- a/libbpf-rs/src/print.rs
+++ b/libbpf-rs/src/print.rs
@@ -2,6 +2,7 @@ use crate::*;
 use std::io::{self, Write};
 use std::os::raw::c_char;
 use std::sync::atomic::{AtomicPtr, Ordering};
+use std::fmt;
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
 #[repr(u32)]
@@ -23,16 +24,31 @@ impl From<libbpf_sys::libbpf_print_level> for PrintLevel {
     }
 }
 
-pub type PrintCallback = fn(PrintLevel, &str);
+pub type PrintCallback = fn(PrintLevel, &dyn fmt::Display);
+
+struct Vsprintf {
+    fmtstr: *const c_char,
+    va_list: *mut libbpf_sys::__va_list_tag,
+}
+
+impl fmt::Display for Vsprintf {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let msg = match unsafe { vsprintf::vsprintf(self.fmtstr, self.va_list) } {
+            Ok(s) => s,
+            Err(e) => format!("Failed to parse libbpf output: {}", e),
+        };
+        f.write_str(&msg)
+    }
+}
 
 /// 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) {
+fn default_callback(lvl: PrintLevel, msg: &dyn fmt::Display) {
     if lvl == PrintLevel::Debug {
         return;
     }
 
-    let _ = io::stderr().write(msg.as_bytes());
+    let _ = io::stderr().write(msg.to_string().as_bytes());
 }
 
 fn from_ptr(ptr: *mut ()) -> PrintCallback {
@@ -42,27 +58,15 @@ fn from_ptr(ptr: *mut ()) -> PrintCallback {
 // There is no AtomicFnPtr. This is a workaround.
 static PRINT_CB: AtomicPtr<()> = AtomicPtr::new(default_callback as *mut ());
 
-/// libbpf's default cb uses vfprintf's return code...which is ignored everywhere. Mimic that for
-/// completeness
 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 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(e) => {
-            cb(
-                PrintLevel::Warn,
-                &format!("Failed to parse libbpf output: {}", e),
-            );
-            -1
-        }
-    }
+    let msg = Vsprintf { fmtstr, va_list };
+    cb(level.into(), &msg);
+    1 // return value is ignored by libbpf
 }
 
 /// Set a callback to receive log messages from libbpf, instead of printing them to stderr. Returns
diff --git a/libbpf-rs/tests/test_print.rs b/libbpf-rs/tests/test_print.rs
index 93655f99..279f5c6e 100644
--- a/libbpf-rs/tests/test_print.rs
+++ b/libbpf-rs/tests/test_print.rs
@@ -8,6 +8,7 @@
 use libbpf_rs::{set_print, ObjectBuilder, PrintLevel};
 use serial_test::serial;
 use std::sync::atomic::{AtomicBool, Ordering};
+use std::fmt::Display;
 
 #[test]
 #[serial]
@@ -15,12 +16,12 @@ fn test_set_print() {
     static CORRECT_LEVEL: AtomicBool = AtomicBool::new(false);
     static CORRECT_MESSAGE: AtomicBool = AtomicBool::new(false);
 
-    fn callback(level: PrintLevel, msg: &str) {
+    fn callback(level: PrintLevel, msg: &dyn Display) {
         if level == PrintLevel::Warn {
             CORRECT_LEVEL.store(true, Ordering::Relaxed);
         }
 
-        if msg.starts_with("libbpf: ") {
+        if msg.to_string().starts_with("libbpf: ") {
             CORRECT_MESSAGE.store(true, Ordering::Relaxed);
         }
     }
@@ -39,10 +40,10 @@ fn test_set_print() {
 #[test]
 #[serial]
 fn test_set_restore_print() {
-    fn callback1(_: PrintLevel, _: &str) {
+    fn callback1(_: PrintLevel, _: &dyn Display) {
         println!("one");
     }
-    fn callback2(_: PrintLevel, _: &str) {
+    fn callback2(_: PrintLevel, _: &dyn Display) {
         println!("two");
     }