bpf: Support BPF_PROG_TYPE_SK_LOOKUP

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>

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

@@ -686,6 +686,38 @@ impl FExit {
     }
 }
 
+pub struct SkLookup {
+    item: ItemFn,
+    name: Option<String>,
+}
+
+impl SkLookup {
+    pub fn from_syn(mut args: Args, item: ItemFn) -> Result<SkLookup> {
+        let name = name_arg(&mut args)?;
+
+        Ok(SkLookup { item, name })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        let section_name = if let Some(name) = &self.name {
+            format!("sk_lookup/{}", name)
+        } else {
+            "sk_lookup".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_sk_lookup) -> u32 {
+                return #fn_name(::aya_bpf::programs::SkLookupContext::new(ctx));
+
+                #item
+            }
+        })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use syn::parse_quote;

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

@@ -2,7 +2,7 @@ mod expand;
 
 use expand::{
     Args, BtfTracePoint, CgroupSkb, CgroupSockAddr, CgroupSysctl, FEntry, FExit, Lsm, Map,
-    PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind,
+    PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkLookup, SkMsg, SkSkb, SkSkbKind,
     SockAddrArgs, SockOps, SocketFilter, TracePoint, Xdp,
 };
 use proc_macro::TokenStream;
@@ -455,3 +455,32 @@ pub fn fexit(attrs: TokenStream, item: TokenStream) -> TokenStream {
         .unwrap_or_else(|err| err.to_compile_error())
         .into()
 }
+
+/// Marks a function as an eBPF Socket Lookup program that can be attached to
+/// a network namespace.
+///
+/// # Minimum kernel version
+///
+/// The minimum kernel version required to use this feature is 5.9
+///
+/// # Examples
+///
+/// ```no_run
+/// use aya_bpf::{macros::sk_lookup, programs::SkLookupContext};
+///
+/// #[sk_lookup(name = "redirect")]
+/// pub fn accept_all(_ctx: SkLookupContext) -> u32 {
+///     // use sk_assign to redirect
+///     return 0
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn sk_lookup(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    let args = parse_macro_input!(attrs as Args);
+    let item = parse_macro_input!(item as ItemFn);
+
+    SkLookup::from_syn(args, item)
+        .and_then(|u| u.expand())
+        .unwrap_or_else(|err| err.to_compile_error())
+        .into()
+}

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

@@ -4,9 +4,9 @@ use aya_bpf_cty::c_void;
 
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_SOCKHASH, bpf_sock_ops},
-    helpers::{bpf_msg_redirect_hash, bpf_sk_redirect_hash, bpf_sock_hash_update},
+    helpers::{bpf_msg_redirect_hash, bpf_sk_redirect_hash, bpf_sock_hash_update, bpf_map_lookup_elem, bpf_sk_assign, bpf_sk_release},
     maps::PinningType,
-    programs::{SkBuffContext, SkMsgContext},
+    programs::{SkBuffContext, SkMsgContext, SkLookupContext},
     BpfContext,
 };
 
@@ -85,4 +85,19 @@ impl<K> SockHash<K> {
             )
         }
     }
+
+    pub fn redirect_sk_lookup(&mut self, ctx: &SkLookupContext, key: K, flags: u64) -> Result<(),()> {
+        unsafe {
+            let sk = bpf_map_lookup_elem(
+                &mut self.def as *mut _ as *mut _,
+                &key as *const _ as *const c_void,
+            );
+            if sk.is_null() {
+                return Err(())
+            }
+            let ret = bpf_sk_assign(ctx.as_ptr() as *mut _, sk, flags);
+            bpf_sk_release(sk);
+            (ret >= 0).then(|| ()).ok_or(())
+        }
+    }
 }

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

@@ -4,9 +4,9 @@ use aya_bpf_cty::c_void;
 
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_SOCKMAP, bpf_sock_ops},
-    helpers::{bpf_msg_redirect_map, bpf_sk_redirect_map, bpf_sock_map_update},
+    helpers::{bpf_msg_redirect_map, bpf_sk_redirect_map, bpf_sock_map_update, bpf_map_lookup_elem, bpf_sk_assign, bpf_sk_release},
     maps::PinningType,
-    programs::{SkBuffContext, SkMsgContext},
+    programs::{SkBuffContext, SkLookupContext, SkMsgContext},
     BpfContext,
 };
 
@@ -80,4 +80,19 @@ impl SockMap {
             flags,
         )
     }
+
+    pub fn redirect_sk_lookup(&mut self, ctx: &SkLookupContext, index: u32, flags: u64) -> Result<(),()> {
+        unsafe {
+            let sk = bpf_map_lookup_elem(
+                &mut self.def as *mut _ as *mut _,
+                &index as *const _ as *const c_void,
+            );
+            if sk.is_null() {
+                return Err(())
+            }
+            let ret = bpf_sk_assign(ctx.as_ptr() as *mut _, sk, flags);
+            bpf_sk_release(sk);
+            (ret >= 0).then(|| ()).ok_or(())
+        }
+    }
 }

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

@@ -5,6 +5,7 @@ pub mod perf_event;
 pub mod probe;
 pub mod raw_tracepoint;
 pub mod sk_buff;
+pub mod sk_lookup;
 pub mod sk_msg;
 pub mod sock_addr;
 pub mod sock_ops;
@@ -20,6 +21,7 @@ pub use perf_event::PerfEventContext;
 pub use probe::ProbeContext;
 pub use raw_tracepoint::RawTracePointContext;
 pub use sk_buff::SkBuffContext;
+pub use sk_lookup::SkLookupContext;
 pub use sk_msg::SkMsgContext;
 pub use sock_addr::SockAddrContext;
 pub use sock_ops::SockOpsContext;

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

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