map support

Cargo.toml

@@ -9,3 +9,6 @@ members = [
     "examples/syscount",
     "examples/syscount/probe",
 ]
+
+[patch.crates-io]
+libbpf-rs = { git = "https://github.com/libbpf/libbpf-rs" }

bpf-helpers/Cargo.toml

@@ -7,4 +7,6 @@ edition = "2018"
 [dependencies]
 bpf-helpers-sys = { version = "0.1.0", path = "../bpf-helpers-sys" }
 bpf-macros = { version = "0.1.0", path = "../bpf-macros" }
+byteorder = { version = "1.4.2", default-features = false }
 cty = "0.2.1"
+zerocopy = { version = "0.3.0", default-features = false }

bpf-helpers/src/lib.rs

@@ -10,6 +10,13 @@ pub use crate::time::*;
 pub use bpf_macros::*;
 pub use cty;
 
+pub type I16 = zerocopy::byteorder::I16<byteorder::NativeEndian>;
+pub type I32 = zerocopy::byteorder::I32<byteorder::NativeEndian>;
+pub type I64 = zerocopy::byteorder::I64<byteorder::NativeEndian>;
+pub type U16 = zerocopy::byteorder::U16<byteorder::NativeEndian>;
+pub type U32 = zerocopy::byteorder::U32<byteorder::NativeEndian>;
+pub type U64 = zerocopy::byteorder::U64<byteorder::NativeEndian>;
+
 #[inline]
 pub fn bpf_trace_printk(msg: &[u8]) -> usize {
     unsafe {

bpf-helpers/src/map.rs

@@ -46,41 +46,51 @@ impl<K, V> HashMap<K, V> {
 
     /// Returns a reference to the value corresponding to the key.
     #[inline]
-    pub fn get(&mut self, key: &K) -> Option<&V> {
-        unsafe {
+    pub fn get<R, F: FnOnce(Option<&mut V>) -> R>(&self, key: &K, f: F) -> R {
+        let rvalue = unsafe {
             let value = bpf_helpers_sys::bpf_map_lookup_elem(
-                &mut self.def as *mut _ as *mut c_void,
+                &self.def as *const _ as *mut c_void,
                 key as *const _ as *const c_void,
             );
             if value.is_null() {
                 None
             } else {
-                Some(&*(value as *const V))
+                Some(&mut *(value as *mut V))
             }
-        }
+        };
+        f(rvalue)
     }
 
+    /// Returns a reference to the value corresponding to the key.
     #[inline]
-    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
-        unsafe {
-            let value = bpf_helpers_sys::bpf_map_lookup_elem(
-                &mut self.def as *mut _ as *mut c_void,
-                key as *const _ as *const c_void,
-            );
-            if value.is_null() {
-                None
+    pub fn get_map<R, F: FnOnce(&mut V) -> R>(&self, key: &K, f: F) -> Option<R> {
+        self.get(key, |value| value.map(f))
+    }
+
+    /// Returns a reference to the value corresponding to the key.
+    #[inline]
+    pub fn get_or_default<R, F: FnOnce(&mut V) -> R>(&self, key: &K, f: F) -> R
+    where
+        V: Default,
+    {
+        self.get(key, |value| {
+            if let Some(value) = value {
+                f(value)
             } else {
-                Some(&mut *(value as *mut V))
+                let mut value = V::default();
+                let res = f(&mut value);
+                self.set(key, &value);
+                res
             }
-        }
+        })
     }
 
     /// Set the `value` in the map for `key`
     #[inline]
-    pub fn set(&mut self, key: &K, value: &V) {
+    pub fn set(&self, key: &K, value: &V) {
         unsafe {
             bpf_helpers_sys::bpf_map_update_elem(
-                &mut self.def as *mut _ as *mut c_void,
+                &self.def as *const _ as *mut c_void,
                 key as *const _ as *const c_void,
                 value as *const _ as *const c_void,
                 bpf_helpers_sys::BPF_ANY.into(),
@@ -90,10 +100,10 @@ impl<K, V> HashMap<K, V> {
 
     /// Delete the entry indexed by `key`
     #[inline]
-    pub fn delete(&mut self, key: &K) {
+    pub fn delete(&self, key: &K) {
         unsafe {
             bpf_helpers_sys::bpf_map_delete_elem(
-                &mut self.def as *mut _ as *mut c_void,
+                &self.def as *const _ as *mut c_void,
                 key as *const _ as *const c_void,
             );
         }

bpf-helpers/src/pid.rs

@@ -1,16 +1,21 @@
-#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
-pub struct PidTgid(u64);
+use byteorder::NativeEndian;
+use zerocopy::byteorder::U64;
+
+#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
+pub struct PidTgid(U64<NativeEndian>);
 
 impl PidTgid {
     pub fn current() -> Self {
-        Self(unsafe { bpf_helpers_sys::bpf_get_current_pid_tgid() })
+        Self(U64::new(unsafe {
+            bpf_helpers_sys::bpf_get_current_pid_tgid()
+        }))
     }
 
     pub fn pid(&self) -> u32 {
-        (self.0 >> 32) as _
+        (self.0.get() >> 32) as _
     }
 
     pub fn tgid(&self) -> u32 {
-        (self.0 & 0xf) as _
+        (self.0.get() & 0xf) as _
     }
 }

bpf-helpers/src/time.rs

@@ -1,18 +1,50 @@
-pub use core::time::Duration;
+use byteorder::NativeEndian;
+use core::ops::{Add, AddAssign};
+use zerocopy::byteorder::U64;
+use zerocopy::{AsBytes, FromBytes, Unaligned};
 
-#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
-pub struct Instant(u64);
+#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq, AsBytes, FromBytes, Unaligned)]
+#[repr(C)]
+pub struct Duration(U64<NativeEndian>);
+
+impl Duration {
+    pub fn from_nanos(nanos: u64) -> Self {
+        Self(U64::new(nanos))
+    }
+
+    pub fn as_nanos(&self) -> u64 {
+        self.0.get()
+    }
+}
+
+impl Add for Duration {
+    type Output = Self;
+
+    fn add(self, other: Self) -> Self {
+        Self(U64::new(self.0.get() + other.0.get()))
+    }
+}
+
+impl AddAssign for Duration {
+    fn add_assign(&mut self, other: Self) {
+        self.0.set(self.0.get() + other.0.get())
+    }
+}
+
+#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, AsBytes, FromBytes, Unaligned)]
+#[repr(C)]
+pub struct Instant(U64<NativeEndian>);
 
 impl Instant {
     pub fn now() -> Self {
-        Self(unsafe { bpf_helpers_sys::bpf_ktime_get_ns() })
+        Self(U64::new(unsafe { bpf_helpers_sys::bpf_ktime_get_ns() }))
     }
 
     pub fn duration_since(&self, earlier: Instant) -> Option<Duration> {
-        if earlier > *self {
+        if earlier.0.get() > self.0.get() {
             return None;
         }
-        Some(Duration::from_nanos(earlier.0 - self.0))
+        Some(Duration::from_nanos(self.0.get() - earlier.0.get()))
     }
 
     pub fn elapsed(&self) -> Duration {

bpf-macros/src/lib.rs

@@ -61,16 +61,11 @@ pub fn program(input: TokenStream) -> TokenStream {
 }
 
 #[proc_macro_attribute]
-pub fn map(attrs: TokenStream, item: TokenStream) -> TokenStream {
+pub fn map(_: TokenStream, item: TokenStream) -> TokenStream {
     let map = parse_macro_input!(item as syn::ItemStatic);
-    let map_name = match parse_macro_input!(attrs as syn::Lit) {
-        syn::Lit::Str(s) => s.value(),
-        _ => panic!("expected string literal"),
-    };
-    let section_name = format!(".maps/{}", map_name);
     let tokens = quote! {
         #[no_mangle]
-        #[link_section = #section_name]
+        #[link_section = "maps"]
         #map
     };
     tokens.into()
@@ -85,11 +80,11 @@ pub fn entry(attrs: TokenStream, item: TokenStream) -> TokenStream {
     };
     let mut event = quote!();
     let arg = match prog_type.as_str() {
-        "kprobe" => quote!(&bpf_helpers::kprobe::pt_regs),
-        "perf_event" => quote!(&bpf_helpers::perf_event::bpf_perf_event_data),
-        "tracing" => quote!(*const core::ffi::c_void),
-        "raw_tracepoint" => quote!(u64),
-        "raw_tracepoint_writable" => quote!(u64),
+        "kprobe" => quote!(bpf_helpers::kprobe::pt_regs),
+        "perf_event" => quote!(bpf_helpers::perf_event::bpf_perf_event_data),
+        "tracing" => quote!(core::ffi::c_void),
+        //"raw_tracepoint" => quote!(u64),
+        //"raw_tracepoint_writable" => quote!(u64),
         tracepoint => {
             let mut iter = tracepoint.split(':');
             let category = iter.next().expect("category");
@@ -103,11 +98,12 @@ pub fn entry(attrs: TokenStream, item: TokenStream) -> TokenStream {
             });
             prog_type = "tracepoint".to_string();
             event = quote! {
+                #[repr(C)]
                 struct #struct_ident {
                     #(#fields)*
                 }
             };
-            quote!(&#struct_ident)
+            quote!(#struct_ident)
         }
     };
     let ident = &prog.sig.ident;
@@ -120,8 +116,9 @@ pub fn entry(attrs: TokenStream, item: TokenStream) -> TokenStream {
         #[link_section = #section_name]
         fn #ident(arg: *const core::ffi::c_void) -> i32 {
             use bpf_helpers::#prog_type::*;
+            #[inline(always)]
             #prog
-            let arg: #arg = unsafe { core::mem::transmute(arg) };
+            let arg = unsafe { &*(arg as *const #arg) };
             #ident(arg);
             0
         }

bpf-probes/Cargo.toml

@@ -7,4 +7,7 @@ edition = "2018"
 [dependencies]
 anyhow = "1.0.38"
 libbpf-rs = "0.6.2"
+libc = "0.2.82"
+log = "0.4.13"
+perf-event-open-sys = "1.0.1"
 thiserror = "1.0.23"