Merge pull request #17 from retis-org/at/elf

Add ELF helpers (& various fixes)

.cirrus.yml

@@ -14,7 +14,7 @@ test_task:
     - rustup component add rustfmt
     - rustup component add clippy
   build_script: cargo build --verbose
-  test_script: cargo test --verbose
+  test_script: cargo test --verbose -F elf
   check_script:
     - cargo fmt --check
     - cargo clippy -- -D warnings

Cargo.toml

@@ -18,9 +18,11 @@ maintenance = { status = "actively-developed" }
 [dependencies]
 anyhow = "1.0"
 byteorder = "1.5"
+elf = { version = "0.7", optional = true }
 
 [dev-dependencies]
 test-case = "3.2"
 
 [features]
+elf = ["dep:elf"]
 test_runtime = []

src/lib.rs

@@ -97,6 +97,8 @@
 //!
 //! ### Feature flags
 //!
+//! - elf: Enable helpers parsing the .BTF section of ELF files in
+//!   `utils::elf`.
 //! - test_runtime: Use the system's runtime BTF files to perform extra
 //!   integration tests.
 

src/utils/collection.rs

@@ -157,11 +157,7 @@ impl BtfCollection {
     /// the `Btf` object contained in the linked `NamedBtf` one.
     pub fn resolve_ids_by_name(&self, name: &str) -> Result<Vec<(&NamedBtf, u32)>> {
         let mut ids = Vec::new();
-
-        let mut base_ids = match self.base.btf.resolve_ids_by_name(name) {
-            Ok(base_ids) => base_ids,
-            _ => Vec::new(), // Id not found in base.
-        };
+        let mut base_ids = self.base.btf.resolve_ids_by_name(name).unwrap_or_default();
 
         for split in self.split.iter() {
             if let Ok(mut mod_ids) = split.btf.resolve_split_ids_by_name(name) {
@@ -186,11 +182,11 @@ impl BtfCollection {
     /// using the `Btf` object contained in the linked `NamedBtf` one.
     pub fn resolve_types_by_name(&self, name: &str) -> Result<Vec<(&NamedBtf, Type)>> {
         let mut types = Vec::new();
-
-        let mut base_types = match self.base.btf.resolve_types_by_name(name) {
-            Ok(base_types) => base_types,
-            _ => Vec::new(), // Id not found in base.
-        };
+        let mut base_types = self
+            .base
+            .btf
+            .resolve_types_by_name(name)
+            .unwrap_or_default();
 
         for split in self.split.iter() {
             if let Ok(mut mod_types) = split.btf.resolve_split_types_by_name(name) {

src/utils/elf.rs

@@ -0,0 +1,79 @@
+use std::{
+    fs::{self, File},
+    path::Path,
+};
+
+use anyhow::{anyhow, bail, Result};
+use elf::{endian::AnyEndian, ElfStream};
+
+use crate::utils::collection::BtfCollection;
+
+/// Extract raw BTF data from the .BTF elf section of the given file. Output can
+/// be used to fed `from_bytes` constructors in this library.
+pub fn extract_btf_from_file<P: AsRef<Path>>(path: P) -> Result<Vec<u8>> {
+    let file = File::open(&path)
+        .map_err(|e| anyhow!("Could not open {}: {e}", path.as_ref().display()))?;
+    let mut elf = ElfStream::<AnyEndian, _>::open_stream(file)?;
+
+    let btf_hdr = match elf.section_header_by_name(".BTF")? {
+        Some(hdr) => *hdr,
+        None => bail!("No BTF section in {}", path.as_ref().display()),
+    };
+
+    let (btf, chdr) = elf.section_data(&btf_hdr)?;
+    if chdr.is_some() {
+        bail!(
+            "Compressed BTF sections are not supported ({})",
+            path.as_ref().display()
+        );
+    }
+
+    Ok(btf.to_vec())
+}
+
+/// Given a directory containing a 'vmlinux' ELF file in its root and optional
+/// '*.ko' ELF modules in the root or any sub-directory (this maps well to a
+/// Linux build directory or /usr/lib/modules/), initialize a BtfCollection
+/// extracting BTF data from the .BTF section of those files.
+pub fn collection_from_kernel_dir<P: AsRef<Path>>(path: P) -> Result<BtfCollection> {
+    let path = path.as_ref();
+    if !path.is_dir() {
+        bail!(
+            "Can't initialize a BTF collection from {}: not a directory",
+            path.display()
+        );
+    }
+
+    // Find the base BTF file and initialize the collection.
+    let vmlinux = path.join("vmlinux");
+    let mut collection = BtfCollection::from_bytes("vmlinux", &extract_btf_from_file(vmlinux)?)?;
+
+    // Traverse the directory looking for modules.
+    fn visit_dir<P: AsRef<Path>>(dir: P, collection: &mut BtfCollection) -> Result<()> {
+        for entry in fs::read_dir(dir)? {
+            let path = entry?.path();
+            let filename = path
+                .file_name()
+                .ok_or_else(|| anyhow!("Could not get module file name"))?
+                .to_str()
+                .ok_or_else(|| anyhow!("Could not convert module name to str"))?;
+
+            if path.is_dir() {
+                visit_dir(path, collection)?;
+            } else if filename.ends_with(".ko") {
+                collection.add_split_btf_from_bytes(
+                    match filename.split_once('.') {
+                        Some((name, _)) => name,
+                        // Should not happen as we already filtered on extensions.
+                        None => bail!("Invalid module file name"),
+                    },
+                    &extract_btf_from_file(&path)?,
+                )?;
+            }
+        }
+        Ok(())
+    }
+    visit_dir(path, &mut collection)?;
+
+    Ok(collection)
+}

src/utils/mod.rs

@@ -2,3 +2,5 @@
 //! use cases.
 
 pub mod collection;
+#[cfg(feature = "elf")]
+pub mod elf;

tests/integration_test.rs

@@ -5,48 +5,75 @@ use test_case::test_case;
 use btf_rs::*;
 
 fn bytes() -> Btf {
-    Btf::from_bytes(&read("tests/data/vmlinux").unwrap()).unwrap()
+    Btf::from_bytes(&read("tests/data/btf/vmlinux").unwrap()).unwrap()
 }
 
 fn file() -> Btf {
-    Btf::from_file("tests/data/vmlinux").unwrap()
+    Btf::from_file("tests/data/btf/vmlinux").unwrap()
+}
+
+#[cfg(feature = "elf")]
+fn elf() -> Btf {
+    Btf::from_bytes(&utils::elf::extract_btf_from_file("tests/data/linux_build/vmlinux").unwrap())
+        .unwrap()
 }
 
 fn split_file() -> Btf {
-    let vmlinux = Btf::from_file("tests/data/vmlinux").unwrap();
-    Btf::from_split_file("tests/data/openvswitch", &vmlinux).unwrap()
+    let vmlinux = Btf::from_file("tests/data/btf/vmlinux").unwrap();
+    Btf::from_split_file("tests/data/btf/openvswitch", &vmlinux).unwrap()
 }
 
 fn split_bytes() -> Btf {
-    let vmlinux = Btf::from_bytes(&read("tests/data/vmlinux").unwrap()).unwrap();
-    Btf::from_split_bytes(&read("tests/data/openvswitch").unwrap(), &vmlinux).unwrap()
+    let vmlinux = Btf::from_bytes(&read("tests/data/btf/vmlinux").unwrap()).unwrap();
+    Btf::from_split_bytes(&read("tests/data/btf/openvswitch").unwrap(), &vmlinux).unwrap()
+}
+
+#[cfg(feature = "elf")]
+fn split_elf() -> Btf {
+    let vmlinux = Btf::from_bytes(
+        &utils::elf::extract_btf_from_file("tests/data/linux_build/vmlinux").unwrap(),
+    )
+    .unwrap();
+    Btf::from_split_bytes(
+        &utils::elf::extract_btf_from_file("tests/data/linux_build/net/openvswitch/openvswitch.ko")
+            .unwrap(),
+        &vmlinux,
+    )
+    .unwrap()
 }
 
 #[test_case(bytes())]
 #[test_case(file())]
+#[cfg_attr(feature = "elf", test_case(elf()))]
 #[test_case(split_file())]
 #[test_case(split_bytes())]
+#[cfg_attr(feature = "elf", test_case(split_elf()))]
 fn resolve_ids_by_name(btf: Btf) {
     // Resolve primitive type.
-    assert_eq!(btf.resolve_ids_by_name("int").unwrap().pop().unwrap(), 21);
+    assert_eq!(btf.resolve_ids_by_name("int").unwrap().pop().unwrap(), 11);
     // Resolve typedef.
-    assert_eq!(btf.resolve_ids_by_name("u64").unwrap().pop().unwrap(), 37);
+    assert_eq!(btf.resolve_ids_by_name("u64").unwrap().pop().unwrap(), 58);
     // Resolve struct.
-    assert_eq!(btf.resolve_ids_by_name("sk_buff").unwrap()[0], 3482);
+    assert_eq!(
+        btf.resolve_ids_by_name("sk_buff").unwrap().pop().unwrap(),
+        4984
+    );
     // Resolve function.
     assert_eq!(
         btf.resolve_ids_by_name("consume_skb")
             .unwrap()
             .pop()
             .unwrap(),
-        36977
+        95474
     );
 }
 
 #[test_case(bytes())]
 #[test_case(file())]
+#[cfg_attr(feature = "elf", test_case(elf()))]
 #[test_case(split_file())]
 #[test_case(split_bytes())]
+#[cfg_attr(feature = "elf", test_case(split_elf()))]
 fn iter_types(btf: Btf) {
     // Iterate without looping ensuring non BtfTypes return None.
     let vmalloc = match btf.resolve_types_by_name("vmalloc").unwrap().pop().unwrap() {
@@ -83,8 +110,10 @@ fn iter_types(btf: Btf) {
 
 #[test_case(bytes())]
 #[test_case(file())]
+#[cfg_attr(feature = "elf", test_case(elf()))]
 #[test_case(split_file())]
 #[test_case(split_bytes())]
+#[cfg_attr(feature = "elf", test_case(split_elf()))]
 fn resolve_types_by_name(btf: Btf) {
     let types = btf.resolve_types_by_name("consume_skb");
     assert!(types.is_ok());
@@ -93,16 +122,20 @@ fn resolve_types_by_name(btf: Btf) {
 
 #[test_case(bytes())]
 #[test_case(file())]
+#[cfg_attr(feature = "elf", test_case(elf()))]
 #[test_case(split_file())]
 #[test_case(split_bytes())]
+#[cfg_attr(feature = "elf", test_case(split_elf()))]
 fn resolve_types_by_name_unknown(btf: Btf) {
     assert!(btf.resolve_types_by_name("not_a_known_function").is_err());
 }
 
 #[test_case(bytes())]
 #[test_case(file())]
+#[cfg_attr(feature = "elf", test_case(elf()))]
 #[test_case(split_file())]
 #[test_case(split_bytes())]
+#[cfg_attr(feature = "elf", test_case(split_elf()))]
 fn check_resolved_type(btf: Btf) {
     let mut r#type = btf.resolve_types_by_name("sk_buff").unwrap();
 
@@ -114,8 +147,10 @@ fn check_resolved_type(btf: Btf) {
 
 #[test_case(bytes())]
 #[test_case(file())]
+#[cfg_attr(feature = "elf", test_case(elf()))]
 #[test_case(split_file())]
 #[test_case(split_bytes())]
+#[cfg_attr(feature = "elf", test_case(split_elf()))]
 fn bijection(btf: Btf) {
     let func = match btf.resolve_types_by_name("vmalloc").unwrap().pop().unwrap() {
         Type::Func(func) => func,
@@ -135,8 +170,10 @@ fn bijection(btf: Btf) {
 
 #[test_case(bytes())]
 #[test_case(file())]
+#[cfg_attr(feature = "elf", test_case(elf()))]
 #[test_case(split_file())]
 #[test_case(split_bytes())]
+#[cfg_attr(feature = "elf", test_case(split_elf()))]
 fn resolve_function(btf: Btf) {
     let func = match btf
         .resolve_types_by_name("kfree_skb_reason")
@@ -207,6 +244,7 @@ fn wrong_file() {
 
 #[test_case(split_file())]
 #[test_case(split_bytes())]
+#[cfg_attr(feature = "elf", test_case(split_elf()))]
 fn resolve_split_struct(btf: Btf) {
     let r#struct = btf
         .resolve_types_by_name("datapath")
@@ -245,6 +283,7 @@ fn resolve_split_struct(btf: Btf) {
 
 #[test_case(split_file())]
 #[test_case(split_bytes())]
+#[cfg_attr(feature = "elf", test_case(split_elf()))]
 fn resolve_split_func(btf: Btf) {
     // Resolve the following function:
     // static int queue_userspace_packet(struct datapath *dp, struct sk_buff *skb,
@@ -337,30 +376,36 @@ fn test_split_files() {
 }
 
 fn btfc_files() -> utils::collection::BtfCollection {
-    let mut btfc = utils::collection::BtfCollection::from_file("tests/data/vmlinux").unwrap();
-    btfc.add_split_btf_from_file("tests/data/openvswitch")
+    let mut btfc = utils::collection::BtfCollection::from_file("tests/data/btf/vmlinux").unwrap();
+    btfc.add_split_btf_from_file("tests/data/btf/openvswitch")
         .unwrap();
     btfc
 }
 
 fn btfc_bytes() -> utils::collection::BtfCollection {
     let mut btfc = utils::collection::BtfCollection::from_bytes(
         "vmlinux",
-        &read("tests/data/vmlinux").unwrap(),
+        &read("tests/data/btf/vmlinux").unwrap(),
     )
     .unwrap();
-    btfc.add_split_btf_from_bytes("openvswitch", &read("tests/data/openvswitch").unwrap())
+    btfc.add_split_btf_from_bytes("openvswitch", &read("tests/data/btf/openvswitch").unwrap())
         .unwrap();
     btfc
 }
 
 fn btfc_dir() -> utils::collection::BtfCollection {
-    utils::collection::BtfCollection::from_dir("tests/data", "vmlinux").unwrap()
+    utils::collection::BtfCollection::from_dir("tests/data/btf", "vmlinux").unwrap()
+}
+
+#[cfg(feature = "elf")]
+fn btfc_elf() -> utils::collection::BtfCollection {
+    utils::elf::collection_from_kernel_dir("tests/data/linux_build").unwrap()
 }
 
 #[test_case(btfc_files())]
 #[test_case(btfc_bytes())]
 #[test_case(btfc_dir())]
+#[cfg_attr(feature = "elf", test_case(btfc_elf()))]
 fn btfc(btfc: utils::collection::BtfCollection) {
     // Resolve a function from vmlinux.
     let mut types = btfc.resolve_types_by_name("vmalloc").unwrap();