aya: add struct_ops support

aya-ebpf-macros/src/lib.rs

@@ -23,6 +23,7 @@ mod sk_msg;
 mod sk_skb;
 mod sock_ops;
 mod socket_filter;
+mod struct_ops;
 mod tc;
 mod tracepoint;
 mod uprobe;
@@ -51,6 +52,7 @@ use sk_msg::SkMsg;
 use sk_skb::{SkSkb, SkSkbKind};
 use sock_ops::SockOps;
 use socket_filter::SocketFilter;
+use struct_ops::StructOps;
 use tc::SchedClassifier;
 use tracepoint::TracePoint;
 use uprobe::{UProbe, UProbeKind};
@@ -601,6 +603,40 @@ pub fn fexit(attrs: TokenStream, item: TokenStream) -> TokenStream {
     .into()
 }
 
+/// Marks a function as an eBPF struct_ops program.
+///
+/// Struct ops programs implement kernel interfaces like `hid_bpf_ops` for
+/// HID device handling or `sched_ext_ops` for custom schedulers.
+///
+/// # Arguments
+///
+/// * `name` - Optional name for the struct_ops callback. If not specified,
+///   the function name is used.
+/// * `sleepable` - Mark the program as sleepable.
+///
+/// # Minimum kernel version
+///
+/// The minimum kernel version required to use this feature is 5.6.
+///
+/// # Example
+///
+/// ```no_run
+/// use aya_ebpf::{macros::struct_ops, programs::StructOpsContext};
+///
+/// #[struct_ops]
+/// fn my_callback(ctx: StructOpsContext) -> i32 {
+///     0
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn struct_ops(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    match StructOps::parse(attrs.into(), item.into()) {
+        Ok(prog) => prog.expand(),
+        Err(err) => err.into_compile_error(),
+    }
+    .into()
+}
+
 /// Marks a function as an eBPF Flow Dissector program.
 ///
 /// Flow dissector is a program type that parses metadata out of the packets.

aya-ebpf-macros/src/struct_ops.rs

@@ -0,0 +1,151 @@
+use std::borrow::Cow;
+
+use proc_macro2::TokenStream;
+use quote::quote;
+use syn::{ItemFn, Result};
+
+use crate::args::{err_on_unknown_args, pop_bool_arg, pop_string_arg};
+
+pub(crate) struct StructOps {
+    item: ItemFn,
+    name: Option<String>,
+    sleepable: bool,
+}
+
+impl StructOps {
+    pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self> {
+        let item = syn::parse2(item)?;
+        let mut args = syn::parse2(attrs)?;
+        let name = pop_string_arg(&mut args, "name");
+        let sleepable = pop_bool_arg(&mut args, "sleepable");
+        err_on_unknown_args(&args)?;
+        Ok(Self {
+            item,
+            name,
+            sleepable,
+        })
+    }
+
+    pub(crate) fn expand(&self) -> TokenStream {
+        let Self {
+            item,
+            name,
+            sleepable,
+        } = self;
+        let ItemFn {
+            attrs: _,
+            vis,
+            sig,
+            block: _,
+        } = item;
+        let section_prefix = if *sleepable {
+            "struct_ops.s"
+        } else {
+            "struct_ops"
+        };
+        let fn_name = &sig.ident;
+        let section_name: Cow<'_, _> = if let Some(name) = name {
+            format!("{section_prefix}/{name}").into()
+        } else {
+            format!("{section_prefix}/{fn_name}").into()
+        };
+        quote! {
+            #[unsafe(no_mangle)]
+            #[unsafe(link_section = #section_name)]
+            #vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 {
+                return #fn_name(::aya_ebpf::programs::StructOpsContext::new(ctx));
+
+                #item
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use syn::parse_quote;
+
+    use super::*;
+
+    #[test]
+    fn test_struct_ops() {
+        let prog = StructOps::parse(
+            parse_quote! {},
+            parse_quote! {
+                fn my_callback(ctx: &mut aya_ebpf::programs::StructOpsContext) -> i32 {
+                    0
+                }
+            },
+        )
+        .unwrap();
+        let expanded = prog.expand();
+        let expected = quote! {
+            #[unsafe(no_mangle)]
+            #[unsafe(link_section = "struct_ops/my_callback")]
+            fn my_callback(ctx: *mut ::core::ffi::c_void) -> i32 {
+                return my_callback(::aya_ebpf::programs::StructOpsContext::new(ctx));
+
+                fn my_callback(ctx: &mut aya_ebpf::programs::StructOpsContext) -> i32 {
+                    0
+                }
+            }
+        };
+        assert_eq!(expected.to_string(), expanded.to_string());
+    }
+
+    #[test]
+    fn test_struct_ops_with_name() {
+        let prog = StructOps::parse(
+            parse_quote! {
+                name = "hid_device_event"
+            },
+            parse_quote! {
+                fn my_handler(ctx: &mut aya_ebpf::programs::StructOpsContext) -> i32 {
+                    0
+                }
+            },
+        )
+        .unwrap();
+        let expanded = prog.expand();
+        let expected = quote! {
+            #[unsafe(no_mangle)]
+            #[unsafe(link_section = "struct_ops/hid_device_event")]
+            fn my_handler(ctx: *mut ::core::ffi::c_void) -> i32 {
+                return my_handler(::aya_ebpf::programs::StructOpsContext::new(ctx));
+
+                fn my_handler(ctx: &mut aya_ebpf::programs::StructOpsContext) -> i32 {
+                    0
+                }
+            }
+        };
+        assert_eq!(expected.to_string(), expanded.to_string());
+    }
+
+    #[test]
+    fn test_struct_ops_sleepable() {
+        let prog = StructOps::parse(
+            parse_quote! {
+                sleepable
+            },
+            parse_quote! {
+                fn my_callback(ctx: &mut aya_ebpf::programs::StructOpsContext) -> i32 {
+                    0
+                }
+            },
+        )
+        .unwrap();
+        let expanded = prog.expand();
+        let expected = quote! {
+            #[unsafe(no_mangle)]
+            #[unsafe(link_section = "struct_ops.s/my_callback")]
+            fn my_callback(ctx: *mut ::core::ffi::c_void) -> i32 {
+                return my_callback(::aya_ebpf::programs::StructOpsContext::new(ctx));
+
+                fn my_callback(ctx: &mut aya_ebpf::programs::StructOpsContext) -> i32 {
+                    0
+                }
+            }
+        };
+        assert_eq!(expected.to_string(), expanded.to_string());
+    }
+}

aya-obj/src/btf/btf.rs

@@ -448,7 +448,67 @@ impl Btf {
         })
     }
 
-    pub(crate) fn type_size(&self, root_type_id: u32) -> Result<usize, BtfError> {
+    /// Returns the index of a struct member by name.
+    ///
+    /// This is useful for struct_ops programs where `expected_attach_type` needs
+    /// to be set to the member index in the kernel struct.
+    pub fn struct_member_index(
+        &self,
+        struct_type_id: u32,
+        member_name: &str,
+    ) -> Result<u32, BtfError> {
+        let ty = self.types.type_by_id(struct_type_id)?;
+        let members = match ty {
+            BtfType::Struct(s) => &s.members,
+            _ => {
+                return Err(BtfError::UnexpectedBtfType {
+                    type_id: struct_type_id,
+                });
+            }
+        };
+        for (index, member) in members.iter().enumerate() {
+            let name = self.string_at(member.name_offset)?;
+            if name == member_name {
+                return Ok(index as u32);
+            }
+        }
+        Err(BtfError::UnknownBtfTypeName {
+            type_name: member_name.to_owned(),
+        })
+    }
+
+    /// Returns the byte offset of a struct member by name.
+    ///
+    /// This is useful for struct_ops where we need to find the offset of the
+    /// `data` field within the kernel wrapper struct (e.g., `bpf_struct_ops_hid_bpf_ops`).
+    pub fn struct_member_byte_offset(
+        &self,
+        struct_type_id: u32,
+        member_name: &str,
+    ) -> Result<u32, BtfError> {
+        let ty = self.types.type_by_id(struct_type_id)?;
+        let members = match ty {
+            BtfType::Struct(s) => &s.members,
+            _ => {
+                return Err(BtfError::UnexpectedBtfType {
+                    type_id: struct_type_id,
+                });
+            }
+        };
+        for member in members.iter() {
+            let name = self.string_at(member.name_offset)?;
+            if name == member_name {
+                // offset is in bits, convert to bytes
+                return Ok(member.offset / 8);
+            }
+        }
+        Err(BtfError::UnknownBtfTypeName {
+            type_name: member_name.to_owned(),
+        })
+    }
+
+    /// Returns the size of a BTF type in bytes.
+    pub fn type_size(&self, root_type_id: u32) -> Result<usize, BtfError> {
         let mut type_id = root_type_id;
         let mut n_elems = 1;
         for () in core::iter::repeat_n((), MAX_RESOLVE_DEPTH) {
@@ -569,8 +629,23 @@ impl Btf {
                         d.name_offset = self.add_string(&fixed_name);
                     }
 
-                    // There are some cases when the compiler does indeed populate the size.
-                    if d.size > 0 {
+                    // .ksyms is a pseudo-section for extern kernel symbols (see EbpfSectionKind::Ksyms).
+                    // It has no real ELF section and the kernel rejects DATASEC with
+                    // size=0 and non-empty entries. Replace with INT to remove it.
+                    const KSYMS_SECTION: &str = ".ksyms";
+                    if name == KSYMS_SECTION {
+                        debug!("{kind} {name}: pseudo-section, replacing with INT");
+                        let old_size = d.type_info_size();
+                        let new_type =
+                            BtfType::Int(Int::new(d.name_offset, 0, IntEncoding::None, 0));
+                        let new_size = new_type.type_info_size();
+                        // Update header to reflect the size change
+                        self.header.type_len =
+                            self.header.type_len - old_size as u32 + new_size as u32;
+                        self.header.str_off = self.header.type_len;
+                        *t = new_type;
+                    } else if d.size > 0 {
+                        // There are some cases when the compiler does indeed populate the size.
                         debug!("{kind} {name}: size fixup not required");
                     } else {
                         // We need to get the size of the section from the ELF file.
@@ -700,6 +775,13 @@ impl Btf {
                     if !features.btf_func {
                         debug!("{kind}: not supported. replacing with TYPEDEF");
                         *t = BtfType::Typedef(Typedef::new(ty.name_offset, ty.btf_type));
+                    } else if ty.linkage() == FuncLinkage::Extern {
+                        // BTF_FUNC_EXTERN is used for kfuncs (kernel functions).
+                        // These are resolved from kernel BTF at load time, so we
+                        // should not include them in the program BTF. Replace with
+                        // TYPEDEF to preserve type info without the problematic linkage.
+                        debug!("{kind} {name}: extern kfunc, replacing with TYPEDEF");
+                        *t = BtfType::Typedef(Typedef::new(ty.name_offset, ty.btf_type));
                     } else if !features.btf_func_global
                         || name == "memset"
                         || name == "memcpy"

aya-obj/src/lib.rs

@@ -41,7 +41,7 @@
 //! let text_sections = std::collections::HashSet::new();
 //! #[cfg(not(feature = "std"))]
 //! let text_sections = hashbrown::HashSet::new();
-//! object.relocate_calls(&text_sections).unwrap();
+//! object.relocate_calls(&text_sections, None).unwrap();
 //! object.relocate_maps(std::iter::empty(), &text_sections).unwrap();
 //!
 //! // Run with rbpf