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 a reader factory that is invoked only for CO-RE relocations or typed
kernel symbols, and cache the BTF after the first successful parse. Keep
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,38 @@ pub struct EbpfLoader<'a> {
     allow_unsupported_maps: bool,
 }
 
+#[derive(Debug)]
+enum TargetBtf<'a> {
+    System,
+    Parsed(Cow<'a, Btf>),
+    Source(BtfSource),
+}
+
+struct BtfSource {
+    read: Box<
+        dyn Fn() -> io::Result<Vec<u8>>
+            // Preserve `EbpfLoader`'s existing auto traits.
+            + Send
+            + Sync
+            + std::panic::RefUnwindSafe
+            + std::panic::UnwindSafe
+            // Keep captured state independent of `EbpfLoader`'s borrowed configuration.
+            + 'static,
+    >,
+}
+
+impl std::fmt::Debug for BtfSource {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("BtfSource").finish_non_exhaustive()
+    }
+}
+
+impl BtfSource {
+    fn read(&self) -> io::Result<Vec<u8>> {
+        (self.read)()
+    }
+}
+
 /// 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 +189,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 +206,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 +226,48 @@ 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 factory for lazily reading the target [BTF](Btf) info.
+    ///
+    /// The factory is called only when the object contains CO-RE relocations or
+    /// typed kernel symbols. 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::File;
+    ///
+    /// use aya::EbpfLoader;
+    ///
+    /// let bpf = EbpfLoader::new()
+    ///     .btf_source(|| File::open("/custom_btf_file"))
+    ///     .load_file("file.o")?;
+    ///
+    /// # Ok::<(), aya::EbpfError>(())
+    /// ```
+    pub fn btf_source<F, R>(&mut self, source: F) -> &mut Self
+    where
+        F: Fn() -> io::Result<R>
+            + Send
+            + Sync
+            + std::panic::RefUnwindSafe
+            + std::panic::UnwindSafe
+            + 'static,
+        R: io::Read,
+    {
+        self.btf = TargetBtf::Source(BtfSource {
+            read: Box::new(move || {
+                let mut reader = source()?;
+                let mut data = Vec::new();
+                reader.read_to_end(&mut data)?;
+                Ok(data)
+            }),
+        });
         self
     }
 
@@ -493,29 +570,51 @@ 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)),
-                Err(err) => {
-                    // CO-RE relocations and strong typed ksyms cannot be resolved without kernel
-                    // BTF, so preserve the original loading error.
-                    if obj.has_btf_relocations() || obj.has_strong_typed_ksyms() {
-                        return Err(err.into());
-                    }
+        let target_btf: Option<&Btf> = if obj.has_btf_relocations() || obj.has_typed_ksyms() {
+            let load_result = match btf {
+                TargetBtf::System => Some(Btf::from_sys_fs().map_err(EbpfError::from)),
+                TargetBtf::Source(source) => Some(
+                    source
+                        .read()
+                        .map_err(EbpfError::BtfSourceError)
+                        .and_then(|data| {
+                            Btf::parse(&data, obj.endianness).map_err(EbpfError::BtfError)
+                        }),
+                ),
+                TargetBtf::Parsed(_) => None,
+            };
+            if let Some(load_result) = load_result {
+                match load_result {
+                    Ok(target_btf) => *btf = TargetBtf::Parsed(Cow::Owned(target_btf)),
+                    Err(err) => {
+                        // 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);
+                        }
 
-                    // 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}"
+                        )
+                    }
                 }
             }
-        }
 
-        let btf = btf.as_deref();
-        if let Some(btf) = btf {
-            obj.relocate_btf(btf)?;
+            match btf {
+                TargetBtf::Parsed(btf) => Some(btf),
+                _ => None,
+            }
+        } else {
+            None
+        };
+
+        if let Some(target_btf) = target_btf {
+            obj.relocate_btf(target_btf)?;
         }
 
-        obj.resolve_externs(btf)?;
+        obj.resolve_externs(target_btf)?;
 
         const fn is_map_of_maps(map_type: bpf_map_type) -> bool {
             matches!(
@@ -1244,6 +1343,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,13 @@
+use std::{
+    io::Cursor,
+    sync::{
+        Arc,
+        atomic::{AtomicUsize, Ordering},
+    },
+};
+
 use aya::{
-    EbpfLoader, Endianness,
+    Ebpf, EbpfLoader, Endianness,
     maps::Array,
     programs::{UProbe, uprobe::UProbeScope},
 };
@@ -54,8 +62,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 || {
+        source_calls.fetch_add(1, Ordering::Relaxed);
+        Ok(Cursor::new(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_read_when_unneeded() {
+    let calls = Arc::new(AtomicUsize::new(0));
+    let source_calls = Arc::clone(&calls);
+    let mut loader = EbpfLoader::new();
+    loader.btf_source(move || {
+        source_calls.fetch_add(1, Ordering::Relaxed);
+        Ok(Cursor::new(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, R>(&mut self, F) -> &mut Self where F: core::ops::function::Fn() -> core::io::error::Result<R> + core::marker::Send + core::marker::Sync + core::panic::unwind_safe::RefUnwindSafe + core::panic::unwind_safe::UnwindSafe + 'static, R: alloc::io::read::Read
 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>