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,8 +14,9 @@ use object::{Endianness, SectionIndex};
 use crate::{
     Object,
     btf::{
-        Array, BtfEnum, BtfKind, BtfMember, BtfType, Const, Enum, FuncInfo, FuncLinkage, Int,
+        Array, BtfEnum, BtfKind, BtfMember, BtfRebaseInfo, BtfRebaseInfoField, BtfType, Const,
-        IntEncoding, LineInfo, Struct, Typedef, Union, VarLinkage,
+        Enum, FuncInfo, FuncLinkage, Int, IntEncoding, LineInfo, Struct, Typedef, Union,
+        VarLinkage,
         info::{FuncSecInfo, LineSecInfo},
         relocation::Relocation,
     },
@@ -268,6 +269,50 @@ impl Btf {
         }
     }
 
+    /// Merges a base BTF and multiple split BTFs into a single BTF.
+    pub fn merge_split_btfs(mut base_btf: Btf, split_btfs: &[Btf]) -> Btf {
+        let strings_rebase_from = base_btf.strings.len() as u32;
+        let types_rebase_from = base_btf.types.types.len() as u32;
+        for split_btf in split_btfs {
+            let Btf {
+                header: split_btf_header,
+                strings: split_btf_strings,
+                types: split_btf_types,
+                _endianness,
+            } = split_btf;
+
+            let rebase_info = BtfRebaseInfo {
+                strings: BtfRebaseInfoField {
+                    rebase_from: strings_rebase_from,
+                    new_offset: base_btf.strings.len() as u32,
+                },
+                types: BtfRebaseInfoField {
+                    rebase_from: types_rebase_from,
+                    new_offset: base_btf.types.types.len() as u32,
+                },
+            };
+
+            // Append the strings from the split BTF. We concatenate it with the
+            // existing strings, and will rebase the types to take the new
+            // offsets into account.
+            base_btf.strings.extend(split_btf_strings);
+
+            // Skip over the Unknown type at offset 0, as this only exists in
+            // the "base" BTF, and not the "split" BTFs. Append the rest of the
+            // types from the split BTF.
+            for ty in split_btf_types.types.iter().skip(1) {
+                base_btf.types.types.push(ty.rebase(&rebase_info));
+            }
+
+            // Update the header.
+            base_btf.header.str_len = base_btf.strings.len() as u32;
+            base_btf.header.type_len += split_btf_header.type_len;
+            base_btf.header.str_off += split_btf_header.type_len;
+        }
+
+        base_btf
+    }
+
     pub(crate) fn is_empty(&self) -> bool {
         // the first one is awlays BtfType::Unknown
         self.types.types.len() < 2
@@ -296,10 +341,34 @@ 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 sys_kernel_btf_path = "/sys/kernel/btf";
+        let dir_iter =
+            std::fs::read_dir(sys_kernel_btf_path).map_err(|error| BtfError::FileError {
+                path: sys_kernel_btf_path.into(),
+                error,
+            })?;
+        for entry in dir_iter {
+            let entry = entry.map_err(|error| BtfError::FileError {
+                path: sys_kernel_btf_path.into(),
+                error,
+            })?;
+            let file_type = entry.file_type().map_err(|error| BtfError::FileError {
+                path: entry.path(),
+                error,
+            })?;
+
+            if !file_type.is_file() {
+                continue;
+            }
+            split_btfs.push(Btf::parse_file(entry.path(), Endianness::default())?);
+        }
+
+        Ok(Self::merge_split_btfs(base_btf, &split_btfs))
     }
 
     /// Loads BTF metadata from the given `path`.

aya-obj/src/btf/types.rs

@@ -7,6 +7,45 @@ use object::Endianness;
 
 use crate::btf::{Btf, BtfError, MAX_RESOLVE_DEPTH};
 
+#[derive(Debug)]
+pub(crate) struct BtfRebaseInfoField {
+    /// Index of the first 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 rebase_from: u32,
+    /// The new starting offset.
+    pub new_offset: u32,
+}
+
+impl BtfRebaseInfoField {
+    fn rebase(&self, offset: u32) -> u32 {
+        let &Self {
+            rebase_from,
+            new_offset,
+        } = self;
+        match offset.checked_sub(rebase_from) {
+            None => offset,
+            Some(offset) => new_offset + offset,
+        }
+    }
+}
+
+pub(crate) struct BtfRebaseInfo {
+    pub strings: BtfRebaseInfoField,
+    pub types: BtfRebaseInfoField,
+}
+
+impl BtfRebaseInfo {
+    fn rebase_str(&self, str_offset: u32) -> u32 {
+        self.strings.rebase(str_offset)
+    }
+
+    fn rebase_type(&self, type_offset: u32) -> u32 {
+        self.types.rebase(type_offset)
+    }
+}
+
 #[derive(Clone, Debug)]
 pub enum BtfType {
     Unknown,
@@ -51,6 +90,19 @@ impl Fwd {
     pub(crate) fn type_info_size(&self) -> usize {
         mem::size_of::<Self>()
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            _unused,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            _unused,
+        }
+    }
 }
 
 #[repr(C)]
@@ -82,6 +134,19 @@ impl Const {
             btf_type,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            btf_type,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            btf_type: rebase_info.rebase_type(btf_type),
+        }
+    }
 }
 
 #[repr(C)]
@@ -104,6 +169,19 @@ impl Volatile {
     pub(crate) fn type_info_size(&self) -> usize {
         mem::size_of::<Self>()
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            btf_type,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            btf_type: rebase_info.rebase_type(btf_type),
+        }
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -125,6 +203,19 @@ impl Restrict {
     pub(crate) fn type_info_size(&self) -> usize {
         mem::size_of::<Self>()
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            _info,
+            btf_type,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            _info,
+            btf_type: rebase_info.rebase_type(btf_type),
+        }
+    }
 }
 
 #[repr(C)]
@@ -156,6 +247,19 @@ impl Ptr {
             btf_type,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            btf_type,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            btf_type: rebase_info.rebase_type(btf_type),
+        }
+    }
 }
 
 #[repr(C)]
@@ -187,6 +291,19 @@ impl Typedef {
             btf_type,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            btf_type,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            btf_type: rebase_info.rebase_type(btf_type),
+        }
+    }
 }
 
 #[repr(C)]
@@ -217,6 +334,19 @@ impl Float {
             size,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            size,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            size,
+        }
+    }
 }
 
 #[repr(C)]
@@ -276,6 +406,19 @@ 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 {
+        let &Self {
+            name_offset,
+            info,
+            btf_type,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            btf_type: rebase_info.rebase_type(btf_type),
+        }
+    }
 }
 
 #[repr(C)]
@@ -307,6 +450,19 @@ impl TypeTag {
             btf_type,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            btf_type,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            btf_type: rebase_info.rebase_type(btf_type),
+        }
+    }
 }
 
 #[repr(u32)]
@@ -391,6 +547,21 @@ impl Int {
     pub(crate) fn bits(&self) -> u32 {
         self.data & 0x000000ff
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            size,
+            data,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            size,
+            data,
+        }
+    }
 }
 
 #[repr(C)]
@@ -404,6 +575,14 @@ impl BtfEnum {
     pub fn new(name_offset: u32, value: u32) -> Self {
         Self { name_offset, value }
     }
+
+    fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self { name_offset, value } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            value,
+        }
+    }
 }
 
 #[repr(C)]
@@ -470,6 +649,21 @@ impl Enum {
             self.info &= !(1 << 31);
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            size,
+            ref variants,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            size,
+            variants: variants.iter().map(|v| v.rebase(rebase_info)).collect(),
+        }
+    }
 }
 
 #[repr(C)]
@@ -488,6 +682,19 @@ impl BtfEnum64 {
             value_high: (value >> 32) as u32,
         }
     }
+
+    fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            value_low,
+            value_high,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            value_low,
+            value_high,
+        }
+    }
 }
 
 #[repr(C)]
@@ -562,6 +769,21 @@ impl Enum64 {
             variants,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            size,
+            ref variants,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            size,
+            variants: variants.iter().map(|v| v.rebase(rebase_info)).collect(),
+        }
+    }
 }
 
 #[repr(C)]
@@ -572,6 +794,21 @@ pub(crate) struct BtfMember {
     pub(crate) offset: u32,
 }
 
+impl BtfMember {
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            btf_type,
+            offset,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            btf_type: rebase_info.rebase_type(btf_type),
+            offset,
+        }
+    }
+}
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct Struct {
@@ -649,6 +886,21 @@ impl Struct {
 
         size as usize
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            size,
+            ref members,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            size,
+            members: members.iter().map(|v| v.rebase(rebase_info)).collect(),
+        }
+    }
 }
 
 #[repr(C)]
@@ -728,6 +980,21 @@ impl Union {
 
         size as usize
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            size,
+            ref members,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            size,
+            members: members.iter().map(|v| v.rebase(rebase_info)).collect(),
+        }
+    }
 }
 
 #[repr(C)]
@@ -738,6 +1005,21 @@ pub(crate) struct BtfArray {
     pub(crate) len: u32,
 }
 
+impl BtfArray {
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            element_type,
+            index_type,
+            len,
+        } = self;
+        Self {
+            element_type: rebase_info.rebase_type(element_type),
+            index_type: rebase_info.rebase_type(index_type),
+            len,
+        }
+    }
+}
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct Array {
@@ -786,6 +1068,21 @@ impl Array {
             },
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            _unused,
+            ref array,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            _unused,
+            array: array.rebase(rebase_info),
+        }
+    }
 }
 
 #[repr(C)]
@@ -795,6 +1092,19 @@ pub struct BtfParam {
     pub btf_type: u32,
 }
 
+impl BtfParam {
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            btf_type,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            btf_type: rebase_info.rebase_type(btf_type),
+        }
+    }
+}
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct FuncProto {
@@ -847,6 +1157,21 @@ impl FuncProto {
             params,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            return_type,
+            ref params,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            return_type: rebase_info.rebase_type(return_type),
+            params: params.iter().map(|v| v.rebase(rebase_info)).collect(),
+        }
+    }
 }
 
 #[repr(u32)]
@@ -912,6 +1237,21 @@ impl Var {
             linkage,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            btf_type,
+            ref linkage,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            btf_type: rebase_info.rebase_type(btf_type),
+            linkage: linkage.clone(),
+        }
+    }
 }
 
 #[repr(C)]
@@ -922,6 +1262,21 @@ pub struct DataSecEntry {
     pub size: u32,
 }
 
+impl DataSecEntry {
+    fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            btf_type,
+            offset,
+            size,
+        } = self;
+        Self {
+            btf_type: rebase_info.rebase_type(btf_type),
+            offset,
+            size,
+        }
+    }
+}
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct DataSec {
@@ -981,6 +1336,21 @@ impl DataSec {
             entries,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            size,
+            ref entries,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            size,
+            entries: entries.iter().map(|v| v.rebase(rebase_info)).collect(),
+        }
+    }
 }
 
 #[repr(C)]
@@ -1026,6 +1396,21 @@ impl DeclTag {
             component_index,
         }
     }
+
+    pub(crate) fn rebase(&self, rebase_info: &BtfRebaseInfo) -> Self {
+        let &Self {
+            name_offset,
+            info,
+            btf_type,
+            component_index,
+        } = self;
+        Self {
+            name_offset: rebase_info.rebase_str(name_offset),
+            info,
+            btf_type: rebase_info.rebase_type(btf_type),
+            component_index,
+        }
+    }
 }
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
@@ -1425,6 +1810,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> {