aya: add support for BPF_PROG_TYPE_SK_REUSEPORT

aya-ebpf-macros/src/lib.rs

@@ -16,6 +16,7 @@ mod perf_event;
 mod raw_tracepoint;
 mod sk_lookup;
 mod sk_msg;
+mod sk_reuseport;
 mod sk_skb;
 mod sock_ops;
 mod socket_filter;
@@ -42,6 +43,7 @@ use proc_macro::TokenStream;
 use raw_tracepoint::RawTracePoint;
 use sk_lookup::SkLookup;
 use sk_msg::SkMsg;
+use sk_reuseport::SkReuseport;
 use sk_skb::{SkSkb, SkSkbKind};
 use sock_ops::SockOps;
 use socket_filter::SocketFilter;
@@ -612,6 +614,67 @@ pub fn sk_lookup(attrs: TokenStream, item: TokenStream) -> TokenStream {
     .into()
 }
 
+/// Marks a function as an eBPF Socket Reuseport program that can be attached to
+/// a socket with SO_REUSEPORT set.
+///
+/// # Minimum kernel version
+///
+/// The minimum kernel version required to use this feature is 4.19
+///
+/// # Examples
+///
+/// Basic usage allowing kernel to handle socket selection:
+///
+/// ```no_run
+/// use aya_ebpf::{macros::sk_reuseport, programs::{SkReuseportContext, SK_PASS, SK_DROP}};
+///
+/// #[sk_reuseport]
+/// pub fn select_socket(_ctx: SkReuseportContext) -> u32 {
+///     // Return SK_DROP to drop packet, SK_PASS to let kernel handle selection
+///     SK_PASS
+/// }
+/// ```
+///
+/// Advanced usage with custom socket selection:
+///
+/// ```no_run
+/// use aya_ebpf::{
+///     macros::{sk_reuseport, map},
+///     programs::{SkReuseportContext, SK_PASS, SK_DROP},
+///     helpers::bpf_sk_select_reuseport,
+///     maps::ReusePortSockArray,
+///     EbpfContext,
+/// };
+///
+/// #[map(name = "socket_map")]
+/// static SOCKET_MAP: ReusePortSockArray = ReusePortSockArray::with_max_entries(4, 0);
+///
+/// #[sk_reuseport]
+/// pub fn load_balance(ctx: SkReuseportContext) -> u32 {
+///     // Use packet hash for consistent load balancing
+///     let socket_idx = ctx.hash() % 4;
+///     
+///     let ret = unsafe {
+///         bpf_sk_select_reuseport(
+///             ctx.as_ptr() as *mut _,
+///             SOCKET_MAP.as_ptr(),
+///             &socket_idx as *const _ as *mut _,
+///             0
+///         )
+///     };
+///     
+///     if ret == 0 { SK_PASS } else { SK_DROP }
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn sk_reuseport(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    match SkReuseport::parse(attrs.into(), item.into()) {
+        Ok(prog) => prog.expand(),
+        Err(err) => err.emit_as_expr_tokens(),
+    }
+    .into()
+}
+
 /// Marks a function as a cgroup device eBPF program that can be attached to a
 /// cgroup.
 ///

aya-ebpf-macros/src/sk_reuseport.rs

@@ -0,0 +1,71 @@
+use proc_macro2::TokenStream;
+use proc_macro2_diagnostics::{Diagnostic, SpanDiagnosticExt as _};
+use quote::quote;
+use syn::{ItemFn, spanned::Spanned as _};
+
+pub(crate) struct SkReuseport {
+    item: ItemFn,
+}
+
+impl SkReuseport {
+    pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self, Diagnostic> {
+        if !attrs.is_empty() {
+            return Err(attrs.span().error("unexpected attribute"));
+        }
+        let item = syn::parse2(item)?;
+        Ok(Self { item })
+    }
+
+    pub(crate) fn expand(&self) -> TokenStream {
+        let Self { item } = self;
+        let ItemFn {
+            attrs: _,
+            vis,
+            sig,
+            block: _,
+        } = item;
+        let fn_name = &sig.ident;
+        quote! {
+            #[unsafe(no_mangle)]
+            #[unsafe(link_section = "sk_reuseport")]
+            #vis fn #fn_name(ctx: *mut ::aya_ebpf::bindings::sk_reuseport_md) -> u32 {
+                return #fn_name(::aya_ebpf::programs::SkReuseportContext::new(ctx));
+
+                #item
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use syn::parse_quote;
+
+    use super::*;
+
+    #[test]
+    fn test_sk_reuseport() {
+        let prog = SkReuseport::parse(
+            parse_quote! {},
+            parse_quote! {
+                fn prog(ctx: &mut ::aya_ebpf::programs::SkReuseportContext) -> u32 {
+                    0
+                }
+            },
+        )
+        .unwrap();
+        let expanded = prog.expand();
+        let expected = quote! {
+            #[unsafe(no_mangle)]
+            #[unsafe(link_section = "sk_reuseport")]
+            fn prog(ctx: *mut ::aya_ebpf::bindings::sk_reuseport_md) -> u32 {
+                return prog(::aya_ebpf::programs::SkReuseportContext::new(ctx));
+
+                fn prog(ctx: &mut ::aya_ebpf::programs::SkReuseportContext) -> u32 {
+                    0
+                }
+            }
+        };
+        assert_eq!(expected.to_string(), expanded.to_string());
+    }
+}

aya-obj/src/obj.rs

@@ -215,7 +215,7 @@ pub struct Function {
 /// - `lwt_in`, `lwt_out`, `lwt_seg6local`, `lwt_xmit`
 /// - `raw_tp.w+`, `raw_tracepoint.w+`
 /// - `action`
-/// - `sk_reuseport/migrate`, `sk_reuseport`
+/// - `sk_reuseport/migrate`
 /// - `syscall`
 /// - `struct_ops+`
 /// - `fmod_ret+`, `fmod_ret.s+`
@@ -268,6 +268,7 @@ pub enum ProgramSection {
     FlowDissector,
     Extension,
     SkLookup,
+    SkReuseport,
     CgroupSock {
         attach_type: CgroupSockAttachType,
     },
@@ -427,6 +428,7 @@ impl FromStr for ProgramSection {
             "flow_dissector" => FlowDissector,
             "freplace" => Extension,
             "sk_lookup" => SkLookup,
+            "sk_reuseport" => SkReuseport,
             "iter" => Iter { sleepable: false },
             "iter.s" => Iter { sleepable: true },
             _ => {

aya/src/bpf.rs

@@ -25,8 +25,8 @@ use crate::{
         BtfTracePoint, CgroupDevice, CgroupSkb, CgroupSkbAttachType, CgroupSock, CgroupSockAddr,
         CgroupSockopt, CgroupSysctl, Extension, FEntry, FExit, FlowDissector, Iter, KProbe,
         LircMode2, Lsm, PerfEvent, ProbeKind, Program, ProgramData, ProgramError, RawTracePoint,
-        SchedClassifier, SkLookup, SkMsg, SkSkb, SkSkbKind, SockOps, SocketFilter, TracePoint,
-        UProbe, Xdp,
+        SchedClassifier, SkLookup, SkMsg, SkReuseport, SkSkb, SkSkbKind, SockOps, SocketFilter,
+        TracePoint, UProbe, Xdp,
     },
     sys::{
         bpf_load_btf, is_bpf_cookie_supported, is_bpf_global_data_supported,
@@ -429,6 +429,7 @@ impl<'a> EbpfLoader<'a> {
                                 | ProgramSection::PerfEvent
                                 | ProgramSection::RawTracePoint
                                 | ProgramSection::SkLookup
+                                | ProgramSection::SkReuseport
                                 | ProgramSection::FlowDissector
                                 | ProgramSection::CgroupSock { attach_type: _ }
                                 | ProgramSection::CgroupDevice => {}
@@ -670,6 +671,9 @@ impl<'a> EbpfLoader<'a> {
                         ProgramSection::SkLookup => Program::SkLookup(SkLookup {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
                         }),
+                        ProgramSection::SkReuseport => Program::SkReuseport(SkReuseport {
+                            data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
+                        }),
                         ProgramSection::CgroupSock { attach_type, .. } => {
                             Program::CgroupSock(CgroupSock {
                                 data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
@@ -716,6 +720,7 @@ fn parse_map(
         BPF_MAP_TYPE_PERCPU_HASH => Map::PerCpuHashMap(map),
         BPF_MAP_TYPE_LRU_PERCPU_HASH => Map::PerCpuLruHashMap(map),
         BPF_MAP_TYPE_PERF_EVENT_ARRAY => Map::PerfEventArray(map),
+        BPF_MAP_TYPE_REUSEPORT_SOCKARRAY => Map::ReusePortSockArray(map),
         BPF_MAP_TYPE_RINGBUF => Map::RingBuf(map),
         BPF_MAP_TYPE_SOCKHASH => Map::SockHash(map),
         BPF_MAP_TYPE_SOCKMAP => Map::SockMap(map),

aya/src/maps/mod.rs

@@ -95,7 +95,7 @@ pub use lpm_trie::LpmTrie;
 pub use perf::PerfEventArray;
 pub use queue::Queue;
 pub use ring_buf::RingBuf;
-pub use sock::{SockHash, SockMap};
+pub use sock::{ReusePortSockArray, SockHash, SockMap};
 pub use stack::Stack;
 pub use stack_trace::StackTraceMap;
 pub use xdp::{CpuMap, DevMap, DevMapHash, XskMap};
@@ -297,6 +297,8 @@ pub enum Map {
     ProgramArray(MapData),
     /// A [`Queue`] map.
     Queue(MapData),
+    /// A [`ReusePortSockArray`] map.
+    ReusePortSockArray(MapData),
     /// A [`RingBuf`] map.
     RingBuf(MapData),
     /// A [`SockHash`] map
@@ -331,6 +333,7 @@ impl Map {
             Self::PerfEventArray(map) => map.obj.map_type(),
             Self::ProgramArray(map) => map.obj.map_type(),
             Self::Queue(map) => map.obj.map_type(),
+            Self::ReusePortSockArray(map) => map.obj.map_type(),
             Self::RingBuf(map) => map.obj.map_type(),
             Self::SockHash(map) => map.obj.map_type(),
             Self::SockMap(map) => map.obj.map_type(),
@@ -362,6 +365,7 @@ impl Map {
             Self::ProgramArray(map) => map.pin(path),
             Self::Queue(map) => map.pin(path),
             Self::RingBuf(map) => map.pin(path),
+            Self::ReusePortSockArray(map) => map.pin(path),
             Self::SockHash(map) => map.pin(path),
             Self::SockMap(map) => map.pin(path),
             Self::Stack(map) => map.pin(path),
@@ -480,6 +484,7 @@ impl_try_from_map!(() {
     DevMapHash,
     PerfEventArray,
     ProgramArray,
+    ReusePortSockArray,
     RingBuf,
     SockMap,
     StackTraceMap,

aya/src/maps/sock/mod.rs

@@ -1,4 +1,5 @@
 //! Socket maps.
+mod reuseport_sock_array;
 mod sock_hash;
 mod sock_map;
 
@@ -7,6 +8,7 @@ use std::{
     os::fd::{AsFd, BorrowedFd},
 };
 
+pub use reuseport_sock_array::ReusePortSockArray;
 pub use sock_hash::SockHash;
 pub use sock_map::SockMap;