From 0970300d1f5659622fa55a18dd7681c608d75b0f Mon Sep 17 00:00:00 2001
From: Hanaasagi <ambiguous404@gmail.com>
Date: Fri, 21 Apr 2023 15:55:27 +0800
Subject: [PATCH] feat(aya-log): check format and value type in proc macro

---
 aya-log-common/src/lib.rs         | 76 +++++++++++++++++++++++++++++++
 aya-log-ebpf-macros/src/expand.rs | 26 ++++++++++-
 aya-log/src/lib.rs                |  5 +-
 bpf/aya-log-ebpf/src/lib.rs       |  7 ++-
 4 files changed, 111 insertions(+), 3 deletions(-)

diff --git a/aya-log-common/src/lib.rs b/aya-log-common/src/lib.rs
index f833f3c1..ce04f210 100644
--- a/aya-log-common/src/lib.rs
+++ b/aya-log-common/src/lib.rs
@@ -35,6 +35,82 @@ pub enum Level {
     Trace,
 }
 
+macro_rules! impl_formatter_for_types {
+    ($trait:path : { $($type:ty),*}) => {
+        $(
+            impl $trait for $type {}
+        )*
+    };
+}
+
+// Any value impl `ToString`
+pub trait DefaultFormatter {}
+impl_formatter_for_types!(
+    DefaultFormatter: {
+        bool,
+        i8, i16, i32, i64, isize,
+        u8, u16, u32, u64, usize,
+        f32, f64,
+        char,
+        str,
+        &str
+    }
+);
+
+pub trait LowerHexFormatter {}
+impl_formatter_for_types!(
+    LowerHexFormatter: {
+        i8, i16, i32, i64, isize,
+        u8, u16, u32, u64, usize
+    }
+);
+
+pub trait LowerHexDebugFormatter {}
+impl LowerHexDebugFormatter for &[u8] {}
+
+pub trait UpperHexFormatter {}
+impl_formatter_for_types!(
+    UpperHexFormatter: {
+        i8, i16, i32, i64, isize,
+        u8, u16, u32, u64, usize
+    }
+);
+
+pub trait UpperHexDebugFormatter {}
+impl UpperHexDebugFormatter for &[u8] {}
+
+pub trait Ipv4Formatter {}
+impl Ipv4Formatter for u32 {}
+
+pub trait Ipv6Formatter {}
+impl Ipv6Formatter for [u8; 16] {}
+impl Ipv6Formatter for [u16; 8] {}
+
+pub trait LowerMacFormatter {}
+impl LowerMacFormatter for [u8; 6] {}
+
+pub trait UpperMacFormatter {}
+impl UpperMacFormatter for [u8; 6] {}
+
+#[inline(always)]
+pub fn check_impl_default<T: DefaultFormatter>(_v: T) {}
+#[inline(always)]
+pub fn check_impl_lower_hex<T: LowerHexFormatter>(_v: T) {}
+#[inline(always)]
+pub fn check_impl_lower_hex_debug<T: LowerHexDebugFormatter>(_v: T) {}
+#[inline(always)]
+pub fn check_impl_upper_hex<T: UpperHexFormatter>(_v: T) {}
+#[inline(always)]
+pub fn check_impl_upper_hex_debug<T: UpperHexDebugFormatter>(_v: T) {}
+#[inline(always)]
+pub fn check_impl_ipv4<T: Ipv4Formatter>(_v: T) {}
+#[inline(always)]
+pub fn check_impl_ipv6<T: Ipv6Formatter>(_v: T) {}
+#[inline(always)]
+pub fn check_impl_lower_mac<T: LowerMacFormatter>(_v: T) {}
+#[inline(always)]
+pub fn check_impl_upper_mac<T: UpperMacFormatter>(_v: T) {}
+
 #[repr(u8)]
 #[derive(Copy, Clone, Debug)]
 pub enum RecordField {
diff --git a/aya-log-ebpf-macros/src/expand.rs b/aya-log-ebpf-macros/src/expand.rs
index eefcbb72..01de1bf9 100644
--- a/aya-log-ebpf-macros/src/expand.rs
+++ b/aya-log-ebpf-macros/src/expand.rs
@@ -86,6 +86,18 @@ fn hint_to_expr(hint: DisplayHint) -> Result<Expr> {
     }
 }
 
+fn hint_to_format_check(hint: DisplayHint) -> Result<Expr> {
+    match hint {
+        DisplayHint::Default => parse_str("::aya_log_ebpf::macro_support::check_impl_default"),
+        DisplayHint::LowerHex => parse_str("::aya_log_ebpf::macro_support::check_impl_lower_hex"),
+        DisplayHint::UpperHex => parse_str("::aya_log_ebpf::macro_support::check_impl_upper_hex"),
+        DisplayHint::Ipv4 => parse_str("::aya_log_ebpf::macro_support::check_impl_ipv4"),
+        DisplayHint::Ipv6 => parse_str("::aya_log_ebpf::macro_support::check_impl_ipv6"),
+        DisplayHint::LowerMac => parse_str("::aya_log_ebpf::macro_support::check_impl_lower_mac"),
+        DisplayHint::UpperMac => parse_str("::aya_log_ebpf::macro_support::check_impl_upper_mac"),
+    }
+}
+
 pub(crate) fn log(args: LogArgs, level: Option<TokenStream>) -> Result<TokenStream> {
     let ctx = args.ctx;
     let target = match args.target {
@@ -115,6 +127,8 @@ pub(crate) fn log(args: LogArgs, level: Option<TokenStream>) -> Result<TokenStre
     let mut arg_i = 0;
 
     let mut values = Vec::new();
+    let mut f_keys = Vec::new();
+    let mut f_values = Vec::new();
     for fragment in fragments {
         match fragment {
             Fragment::Literal(s) => {
@@ -126,7 +140,10 @@ pub(crate) fn log(args: LogArgs, level: Option<TokenStream>) -> Result<TokenStre
                     None => return Err(Error::new(format_string.span(), "no arguments provided")),
                 };
                 values.push(hint_to_expr(p.hint)?);
-                values.push(arg);
+                values.push(arg.clone());
+
+                f_keys.push(hint_to_format_check(p.hint)?);
+                f_values.push(arg.clone());
                 arg_i += 1;
             }
         }
@@ -135,8 +152,15 @@ pub(crate) fn log(args: LogArgs, level: Option<TokenStream>) -> Result<TokenStre
     let num_args = values.len();
     let values_iter = values.iter();
 
+    let f_keys = f_keys.iter();
+    let f_values = f_values.iter();
+
     Ok(quote! {
         {
+            #(
+                #f_keys(#f_values);
+            )*
+
             if let Some(buf_ptr) = unsafe { ::aya_log_ebpf::AYA_LOG_BUF.get_ptr_mut(0) } {
                 let buf = unsafe { &mut *buf_ptr };
                 if let Ok(header_len) = ::aya_log_ebpf::write_record_header(
diff --git a/aya-log/src/lib.rs b/aya-log/src/lib.rs
index 92d686bb..ff88df5c 100644
--- a/aya-log/src/lib.rs
+++ b/aya-log/src/lib.rs
@@ -126,7 +126,10 @@ impl BpfLogger {
                     #[allow(clippy::needless_range_loop)]
                     for i in 0..events.read {
                         let buf = &mut buffers[i];
-                        log_buf(buf, &*log).unwrap();
+                        match log_buf(buf, &*log) {
+                            Ok(()) => {}
+                            Err(e) => error!("internal error => {:?}", e),
+                        }
                     }
                 }
             });
diff --git a/bpf/aya-log-ebpf/src/lib.rs b/bpf/aya-log-ebpf/src/lib.rs
index 526aed48..777b8771 100644
--- a/bpf/aya-log-ebpf/src/lib.rs
+++ b/bpf/aya-log-ebpf/src/lib.rs
@@ -22,6 +22,11 @@ pub static mut AYA_LOGS: PerfEventByteArray = PerfEventByteArray::new(0);
 
 #[doc(hidden)]
 pub mod macro_support {
-    pub use aya_log_common::{DisplayHint, Level, LOG_BUF_CAPACITY};
+    pub use aya_log_common::{
+        check_impl_default, check_impl_ipv4, check_impl_ipv6, check_impl_lower_hex,
+        check_impl_lower_mac, check_impl_upper_hex, check_impl_upper_mac, DefaultFormatter,
+        DisplayHint, Ipv4Formatter, Ipv6Formatter, Level, LowerHexFormatter, LowerMacFormatter,
+        UpperHexFormatter, UpperMacFormatter, LOG_BUF_CAPACITY,
+    };
     pub use aya_log_ebpf_macros::log;
 }