bpf: add more bindings

Initial support for Array, HashMap and SockHash maps, and for SkSkb,
SkMsg, SockOps and XDP programs.

bpf/aya-bpf-macros/src/expand.rs

@@ -89,6 +89,98 @@ impl Probe {
     }
 }
 
+pub struct SockOps {
+    item: ItemFn,
+    name: Option<String>,
+}
+
+impl SockOps {
+    pub fn from_syn(args: Args, item: ItemFn) -> Result<SockOps> {
+        let name = name_arg(&args)?;
+
+        Ok(SockOps { item, name })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        let section_name = if let Some(name) = &self.name {
+            format!("sockops/{}", name)
+        } else {
+            "sockops".to_owned()
+        };
+        let fn_name = &self.item.sig.ident;
+        let item = &self.item;
+        Ok(quote! {
+            #[no_mangle]
+            #[link_section = #section_name]
+            fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sock_ops) -> u32 {
+                return #fn_name(::aya_bpf::programs::SockOpsContext::new(ctx));
+
+                #item
+            }
+        })
+    }
+}
+
+pub struct SkMsg {
+    item: ItemFn,
+    name: String,
+}
+
+impl SkMsg {
+    pub fn from_syn(args: Args, item: ItemFn) -> Result<SkMsg> {
+        let name = name_arg(&args)?.unwrap_or_else(|| item.sig.ident.to_string());
+
+        Ok(SkMsg { item, name })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        let section_name = format!("sk_msg/{}", 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 ::aya_bpf::bindings::sk_msg_md) -> u32 {
+                return #fn_name(::aya_bpf::programs::SkMsgContext::new(ctx));
+
+                #item
+            }
+        })
+    }
+}
+
+pub struct Xdp {
+    item: ItemFn,
+    name: Option<String>,
+}
+
+impl Xdp {
+    pub fn from_syn(args: Args, item: ItemFn) -> Result<Xdp> {
+        let name = name_arg(&args)?;
+
+        Ok(Xdp { item, name })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        let section_name = if let Some(name) = &self.name {
+            format!("xdp/{}", name)
+        } else {
+            "xdp".to_owned()
+        };
+        let fn_name = &self.item.sig.ident;
+        let item = &self.item;
+        Ok(quote! {
+            #[no_mangle]
+            #[link_section = #section_name]
+            fn #fn_name(ctx: *mut ::aya_bpf::bindings::xdp_md) -> u32 {
+                return #fn_name(::aya_bpf::programs::XdpContext::new(ctx));
+
+                #item
+            }
+        })
+    }
+}
+
 fn name_arg(args: &Args) -> Result<Option<String>> {
     for arg in &args.args {
         if arg.name == "name" {

bpf/aya-bpf-macros/src/lib.rs

@@ -1,6 +1,6 @@
 mod expand;
 
-use expand::{Args, Map, Probe, ProbeKind};
+use expand::{Args, Map, Probe, ProbeKind, SkMsg, SockOps, Xdp};
 use proc_macro::TokenStream;
 use syn::{parse_macro_input, ItemFn, ItemStatic};
 
@@ -35,6 +35,39 @@ pub fn uretprobe(attrs: TokenStream, item: TokenStream) -> TokenStream {
     probe(ProbeKind::URetProbe, attrs, item)
 }
 
+#[proc_macro_attribute]
+pub fn sock_ops(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    let args = parse_macro_input!(attrs as Args);
+    let item = parse_macro_input!(item as ItemFn);
+
+    SockOps::from_syn(args, item)
+        .and_then(|u| u.expand())
+        .unwrap_or_else(|err| err.to_compile_error())
+        .into()
+}
+
+#[proc_macro_attribute]
+pub fn sk_msg(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    let args = parse_macro_input!(attrs as Args);
+    let item = parse_macro_input!(item as ItemFn);
+
+    SkMsg::from_syn(args, item)
+        .and_then(|u| u.expand())
+        .unwrap_or_else(|err| err.to_compile_error())
+        .into()
+}
+
+#[proc_macro_attribute]
+pub fn xdp(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    let args = parse_macro_input!(attrs as Args);
+    let item = parse_macro_input!(item as ItemFn);
+
+    Xdp::from_syn(args, item)
+        .and_then(|u| u.expand())
+        .unwrap_or_else(|err| err.to_compile_error())
+        .into()
+}
+
 fn probe(kind: ProbeKind, attrs: TokenStream, item: TokenStream) -> TokenStream {
     let args = parse_macro_input!(attrs as Args);
     let item = parse_macro_input!(item as ItemFn);

bpf/aya-bpf/src/lib.rs

@@ -24,12 +24,10 @@ pub trait BpfContext {
         bpf_get_current_comm()
     }
 
-    #[inline]
     fn pid(&self) -> u32 {
         bpf_get_current_pid_tgid() as u32
     }
 
-    #[inline]
     fn tgid(&self) -> u32 {
         (bpf_get_current_pid_tgid() >> 32) as u32
     }

bpf/aya-bpf/src/maps/array.rs

@@ -0,0 +1,42 @@
+use core::{marker::PhantomData, mem};
+
+use aya_bpf_cty::c_void;
+
+use crate::{
+    bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_ARRAY},
+    helpers::bpf_map_lookup_elem,
+};
+
+#[repr(transparent)]
+pub struct Array<T> {
+    def: bpf_map_def,
+    _t: PhantomData<T>,
+}
+
+impl<T> Array<T> {
+    pub const fn with_max_entries(max_entries: u32, flags: u32) -> Array<T> {
+        Array {
+            def: bpf_map_def {
+                type_: BPF_MAP_TYPE_ARRAY,
+                key_size: mem::size_of::<u32>() as u32,
+                value_size: mem::size_of::<T>() as u32,
+                max_entries,
+                map_flags: flags,
+            },
+            _t: PhantomData,
+        }
+    }
+
+    pub unsafe fn get(&mut self, index: u32) -> Option<&T> {
+        let value = bpf_map_lookup_elem(
+            &mut self.def as *mut _ as *mut _,
+            &index as *const _ as *const c_void,
+        );
+        if value.is_null() {
+            None
+        } else {
+            // FIXME: alignment
+            Some(&*(value as *const T))
+        }
+    }
+}

bpf/aya-bpf/src/maps/hash_map.rs

@@ -0,0 +1,44 @@
+use core::{marker::PhantomData, mem};
+
+use aya_bpf_cty::c_void;
+
+use crate::{
+    bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_HASH},
+    helpers::bpf_map_lookup_elem,
+};
+
+#[repr(transparent)]
+pub struct HashMap<K, V> {
+    def: bpf_map_def,
+    _k: PhantomData<K>,
+    _v: PhantomData<V>,
+}
+
+impl<K, V> HashMap<K, V> {
+    pub const fn with_max_entries(max_entries: u32, flags: u32) -> HashMap<K, V> {
+        HashMap {
+            def: bpf_map_def {
+                type_: BPF_MAP_TYPE_HASH,
+                key_size: mem::size_of::<K>() as u32,
+                value_size: mem::size_of::<V>() as u32,
+                max_entries,
+                map_flags: flags,
+            },
+            _k: PhantomData,
+            _v: PhantomData,
+        }
+    }
+
+    pub unsafe fn get(&mut self, key: &K) -> Option<&V> {
+        let value = bpf_map_lookup_elem(
+            &mut self.def as *mut _ as *mut _,
+            key as *const _ as *const c_void,
+        );
+        if value.is_null() {
+            None
+        } else {
+            // FIXME: alignment
+            Some(&*(value as *const V))
+        }
+    }
+}

bpf/aya-bpf/src/maps/mod.rs

@@ -1,3 +1,9 @@
+pub mod array;
+pub mod hash_map;
 pub mod perf_map;
+pub mod sock_hash;
 
+pub use array::Array;
+pub use hash_map::HashMap;
 pub use perf_map::PerfMap;
+pub use sock_hash::SockHash;

bpf/aya-bpf/src/maps/sock_hash.rs

@@ -0,0 +1,48 @@
+use core::{marker::PhantomData, mem};
+
+use aya_bpf_cty::c_void;
+
+use crate::{
+    bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_SOCKHASH, bpf_sock_ops},
+    helpers::bpf_sock_hash_update,
+};
+
+#[repr(transparent)]
+pub struct SockHash<K> {
+    def: bpf_map_def,
+    _k: PhantomData<K>,
+}
+
+impl<K> SockHash<K> {
+    pub const fn with_max_entries(max_entries: u32, flags: u32) -> SockHash<K> {
+        SockHash {
+            def: bpf_map_def {
+                type_: BPF_MAP_TYPE_SOCKHASH,
+                key_size: mem::size_of::<K>() as u32,
+                value_size: mem::size_of::<u32>() as u32,
+                max_entries,
+                map_flags: flags,
+            },
+            _k: PhantomData,
+        }
+    }
+
+    pub unsafe fn update(
+        &mut self,
+        key: &mut K,
+        sk_ops: *mut bpf_sock_ops,
+        flags: u64,
+    ) -> Result<(), i64> {
+        let ret = bpf_sock_hash_update(
+            sk_ops,
+            &mut self.def as *mut _ as *mut _,
+            key as *mut _ as *mut c_void,
+            flags,
+        );
+        if ret < 0 {
+            Err(ret)
+        } else {
+            Ok(())
+        }
+    }
+}

bpf/aya-bpf/src/programs/mod.rs

@@ -1,3 +1,11 @@
 pub mod probe;
+pub mod sk_msg;
+pub mod sk_skb;
+pub mod sock_ops;
+pub mod xdp;
 
 pub use probe::ProbeContext;
+pub use sk_msg::SkMsgContext;
+pub use sk_skb::SkSkbContext;
+pub use sock_ops::SockOpsContext;
+pub use xdp::XdpContext;

bpf/aya-bpf/src/programs/probe.rs

@@ -3,7 +3,7 @@ use core::ffi::c_void;
 use crate::{bindings::pt_regs, BpfContext};
 
 pub struct ProbeContext {
-    regs: *mut pt_regs,
+    pub regs: *mut pt_regs,
 }
 
 impl ProbeContext {

bpf/aya-bpf/src/programs/sk_msg.rs

@@ -0,0 +1,39 @@
+use core::ffi::c_void;
+
+use crate::{
+    bindings::sk_msg_md,
+    helpers::{bpf_msg_pop_data, bpf_msg_push_data},
+    BpfContext,
+};
+
+pub struct SkMsgContext {
+    msg: *mut sk_msg_md,
+}
+
+impl SkMsgContext {
+    pub fn new(msg: *mut sk_msg_md) -> SkMsgContext {
+        SkMsgContext { msg }
+    }
+
+    pub fn data(&self) -> usize {
+        unsafe { (*self.msg).__bindgen_anon_1.data as usize }
+    }
+
+    pub fn data_end(&self) -> usize {
+        unsafe { (*self.msg).__bindgen_anon_2.data_end as usize }
+    }
+
+    pub fn push_data(&self, start: u32, len: u32, flags: u64) -> i64 {
+        unsafe { bpf_msg_push_data(self.msg, start, len as u32, flags) }
+    }
+
+    pub fn pop_data(&self, start: u32, len: u32, flags: u64) -> i64 {
+        unsafe { bpf_msg_pop_data(self.msg, start, len as u32, flags) }
+    }
+}
+
+impl BpfContext for SkMsgContext {
+    fn as_ptr(&self) -> *mut c_void {
+        self.msg as *mut _
+    }
+}

bpf/aya-bpf/src/programs/sk_skb.rs

@@ -0,0 +1,19 @@
+use core::ffi::c_void;
+
+use crate::{bindings::__sk_buff, BpfContext};
+
+pub struct SkSkbContext {
+    skb: *mut __sk_buff,
+}
+
+impl SkSkbContext {
+    pub fn new(skb: *mut __sk_buff) -> SkSkbContext {
+        SkSkbContext { skb }
+    }
+}
+
+impl BpfContext for SkSkbContext {
+    fn as_ptr(&self) -> *mut c_void {
+        self.skb as *mut _
+    }
+}

bpf/aya-bpf/src/programs/sock_ops.rs

@@ -0,0 +1,23 @@
+use core::ffi::c_void;
+
+use crate::{bindings::bpf_sock_ops, BpfContext};
+
+pub struct SockOpsContext {
+    ops: *mut bpf_sock_ops,
+}
+
+impl SockOpsContext {
+    pub fn new(ops: *mut bpf_sock_ops) -> SockOpsContext {
+        SockOpsContext { ops }
+    }
+
+    pub fn op(&self) -> u32 {
+        unsafe { (*self.ops).op }
+    }
+}
+
+impl BpfContext for SockOpsContext {
+    fn as_ptr(&self) -> *mut c_void {
+        self.ops as *mut _
+    }
+}

bpf/aya-bpf/src/programs/xdp.rs

@@ -0,0 +1,27 @@
+use core::ffi::c_void;
+
+use crate::{bindings::xdp_md, BpfContext};
+
+pub struct XdpContext {
+    ctx: *mut xdp_md,
+}
+
+impl XdpContext {
+    pub fn new(ctx: *mut xdp_md) -> XdpContext {
+        XdpContext { ctx }
+    }
+
+    pub fn data(&self) -> usize {
+        unsafe { (*self.ctx).data as usize }
+    }
+
+    pub fn data_end(&self) -> usize {
+        unsafe { (*self.ctx).data_end as usize }
+    }
+}
+
+impl BpfContext for XdpContext {
+    fn as_ptr(&self) -> *mut c_void {
+        self.ctx as *mut _
+    }
+}