Add support for kernel module split BTF

The kernel splits kernel modules into dedicated BTF files, using "split
BTF mode". In split BTF mode, each split BTF may reference the "base
BTF" (AKA the vmlinux BTF), but may not reference each-other. The way
it works is fairly simple: you can combine a base BTF and split BTF
by concatenating their strings and types together, creating a working
merged BTF.

Where things get dicey is when merging a base BTF with multiple split
BTFs. Because every split BTF will start from the same offset, their
string offsets and type IDs will need to be "rebased" to a new location
when merging it into the output BTF.

This commit adds support for merging one base BTF and multiple split BTF
into a single merged BTF, allowing eBPFs to manipulate structs that come
from kernel modules. It also reworks Btf::from_sys_fs to make use of
this capacity.

aya-obj/src/btf/btf.rs

@@ -14,7 +14,7 @@ use object::{Endianness, SectionIndex};
 use crate::{
     Object,
     btf::{
-        Array, BtfEnum, BtfKind, BtfMember, BtfType, Const, Enum, FuncInfo, FuncLinkage, Int,
+        Array, BtfEnum, BtfKind, BtfMember, BtfRebaseInfo, BtfType, Const, Enum, FuncInfo, FuncLinkage, Int,
         IntEncoding, LineInfo, Struct, Typedef, Union, VarLinkage,
         info::{FuncSecInfo, LineSecInfo},
         relocation::Relocation,
@@ -268,6 +268,34 @@ impl Btf {
         }
     }
 
+    /// Merges a base BTF and multiple split BTFs into a single BTF.
+    pub fn merge_split_btfs(base_btf: &Btf, split_btfs: &[Btf]) -> Btf {
+        // Create a Btf from all the btfs.
+        let mut out_btf = base_btf.clone();
+
+        let mut rebase_info = BtfRebaseInfo {
+            str_rebase_from: out_btf.strings.len() as u32,
+            types_rebase_from: out_btf.types.types.len() as u32,
+            str_new_offset: out_btf.strings.len() as u32,
+            types_new_offset: out_btf.types.types.len() as u32,
+        };
+        for btf in split_btfs {
+            rebase_info.str_new_offset = out_btf.strings.len() as u32;
+            rebase_info.types_new_offset = out_btf.types.types.len() as u32;
+            out_btf.strings.extend(&btf.strings);
+            out_btf.header.str_len = out_btf.strings.len() as u32;
+            // Skip over the Unknown type at offset 0, as this only exists in
+            // the "base" BTF, and not the "split" BTFs.
+            for ty in btf.types.types.iter().skip(1) {
+                out_btf.types.types.push(ty.rebase(&rebase_info));
+            }
+            out_btf.header.type_len += btf.header.type_len;
+            out_btf.header.str_off += btf.header.type_len;
+        }
+
+        out_btf
+    }
+
     pub(crate) fn is_empty(&self) -> bool {
         // the first one is awlays BtfType::Unknown
         self.types.types.len() < 2
@@ -296,10 +324,36 @@ impl Btf {
         type_id as u32
     }
 
-    /// Loads BTF metadata from `/sys/kernel/btf/vmlinux`.
+    /// Loads BTF metadata from `/sys/kernel/btf`.
     #[cfg(feature = "std")]
     pub fn from_sys_fs() -> Result<Btf, BtfError> {
-        Btf::parse_file("/sys/kernel/btf/vmlinux", Endianness::default())
+        let base_btf = Btf::parse_file("/sys/kernel/btf/vmlinux", Endianness::default())?;
+        let mut split_btfs = vec![];
+        let dir_iter = std::fs::read_dir("/sys/kernel/btf")
+            .ok()
+            .into_iter()
+            .flatten()
+            .filter_map(|v| v.ok())
+            .filter(|v| v.file_name() != "vmlinux");
+        for entry in dir_iter {
+            match entry.file_type() {
+                Ok(v) if !v.is_file() => continue,
+                Err(_err) => continue,
+                Ok(_v) => (),
+            }
+            match Btf::parse_file(entry.path(), Endianness::default()) {
+                Ok(v) => split_btfs.push(v),
+                // Ignore errors - the goal is to enhance the base BTF with as
+                // many split BTFs we can.
+                Err(_err) => (),
+            }
+        }
+
+        if split_btfs.len() > 0 {
+            Ok(Self::merge_split_btfs(&base_btf, &split_btfs))
+        } else {
+            Ok(base_btf)
+        }
     }
 
     /// Loads BTF metadata from the given `path`.

aya-obj/src/btf/types.rs

@@ -7,6 +7,42 @@ use object::Endianness;
 
 use crate::btf::{Btf, BtfError, MAX_RESOLVE_DEPTH};
 
+#[derive(Debug)]
+pub(crate) struct BtfRebaseInfo {
+    /// Index of the first string offset to allow rebasing from.
+    ///
+    /// Offsets below this value are considered to be part of the "base BTF",
+    /// and as such should not be relocated.
+    pub(crate) str_rebase_from: u32,
+    /// Index of the first type to allow rebasing from.
+    ///
+    /// Indices below this value are considered to be part of the "base BTF",
+    /// and as such should not be relocated.
+    pub(crate) types_rebase_from: u32,
+    /// The new starting offset for strings.
+    pub(crate) str_new_offset: u32,
+    /// The new starting index for types.
+    pub(crate) types_new_offset: u32,
+}
+
+impl BtfRebaseInfo {
+    fn rebase_str(&self, str_offset: u32) -> u32 {
+        if str_offset < self.str_rebase_from {
+            str_offset
+        } else {
+            str_offset - self.str_rebase_from + self.str_new_offset
+        }
+    }
+
+    fn rebase_type(&self, type_offset: u32) -> u32 {
+        if type_offset < self.types_rebase_from {
+            type_offset
+        } else {
+            type_offset - self.types_rebase_from + self.types_new_offset
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 pub enum BtfType {
     Unknown,
@@ -51,6 +87,14 @@ impl Fwd {
     pub(crate) fn type_info_size(&self) -> usize {
         mem::size_of::<Self>()
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Fwd {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            _unused: self._unused,
+        }
+    }
 }
 
 #[repr(C)]
@@ -82,6 +126,14 @@ impl Const {
             btf_type,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Const {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            btf_type: rebase_info.rebase_type(self.btf_type),
+        }
+    }
 }
 
 #[repr(C)]
@@ -104,6 +156,14 @@ impl Volatile {
     pub(crate) fn type_info_size(&self) -> usize {
         mem::size_of::<Self>()
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Volatile {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            btf_type: rebase_info.rebase_type(self.btf_type),
+        }
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -125,6 +185,14 @@ impl Restrict {
     pub(crate) fn type_info_size(&self) -> usize {
         mem::size_of::<Self>()
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Restrict {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            _info: self._info,
+            btf_type: rebase_info.rebase_type(self.btf_type),
+        }
+    }
 }
 
 #[repr(C)]
@@ -156,6 +224,14 @@ impl Ptr {
             btf_type,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Ptr {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            btf_type: rebase_info.rebase_type(self.btf_type),
+        }
+    }
 }
 
 #[repr(C)]
@@ -187,6 +263,14 @@ impl Typedef {
             btf_type,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Typedef {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            btf_type: rebase_info.rebase_type(self.btf_type),
+        }
+    }
 }
 
 #[repr(C)]
@@ -217,6 +301,14 @@ impl Float {
             size,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Float {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            size: self.size,
+        }
+    }
 }
 
 #[repr(C)]
@@ -276,6 +368,14 @@ impl Func {
     pub(crate) fn set_linkage(&mut self, linkage: FuncLinkage) {
         self.info = (self.info & 0xFFFF0000) | (linkage as u32) & 0xFFFF;
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Func {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            btf_type: rebase_info.rebase_type(self.btf_type),
+        }
+    }
 }
 
 #[repr(C)]
@@ -307,6 +407,14 @@ impl TypeTag {
             btf_type,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        TypeTag {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            btf_type: rebase_info.rebase_type(self.btf_type),
+        }
+    }
 }
 
 #[repr(u32)]
@@ -391,6 +499,15 @@ impl Int {
     pub(crate) fn bits(&self) -> u32 {
         self.data & 0x000000ff
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Int {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            size: self.size,
+            data: self.data,
+        }
+    }
 }
 
 #[repr(C)]
@@ -404,6 +521,13 @@ impl BtfEnum {
     pub fn new(name_offset: u32, value: u32) -> Self {
         Self { name_offset, value }
     }
+
+    fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        BtfEnum {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            value: self.value,
+        }
+    }
 }
 
 #[repr(C)]
@@ -470,6 +594,15 @@ impl Enum {
             self.info &= !(1 << 31);
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Enum {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            size: self.size,
+            variants: self.variants.iter().map(|v| v.rebase(rebase_info)).collect(),
+        }
+    }
 }
 
 #[repr(C)]
@@ -488,6 +621,14 @@ impl BtfEnum64 {
             value_high: (value >> 32) as u32,
         }
     }
+
+    fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        BtfEnum64 {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            value_low: self.value_low,
+            value_high: self.value_high,
+        }
+    }
 }
 
 #[repr(C)]
@@ -562,6 +703,15 @@ impl Enum64 {
             variants,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Enum64 {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            size: self.size,
+            variants: self.variants.iter().map(|v| v.rebase(rebase_info)).collect(),
+        }
+    }
 }
 
 #[repr(C)]
@@ -572,6 +722,16 @@ pub(crate) struct BtfMember {
     pub(crate) offset: u32,
 }
 
+impl BtfMember {
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        BtfMember {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            btf_type: rebase_info.rebase_type(self.btf_type),
+            offset: self.offset,
+        }
+    }
+}
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct Struct {
@@ -649,6 +809,15 @@ impl Struct {
 
         size as usize
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Struct {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            size: self.size,
+            members: self.members.iter().map(|v| v.rebase(rebase_info)).collect(),
+        }
+    }
 }
 
 #[repr(C)]
@@ -728,6 +897,15 @@ impl Union {
 
         size as usize
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Union {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            size: self.size,
+            members: self.members.iter().map(|v| v.rebase(rebase_info)).collect(),
+        }
+    }
 }
 
 #[repr(C)]
@@ -786,6 +964,19 @@ impl Array {
             },
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Array {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            _unused: self._unused,
+            array: BtfArray {
+                element_type: rebase_info.rebase_type(self.array.element_type),
+                index_type: rebase_info.rebase_type(self.array.index_type),
+                len: self.array.len,
+            },
+        }
+    }
 }
 
 #[repr(C)]
@@ -795,6 +986,15 @@ pub struct BtfParam {
     pub btf_type: u32,
 }
 
+impl BtfParam {
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        BtfParam {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            btf_type: rebase_info.rebase_type(self.btf_type),
+        }
+    }
+}
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct FuncProto {
@@ -847,6 +1047,15 @@ impl FuncProto {
             params,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        FuncProto {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            return_type: rebase_info.rebase_type(self.return_type),
+            params: self.params.iter().map(|v| v.rebase(rebase_info)).collect(),
+        }
+    }
 }
 
 #[repr(u32)]
@@ -912,6 +1121,15 @@ impl Var {
             linkage,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        Var {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            btf_type: rebase_info.rebase_type(self.btf_type),
+            linkage: self.linkage.clone(),
+        }
+    }
 }
 
 #[repr(C)]
@@ -922,6 +1140,16 @@ pub struct DataSecEntry {
     pub size: u32,
 }
 
+impl DataSecEntry {
+    fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        DataSecEntry {
+            btf_type: rebase_info.rebase_type(self.btf_type),
+            offset: self.offset,
+            size: self.size,
+        }
+    }
+}
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct DataSec {
@@ -981,6 +1209,15 @@ impl DataSec {
             entries,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        DataSec {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            size: self.size,
+            entries: self.entries.iter().map(|v| v.rebase(rebase_info)).collect(),
+        }
+    }
 }
 
 #[repr(C)]
@@ -1026,6 +1263,15 @@ impl DeclTag {
             component_index,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        DeclTag {
+            name_offset: rebase_info.rebase_str(self.name_offset),
+            info: self.info,
+            btf_type: rebase_info.rebase_type(self.btf_type),
+            component_index: self.component_index,
+        }
+    }
 }
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
@@ -1425,6 +1671,31 @@ impl BtfType {
             (BtfKind::Enum, BtfKind::Enum64) | (BtfKind::Enum64, BtfKind::Enum)
         )
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        match self {
+            BtfType::Unknown => BtfType::Unknown,
+            BtfType::Fwd(t) => BtfType::Fwd(t.rebase(rebase_info)),
+            BtfType::Const(t) => BtfType::Const(t.rebase(rebase_info)),
+            BtfType::Volatile(t) => BtfType::Volatile(t.rebase(rebase_info)),
+            BtfType::Restrict(t) => BtfType::Restrict(t.rebase(rebase_info)),
+            BtfType::Ptr(t) => BtfType::Ptr(t.rebase(rebase_info)),
+            BtfType::Typedef(t) => BtfType::Typedef(t.rebase(rebase_info)),
+            BtfType::Func(t) => BtfType::Func(t.rebase(rebase_info)),
+            BtfType::Int(t) => BtfType::Int(t.rebase(rebase_info)),
+            BtfType::Float(t) => BtfType::Float(t.rebase(rebase_info)),
+            BtfType::Enum(t) => BtfType::Enum(t.rebase(rebase_info)),
+            BtfType::Enum64(t) => BtfType::Enum64(t.rebase(rebase_info)),
+            BtfType::Array(t) => BtfType::Array(t.rebase(rebase_info)),
+            BtfType::Struct(t) => BtfType::Struct(t.rebase(rebase_info)),
+            BtfType::Union(t) => BtfType::Union(t.rebase(rebase_info)),
+            BtfType::FuncProto(t) => BtfType::FuncProto(t.rebase(rebase_info)),
+            BtfType::Var(t) => BtfType::Var(t.rebase(rebase_info)),
+            BtfType::DataSec(t) => BtfType::DataSec(t.rebase(rebase_info)),
+            BtfType::DeclTag(t) => BtfType::DeclTag(t.rebase(rebase_info)),
+            BtfType::TypeTag(t) => BtfType::TypeTag(t.rebase(rebase_info)),
+        }
+    }
 }
 
 fn type_kind(info: u32) -> Result<BtfKind, BtfError> {