From 55ba0538f277db3a7e0459a54ca8af027f851723 Mon Sep 17 00:00:00 2001
From: Tw <wei.tan@intel.com>
Date: Tue, 6 Jul 2021 11:24:00 +0800
Subject: [PATCH] bpf: add support for tracepoint program (#29)

This patch add initial support for tracepoint program type.
Hope you enjoy.

Signed-off-by: Tw <wei.tan@intel.com>
---
 bpf/aya-bpf-macros/src/expand.rs       | 29 ++++++++++++++++++++++++++
 bpf/aya-bpf-macros/src/lib.rs          | 13 +++++++++++-
 bpf/aya-bpf/src/maps/hash_map.rs       |  2 +-
 bpf/aya-bpf/src/programs/mod.rs        |  2 ++
 bpf/aya-bpf/src/programs/tracepoint.rs | 22 +++++++++++++++++++
 5 files changed, 66 insertions(+), 2 deletions(-)
 create mode 100644 bpf/aya-bpf/src/programs/tracepoint.rs

diff --git a/bpf/aya-bpf-macros/src/expand.rs b/bpf/aya-bpf-macros/src/expand.rs
index b079026e..44cb8c5b 100644
--- a/bpf/aya-bpf-macros/src/expand.rs
+++ b/bpf/aya-bpf-macros/src/expand.rs
@@ -293,3 +293,32 @@ impl std::fmt::Display for ProbeKind {
         }
     }
 }
+
+pub struct TracePoint {
+    item: ItemFn,
+    name: String,
+}
+
+impl TracePoint {
+    pub fn from_syn(mut args: Args, item: ItemFn) -> Result<TracePoint> {
+        let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string());
+
+        Ok(TracePoint { item, name })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        let section_name = format!("tp/{}", self.name);
+        let fn_name = &self.item.sig.ident;
+        let item = &self.item;
+        Ok(quote! {
+            #[no_mangle]
+            #[link_section = #section_name]
+            fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 {
+               let _ = #fn_name(::aya_bpf::programs::TracePointContext::new(ctx));
+               return 0;
+
+               #item
+            }
+        })
+    }
+}
diff --git a/bpf/aya-bpf-macros/src/lib.rs b/bpf/aya-bpf-macros/src/lib.rs
index 8c990dd8..9ed55673 100644
--- a/bpf/aya-bpf-macros/src/lib.rs
+++ b/bpf/aya-bpf-macros/src/lib.rs
@@ -1,6 +1,6 @@
 mod expand;
 
-use expand::{Args, Map, Probe, ProbeKind, SchedClassifier, SkMsg, SockOps, Xdp};
+use expand::{Args, Map, Probe, ProbeKind, SchedClassifier, SkMsg, SockOps, TracePoint, Xdp};
 use proc_macro::TokenStream;
 use syn::{parse_macro_input, ItemFn, ItemStatic};
 
@@ -101,3 +101,14 @@ fn probe(kind: ProbeKind, attrs: TokenStream, item: TokenStream) -> TokenStream
         .unwrap_or_else(|err| err.to_compile_error())
         .into()
 }
+
+#[proc_macro_attribute]
+pub fn tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    let args = parse_macro_input!(attrs as Args);
+    let item = parse_macro_input!(item as ItemFn);
+
+    TracePoint::from_syn(args, item)
+        .and_then(|u| u.expand())
+        .unwrap_or_else(|err| err.to_compile_error())
+        .into()
+}
diff --git a/bpf/aya-bpf/src/maps/hash_map.rs b/bpf/aya-bpf/src/maps/hash_map.rs
index 9afc41a4..3f3a803d 100644
--- a/bpf/aya-bpf/src/maps/hash_map.rs
+++ b/bpf/aya-bpf/src/maps/hash_map.rs
@@ -4,7 +4,7 @@ use aya_bpf_cty::{c_long, c_void};
 
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_HASH},
-    helpers::{bpf_map_lookup_elem, bpf_map_update_elem, bpf_map_delete_elem},
+    helpers::{bpf_map_delete_elem, bpf_map_lookup_elem, bpf_map_update_elem},
 };
 
 #[repr(transparent)]
diff --git a/bpf/aya-bpf/src/programs/mod.rs b/bpf/aya-bpf/src/programs/mod.rs
index feb9ac1d..979fe11a 100644
--- a/bpf/aya-bpf/src/programs/mod.rs
+++ b/bpf/aya-bpf/src/programs/mod.rs
@@ -2,10 +2,12 @@ pub mod probe;
 pub mod sk_msg;
 pub mod sk_skb;
 pub mod sock_ops;
+pub mod tracepoint;
 pub mod xdp;
 
 pub use probe::ProbeContext;
 pub use sk_msg::SkMsgContext;
 pub use sk_skb::SkSkbContext;
 pub use sock_ops::SockOpsContext;
+pub use tracepoint::TracePointContext;
 pub use xdp::XdpContext;
diff --git a/bpf/aya-bpf/src/programs/tracepoint.rs b/bpf/aya-bpf/src/programs/tracepoint.rs
new file mode 100644
index 00000000..29db68b4
--- /dev/null
+++ b/bpf/aya-bpf/src/programs/tracepoint.rs
@@ -0,0 +1,22 @@
+use crate::{helpers::bpf_probe_read, BpfContext};
+use core::ffi::c_void;
+
+pub struct TracePointContext {
+    ctx: *mut c_void,
+}
+
+impl TracePointContext {
+    pub fn new(ctx: *mut c_void) -> TracePointContext {
+        TracePointContext { ctx }
+    }
+
+    pub unsafe fn read_at<T>(&self, offset: usize) -> Result<T, i64> {
+        bpf_probe_read(self.ctx.add(offset) as *const T)
+    }
+}
+
+impl BpfContext for TracePointContext {
+    fn as_ptr(&self) -> *mut c_void {
+        self.ctx
+    }
+}