Merge pull request #265 from dave-tucker/sklookup

Add support for BPF_PROG_TYPE_SK_LOOKUP

aya/src/bpf.rs

@@ -24,8 +24,8 @@ use crate::{
     programs::{
         BtfTracePoint, CgroupSkb, CgroupSkbAttachType, CgroupSockAddr, CgroupSockopt, CgroupSysctl,
         Extension, FEntry, FExit, KProbe, LircMode2, Lsm, PerfEvent, ProbeKind, Program,
-        ProgramData, ProgramError, RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind,
+        ProgramData, ProgramError, RawTracePoint, SchedClassifier, SkLookup, SkMsg, SkSkb,
-        SockOps, SocketFilter, TracePoint, UProbe, Xdp,
+        SkSkbKind, SockOps, SocketFilter, TracePoint, UProbe, Xdp,
     },
     sys::{
         bpf_load_btf, bpf_map_freeze, bpf_map_update_elem_ptr, is_btf_datasec_supported,
@@ -521,6 +521,9 @@ impl<'a> BpfLoader<'a> {
                         ProgramSection::Extension { .. } => Program::Extension(Extension {
                             data: ProgramData::new(prog_name, obj, btf_fd),
                         }),
+                        ProgramSection::SkLookup { .. } => Program::SkLookup(SkLookup {
+                            data: ProgramData::new(prog_name, obj, btf_fd),
+                        }),
                     }
                 };
                 (name, program)

aya/src/obj/mod.rs

@@ -185,6 +185,9 @@ pub enum ProgramSection {
     Extension {
         name: String,
     },
+    SkLookup {
+        name: String,
+    },
 }
 
 impl ProgramSection {
@@ -216,6 +219,7 @@ impl ProgramSection {
             ProgramSection::FEntry { name } => name,
             ProgramSection::FExit { name } => name,
             ProgramSection::Extension { name } => name,
+            ProgramSection::SkLookup { name } => name,
         }
     }
 }
@@ -361,6 +365,7 @@ impl FromStr for ProgramSection {
             "fentry" => FEntry { name },
             "fexit" => FExit { name },
             "freplace" => Extension { name },
+            "sk_lookup" => SkLookup { name },
             _ => {
                 return Err(ParseError::InvalidProgramSection {
                     section: section.to_owned(),

aya/src/programs/mod.rs

@@ -50,11 +50,12 @@ pub mod lsm;
 pub mod perf_attach;
 pub mod perf_event;
 mod probe;
-pub mod raw_trace_point;
+mod raw_trace_point;
-pub mod sk_msg;
+mod sk_lookup;
-pub mod sk_skb;
+mod sk_msg;
-pub mod sock_ops;
+mod sk_skb;
-pub mod socket_filter;
+mod sock_ops;
+mod socket_filter;
 pub mod tc;
 pub mod tp_btf;
 pub mod trace_point;
@@ -88,6 +89,7 @@ use perf_attach::*;
 pub use perf_event::{PerfEvent, PerfEventScope, PerfTypeId, SamplePolicy};
 pub use probe::ProbeKind;
 pub use raw_trace_point::RawTracePoint;
+pub use sk_lookup::SkLookup;
 pub use sk_msg::SkMsg;
 pub use sk_skb::{SkSkb, SkSkbKind};
 pub use sock_ops::SockOps;
@@ -261,6 +263,8 @@ pub enum Program {
     FExit(FExit),
     /// A [`Extension`] program
     Extension(Extension),
+    /// A [`SkLookup`] program
+    SkLookup(SkLookup),
 }
 
 impl Program {
@@ -289,6 +293,7 @@ impl Program {
             Program::FExit(_) => BPF_PROG_TYPE_TRACING,
             Program::Extension(_) => BPF_PROG_TYPE_EXT,
             Program::CgroupSockAddr(_) => BPF_PROG_TYPE_CGROUP_SOCK_ADDR,
+            Program::SkLookup(_) => BPF_PROG_TYPE_SK_LOOKUP,
         }
     }
 
@@ -316,6 +321,7 @@ impl Program {
             Program::FExit(p) => p.data.pin(path),
             Program::Extension(p) => p.data.pin(path),
             Program::CgroupSockAddr(p) => p.data.pin(path),
+            Program::SkLookup(p) => p.data.pin(path),
         }
     }
 }
@@ -523,6 +529,7 @@ impl ProgramFd for Program {
             Program::FExit(p) => p.data.fd,
             Program::Extension(p) => p.data.fd,
             Program::CgroupSockAddr(p) => p.data.fd,
+            Program::SkLookup(p) => p.data.fd,
         }
     }
 }
@@ -572,6 +579,7 @@ impl_program_fd!(
     FExit,
     Extension,
     CgroupSockAddr,
+    SkLookup,
 );
 
 macro_rules! impl_try_from_program {
@@ -624,6 +632,7 @@ impl_try_from_program!(
     FExit,
     Extension,
     CgroupSockAddr,
+    SkLookup,
 );
 
 /// Provides information about a loaded program, like name, id and statistics

aya/src/programs/sk_lookup.rs

@@ -0,0 +1,103 @@
+use std::os::unix::prelude::{AsRawFd, RawFd};
+
+use crate::{
+    generated::{bpf_attach_type::BPF_SK_LOOKUP, bpf_prog_type::BPF_PROG_TYPE_SK_LOOKUP},
+    programs::{define_link_wrapper, load_program, FdLinkId, OwnedLink, ProgramData, ProgramError},
+    sys::bpf_link_create,
+};
+
+use super::links::FdLink;
+
+/// A program used to redirect incoming packets to a local socket.
+///
+/// [`SkLookup`] programs are attached to network namespaces to provide programmable
+/// socket lookup for TCP/UDP when a packet is to be delievered locally.
+///
+/// You may attach multiple programs to the same namespace and they are executed
+/// in the order they were attached.
+///
+/// # Minimum kernel version
+///
+/// The minimum kernel version required to use this feature is 5.9.
+///
+/// # Examples
+///
+/// ```no_run
+/// # #[derive(Debug, thiserror::Error)]
+/// # enum Error {
+/// #     #[error(transparent)]
+/// #     IO(#[from] std::io::Error),
+/// #     #[error(transparent)]
+/// #     Map(#[from] aya::maps::MapError),
+/// #     #[error(transparent)]
+/// #     Program(#[from] aya::programs::ProgramError),
+/// #     #[error(transparent)]
+/// #     Bpf(#[from] aya::BpfError)
+/// # }
+/// # let mut bpf = aya::Bpf::load(&[])?;
+/// use std::fs::File;
+/// use std::convert::TryInto;
+/// use aya::programs::SkLookup;
+///
+/// let file = File::open("/var/run/netns/test")?;
+/// let program: &mut SkLookup = bpf.program_mut("sk_lookup").unwrap().try_into()?;
+/// program.load()?;
+/// program.attach(file)?;
+/// # Ok::<(), Error>(())
+/// ```
+#[derive(Debug)]
+#[doc(alias = "BPF_PROG_TYPE_SK_LOOKUP")]
+pub struct SkLookup {
+    pub(crate) data: ProgramData<SkLookupLink>,
+}
+
+impl SkLookup {
+    /// Loads the program inside the kernel.
+    pub fn load(&mut self) -> Result<(), ProgramError> {
+        self.data.expected_attach_type = Some(BPF_SK_LOOKUP);
+        load_program(BPF_PROG_TYPE_SK_LOOKUP, &mut self.data)
+    }
+
+    /// Attaches the program to the given network namespace.
+    ///
+    /// The returned value can be used to detach, see [SkLookup::detach].
+    pub fn attach<T: AsRawFd>(&mut self, netns: T) -> Result<SkLookupLinkId, ProgramError> {
+        let prog_fd = self.data.fd_or_err()?;
+        let netns_fd = netns.as_raw_fd();
+
+        let link_fd = bpf_link_create(prog_fd, netns_fd, BPF_SK_LOOKUP, None, 0).map_err(
+            |(_, io_error)| ProgramError::SyscallError {
+                call: "bpf_link_create".to_owned(),
+                io_error,
+            },
+        )? as RawFd;
+        self.data.links.insert(SkLookupLink(FdLink::new(link_fd)))
+    }
+
+    /// Takes ownership of the link referenced by the provided link_id.
+    ///
+    /// The link will be detached on `Drop` and the caller is now responsible
+    /// for managing its lifetime.
+    pub fn forget_link(
+        &mut self,
+        link_id: SkLookupLinkId,
+    ) -> Result<OwnedLink<SkLookupLink>, ProgramError> {
+        Ok(OwnedLink::new(self.data.forget_link(link_id)?))
+    }
+
+    /// Detaches the program.
+    ///
+    /// See [SkLookup::attach].
+    pub fn detach(&mut self, link_id: SkLookupLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
+}
+
+define_link_wrapper!(
+    /// The link used by [SkLookup] programs.
+    SkLookupLink,
+    /// The type returned by [SkLookup::attach]. Can be passed to [SkLookup::detach].
+    SkLookupLinkId,
+    FdLink,
+    FdLinkId
+);

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

@@ -743,6 +743,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,8 +2,8 @@ mod expand;
 
 use expand::{
     Args, BtfTracePoint, CgroupSkb, CgroupSockAddr, CgroupSockopt, CgroupSysctl, FEntry, FExit,
-    Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind,
+    Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkLookup, SkMsg, SkSkb,
-    SockAddrArgs, SockOps, SocketFilter, SockoptArgs, TracePoint, Xdp,
+    SkSkbKind, SockAddrArgs, SockOps, SocketFilter, SockoptArgs, TracePoint, Xdp,
 };
 use proc_macro::TokenStream;
 use syn::{parse_macro_input, ItemFn, ItemStatic};
@@ -467,3 +467,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,12 @@ 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_map_lookup_elem, bpf_msg_redirect_hash, bpf_sk_assign, bpf_sk_redirect_hash,
+        bpf_sk_release, bpf_sock_hash_update,
+    },
     maps::PinningType,
-    programs::{SkBuffContext, SkMsgContext},
+    programs::{SkBuffContext, SkLookupContext, SkMsgContext},
     BpfContext,
 };
 
@@ -85,4 +88,24 @@ impl<K> SockHash<K> {
             )
         }
     }
+
+    pub fn redirect_sk_lookup(
+        &mut self,
+        ctx: &SkLookupContext,
+        key: K,
+        flags: u64,
+    ) -> Result<(), u32> {
+        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(1);
+            }
+            let ret = bpf_sk_assign(ctx.as_ptr() as *mut _, sk, flags);
+            bpf_sk_release(sk);
+            (ret >= 0).then(|| ()).ok_or(1)
+        }
+    }
 }

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

@@ -4,9 +4,12 @@ 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_map_lookup_elem, bpf_msg_redirect_map, bpf_sk_assign, bpf_sk_redirect_map,
+        bpf_sk_release, bpf_sock_map_update,
+    },
     maps::PinningType,
-    programs::{SkBuffContext, SkMsgContext},
+    programs::{SkBuffContext, SkLookupContext, SkMsgContext},
     BpfContext,
 };
 
@@ -80,4 +83,24 @@ impl SockMap {
             flags,
         )
     }
+
+    pub fn redirect_sk_lookup(
+        &mut self,
+        ctx: &SkLookupContext,
+        index: u32,
+        flags: u64,
+    ) -> Result<(), u32> {
+        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(1);
+            }
+            let ret = bpf_sk_assign(ctx.as_ptr() as *mut _, sk, flags);
+            bpf_sk_release(sk);
+            (ret >= 0).then(|| ()).ok_or(1)
+        }
+    }
 }

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;
@@ -21,6 +22,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 _
+    }
+}