aya: lazily load custom BTF sources

Custom BTF currently has to be read and parsed before configuring an
EbpfLoader, even when the object does not require target BTF.

Add btf_source(), which accepts a callback invoked only for objects with
CO-RE relocations or typed kernel symbols. The callback receives a parser
configured for the object endianness, allowing callers to parse borrowed
data such as a subslice of an mmap without an intermediate allocation.
Cache the first successfully parsed BTF for reuse by the loader, while
keeping btf(&Btf) for callers that share one parsed value across multiple
loaders.

Author: swananan

aya/src/bpf.rs

@@ -113,7 +113,7 @@ pub fn features() -> &'static Features {
 /// ```
 #[derive(Debug)]
 pub struct EbpfLoader<'a> {
-    btf: Option<Cow<'a, Btf>>,
+    btf: TargetBtf<'a>,
     default_map_pin_directory: Option<PathBuf>,
     globals: HashMap<&'a str, (&'a [u8], bool)>,
     // Max entries overrides the max_entries field of the map that matches the provided name
@@ -128,6 +128,34 @@ pub struct EbpfLoader<'a> {
     allow_unsupported_maps: bool,
 }
 
+#[derive(Debug)]
+enum TargetBtf<'a> {
+    System,
+    Parsed(Cow<'a, Btf>),
+    Source(BtfSource),
+}
+
+type BtfParser = dyn Fn(&[u8]) -> Result<Btf, BtfError>;
+// Use `Fn` instead of `FnOnce`: when target BTF loading fails for an object with
+// only weak typed ksyms, loading continues, so a reused loader must be able to
+// retry the source.
+type BtfSourceFn = dyn Fn(&BtfParser) -> Result<Btf, EbpfError>
+    // Preserve `EbpfLoader`'s existing auto traits.
+    + Send
+    + Sync
+    + std::panic::RefUnwindSafe
+    + std::panic::UnwindSafe
+    // Keep captured state independent of `EbpfLoader`'s borrowed configuration.
+    + 'static;
+
+struct BtfSource(Box<BtfSourceFn>);
+
+impl std::fmt::Debug for BtfSource {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("BtfSource").finish_non_exhaustive()
+    }
+}
+
 /// Builder style API for advanced loading of eBPF programs.
 #[deprecated(since = "0.13.0", note = "use `EbpfLoader` instead")]
 pub type BpfLoader<'a> = EbpfLoader<'a>;
@@ -157,7 +185,7 @@ impl<'a> EbpfLoader<'a> {
     /// Creates a new loader instance.
     pub fn new() -> Self {
         Self {
-            btf: None,
+            btf: TargetBtf::System,
             default_map_pin_directory: None,
             globals: HashMap::new(),
             max_entries: HashMap::new(),
@@ -174,6 +202,10 @@ impl<'a> EbpfLoader<'a> {
     /// the object contains CO-RE relocations or typed kernel symbols. Use this
     /// method to load `BTF` from a custom location.
     ///
+    /// The parsed `BTF` can be shared across multiple loaders, avoiding repeated
+    /// parsing. To lazily read and parse target `BTF`, use
+    /// [`EbpfLoader::btf_source`].
+    ///
     /// # Example
     ///
     /// ```no_run
@@ -190,7 +222,45 @@ impl<'a> EbpfLoader<'a> {
     /// # Ok::<(), aya::EbpfError>(())
     /// ```
     pub fn btf(&mut self, btf: &'a Btf) -> &mut Self {
-        self.btf = Some(Cow::Borrowed(btf));
+        self.btf = TargetBtf::Parsed(Cow::Borrowed(btf));
+        self
+    }
+
+    /// Sets a callback for lazily loading the target [BTF](Btf) info.
+    ///
+    /// The callback is called only when the object contains CO-RE relocations or
+    /// typed kernel symbols. It receives a parser configured for the object's
+    /// endianness, allowing BTF to be parsed from a borrowed byte slice. Once
+    /// parsed, the target `BTF` is cached and reused by this loader. To share one
+    /// parsed value across multiple loaders, use [`EbpfLoader::btf`] instead.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// use std::fs;
+    ///
+    /// use aya::{EbpfError, EbpfLoader};
+    ///
+    /// let bpf = EbpfLoader::new()
+    ///     .btf_source(|parse| {
+    ///         let data =
+    ///             fs::read("/custom_btf_file").map_err(EbpfError::BtfSourceError)?;
+    ///         Ok(parse(&data)?)
+    ///     })
+    ///     .load_file("file.o")?;
+    ///
+    /// # Ok::<(), aya::EbpfError>(())
+    /// ```
+    pub fn btf_source<F>(&mut self, source: F) -> &mut Self
+    where
+        F: Fn(&BtfParser) -> Result<Btf, EbpfError>
+            + Send
+            + Sync
+            + std::panic::RefUnwindSafe
+            + std::panic::UnwindSafe
+            + 'static,
+    {
+        self.btf = TargetBtf::Source(BtfSource(Box::new(source)));
         self
     }
 
@@ -493,30 +563,48 @@ impl<'a> EbpfLoader<'a> {
             None
         };
 
-        if btf.is_none() && (obj.has_btf_relocations() || obj.has_typed_ksyms()) {
-            match Btf::from_sys_fs() {
-                Ok(kernel_btf) => *btf = Some(Cow::Owned(kernel_btf)),
+        if obj.has_btf_relocations() || obj.has_typed_ksyms() {
+            let endianness = obj.endianness;
+            let target_btf: Result<Cow<'_, Btf>, EbpfError> = match btf {
+                TargetBtf::System => Btf::from_sys_fs().map(Cow::Owned).map_err(EbpfError::from),
+                TargetBtf::Parsed(btf) => Ok(Cow::Borrowed(btf)),
+                TargetBtf::Source(BtfSource(source)) => {
+                    source(&move |data| Btf::parse(data, endianness)).map(Cow::Owned)
+                }
+            };
+
+            match target_btf {
+                Ok(target_btf) => {
+                    let result = obj
+                        .relocate_btf(&target_btf)
+                        .map_err(EbpfError::from)
+                        .and_then(|()| {
+                            obj.resolve_externs(Some(&target_btf))
+                                .map_err(EbpfError::from)
+                        });
+
+                    if let Cow::Owned(target_btf) = target_btf {
+                        *btf = TargetBtf::Parsed(Cow::Owned(target_btf));
+                    }
+
+                    result?;
+                }
                 Err(err) => {
-                    // CO-RE relocations and strong typed ksyms cannot be resolved without kernel
+                    // CO-RE relocations and strong typed ksyms cannot be resolved without target
                     // BTF, so preserve the original loading error.
                     if obj.has_btf_relocations() || obj.has_strong_typed_ksyms() {
-                        return Err(err.into());
+                        return Err(err);
                     }
 
-                    // Unresolved weak typed ksyms are allowed and handled during extern
-                    // relocation.
-                    warn!("kernel BTF is unavailable; weak typed ksyms will be unresolved: {err}")
+                    // Unresolved weak typed ksyms are allowed and handled during extern relocation.
+                    warn!("target BTF is unavailable; weak typed ksyms will be unresolved: {err}");
+                    obj.resolve_externs(None)?;
                 }
             }
+        } else {
+            obj.resolve_externs(None)?;
         }
 
-        let btf = btf.as_deref();
-        if let Some(btf) = btf {
-            obj.relocate_btf(btf)?;
-        }
-
-        obj.resolve_externs(btf)?;
-
         const fn is_map_of_maps(map_type: bpf_map_type) -> bool {
             matches!(
                 map_type,
@@ -928,6 +1016,8 @@ const fn adjust_to_page_size(byte_size: u32, page_size: u32) -> u32 {
 mod tests {
     use aya_obj::generated::bpf_map_type::*;
 
+    use super::{Btf, BtfSource, EbpfLoader, TargetBtf};
+
     const PAGE_SIZE: u32 = 4096;
     const NUM_CPUS: u32 = 4;
 
@@ -975,6 +1065,26 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn test_btf_source_can_parse_borrowed_subslice() {
+        let mut loader = EbpfLoader::new();
+        loader.btf_source(|parse| {
+            // This models an mmap opened by the callback: its backing storage
+            // can be dropped once the borrowed subslice has been parsed.
+            let data = [0, 1, 2, 3];
+            Ok(parse(&data[1..3])?)
+        });
+
+        let TargetBtf::Source(BtfSource(source)) = &loader.btf else {
+            panic!("expected a BTF source");
+        };
+        source(&|data| {
+            assert_eq!(data, [1, 2]);
+            Ok(Btf::new())
+        })
+        .unwrap();
+    }
 }
 
 impl Default for EbpfLoader<'_> {
@@ -1244,6 +1354,10 @@ pub enum EbpfError {
         error: io::Error,
     },
 
+    /// Error reading target BTF
+    #[error("error reading BTF source")]
+    BtfSourceError(#[source] io::Error),
+
     /// Unexpected pinning type
     #[error("unexpected pinning type {name}")]
     UnexpectedPinningType {

test/integration-test/src/tests/btf_relocations.rs

@@ -1,5 +1,10 @@
+use std::sync::{
+    Arc,
+    atomic::{AtomicUsize, Ordering},
+};
+
 use aya::{
-    EbpfLoader, Endianness,
+    Ebpf, EbpfLoader, Endianness,
     maps::Array,
     programs::{UProbe, uprobe::UProbeScope},
 };
@@ -54,8 +59,46 @@ use rstest::rstest;
 fn relocation_tests(#[case] bpf: &[u8], #[case] btf: &[u8], #[case] expected: u64) {
     let btf = Btf::parse(btf, Endianness::default()).unwrap();
 
-    let mut bpf = EbpfLoader::new().btf(&btf).load(bpf).unwrap();
+    let bpf = EbpfLoader::new().btf(&btf).load(bpf).unwrap();
+
+    assert_relocation(bpf, expected);
+}
+
+#[test_log::test]
+fn lazy_btf_source_is_cached() {
+    let calls = Arc::new(AtomicUsize::new(0));
+    let source_calls = Arc::clone(&calls);
+    let mut loader = EbpfLoader::new();
+    loader.btf_source(move |parse| {
+        source_calls.fetch_add(1, Ordering::Relaxed);
+        Ok(parse(crate::FIELD_RELOC_BTF)?)
+    });
+
+    let bpf = loader.load(crate::FIELD_RELOC_BPF).unwrap();
+    let _second_bpf = loader.load(crate::FIELD_RELOC_BPF).unwrap();
+
+    assert_eq!(calls.load(Ordering::Relaxed), 1);
+    assert_relocation(bpf, 1);
+}
+
+#[test_log::test]
+fn btf_source_is_not_called_when_unneeded() {
+    let calls = Arc::new(AtomicUsize::new(0));
+    let source_calls = Arc::clone(&calls);
+    let mut loader = EbpfLoader::new();
+    loader.btf_source(move |parse| {
+        source_calls.fetch_add(1, Ordering::Relaxed);
+        Ok(parse(crate::FIELD_RELOC_BTF)?)
+    });
+
+    // RINGBUF_BTF contains BTF map definitions but no CO-RE relocations or typed ksyms. If that
+    // changes, replace it with an object that does not require target BTF.
+    let _bpf = loader.load(crate::RINGBUF_BTF).unwrap();
+
+    assert_eq!(calls.load(Ordering::Relaxed), 0);
+}
 
+fn assert_relocation(mut bpf: Ebpf, expected: u64) {
     let program: &mut UProbe = bpf.program_mut("program").unwrap().try_into().unwrap();
     program.load().unwrap();
     program

xtask/public-api/aya.txt

@@ -7840,6 +7840,7 @@ pub macro aya::include_bytes_aligned!
 pub enum aya::EbpfError
 pub aya::EbpfError::BtfError(aya_obj::btf::btf::BtfError)
 pub aya::EbpfError::BtfRelocationError(aya_obj::btf::relocation::BtfRelocationError)
+pub aya::EbpfError::BtfSourceError(core::io::error::Error)
 pub aya::EbpfError::FileError
 pub aya::EbpfError::FileError::error: core::io::error::Error
 pub aya::EbpfError::FileError::path: std::path::PathBuf
@@ -7905,6 +7906,7 @@ pub struct aya::EbpfLoader<'a>
 impl<'a> aya::EbpfLoader<'a>
 pub const fn aya::EbpfLoader<'a>::allow_unsupported_maps(&mut self) -> &mut Self
 pub fn aya::EbpfLoader<'a>::btf(&mut self, &'a aya_obj::btf::btf::Btf) -> &mut Self
+pub fn aya::EbpfLoader<'a>::btf_source<F>(&mut self, F) -> &mut Self where F: core::ops::function::Fn(&dyn core::ops::function::Fn(&[u8]) -> core::result::Result<aya_obj::btf::btf::Btf, aya_obj::btf::btf::BtfError>) -> core::result::Result<aya_obj::btf::btf::Btf, aya::EbpfError> + core::marker::Send + core::marker::Sync + core::panic::unwind_safe::RefUnwindSafe + core::panic::unwind_safe::UnwindSafe + 'static
 pub fn aya::EbpfLoader<'a>::default_map_pin_directory<P: core::convert::AsRef<std::path::Path>>(&mut self, P) -> &mut Self
 pub fn aya::EbpfLoader<'a>::extension(&mut self, &'a str) -> &mut Self
 pub fn aya::EbpfLoader<'a>::load(&mut self, &[u8]) -> core::result::Result<aya::Ebpf, aya::EbpfError>