tree: add support for regex resolution for ids and types

Signed-off-by: Antoine Tenart <[email protected]>

.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 -F elf
+  test_script: cargo test --verbose -F elf,regex
   check_script:
     - cargo fmt --check
     - cargo clippy -- -D warnings

Cargo.toml

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

src/btf.rs

@@ -95,6 +95,33 @@ impl Btf {
         self.obj.resolve_ids_by_name(name)
     }
 
+    /// Find a list of BTF ids whose names match a regex.
+    #[cfg(feature = "regex")]
+    pub fn resolve_ids_by_regex(&self, re: &regex::Regex) -> Result<Vec<u32>> {
+        let mut ids = Vec::new();
+
+        if let Some(base) = &self.base {
+            if let Ok(mut ids_base) = base.resolve_ids_by_regex(re) {
+                ids.append(&mut ids_base);
+            }
+        }
+        if let Ok(mut ids_obj) = self.resolve_split_ids_by_regex(re) {
+            ids.append(&mut ids_obj);
+        }
+
+        if ids.is_empty() {
+            bail!("No id linked to regex {re}");
+        }
+        Ok(ids)
+    }
+
+    /// Find a list of BTF ids whose names match a regex, using the split BTF
+    /// definition only. For internal use only.
+    #[cfg(feature = "regex")]
+    pub(crate) fn resolve_split_ids_by_regex(&self, re: &regex::Regex) -> Result<Vec<u32>> {
+        self.obj.resolve_ids_by_regex(re)
+    }
+
     /// Find a BTF type using its id as a key.
     pub fn resolve_type_by_id(&self, id: u32) -> Result<Type> {
         match &self.base {
@@ -132,6 +159,35 @@ impl Btf {
         self.obj.resolve_types_by_name(name)
     }
 
+    /// Find a list of BTF types using a regex describing their name as a key.
+    #[cfg(feature = "regex")]
+    pub fn resolve_types_by_regex(&self, re: &regex::Regex) -> Result<Vec<Type>> {
+        let mut types = Vec::new();
+
+        if let Some(base) = &self.base {
+            if let Ok(mut types_base) = base.resolve_types_by_regex(re) {
+                types.append(&mut types_base);
+            }
+        }
+        if let Ok(mut types_obj) = self.resolve_split_types_by_regex(re) {
+            types.append(&mut types_obj);
+        }
+
+        if types.is_empty() {
+            // Keep "id" and not "type" below to be consitent with
+            // BtfObj::resolve_types_by_name.
+            bail!("No id linked to regex {re}");
+        }
+        Ok(types)
+    }
+
+    /// Find a list of BTF types using a regex describing their name as a key,
+    /// using the split BTF definition only. For internal use only.
+    #[cfg(feature = "regex")]
+    pub fn resolve_split_types_by_regex(&self, re: &regex::Regex) -> Result<Vec<Type>> {
+        self.obj.resolve_types_by_regex(re)
+    }
+
     /// Resolve a name referenced by a Type which is defined in the current BTF
     /// object.
     pub fn resolve_name<T: BtfType + ?Sized>(&self, r#type: &T) -> Result<String> {

src/lib.rs

@@ -99,6 +99,7 @@
 //!
 //! - elf: Enable helpers parsing the .BTF section of ELF files in
 //!   `utils::elf`.
+//! - regex: Enable name resolutions by regex (`regex::Regex`).
 //! - test_runtime: Use the system's runtime BTF files to perform extra
 //!   integration tests.
 

src/obj.rs

@@ -163,6 +163,20 @@ impl BtfObj {
         }
     }
 
+    /// Find a list of BTF ids whose names match a regex.
+    #[cfg(feature = "regex")]
+    pub(super) fn resolve_ids_by_regex(&self, re: &regex::Regex) -> Result<Vec<u32>> {
+        Ok(self
+            .strings
+            .iter()
+            .filter_map(|(name, ids)| match re.is_match(name) {
+                true => Some(ids.clone()),
+                false => None,
+            })
+            .flatten()
+            .collect::<Vec<u32>>())
+    }
+
     /// Find a BTF type using its id as a key.
     pub(super) fn resolve_type_by_id(&self, id: u32) -> Result<Type> {
         match self.types.get(&id) {
@@ -180,6 +194,16 @@ impl BtfObj {
         Ok(types)
     }
 
+    /// Find a list of BTF types using a regex describing their name as a key.
+    #[cfg(feature = "regex")]
+    pub(super) fn resolve_types_by_regex(&self, re: &regex::Regex) -> Result<Vec<Type>> {
+        let mut types = Vec::new();
+        for id in self.resolve_ids_by_regex(re)? {
+            types.push(self.resolve_type_by_id(id)?);
+        }
+        Ok(types)
+    }
+
     /// Resolve a name referenced by a Type which is defined in the current BTF
     /// object.
     pub(super) fn resolve_name<T: BtfType + ?Sized>(&self, r#type: &T) -> Result<String> {

src/utils/collection.rs

@@ -178,6 +178,37 @@ impl BtfCollection {
         Ok(ids)
     }
 
+    /// Find a list of BTF ids whose names match a regex.
+    ///
+    /// Matching ids can be found in multiple underlying BTF, thus this function
+    /// returns a list of tuples containing each a reference to `NamedBtf`
+    /// (representing the BTF where a match was found) and the id. Further
+    /// lookups must be done using the `Btf` object contained in the linked
+    /// `NamedBtf` one.
+    #[cfg(feature = "regex")]
+    pub fn resolve_ids_by_regex(&self, re: &regex::Regex) -> Result<Vec<(&NamedBtf, u32)>> {
+        let mut ids = self
+            .base
+            .btf
+            .resolve_ids_by_regex(re)
+            .unwrap_or_default()
+            .drain(..)
+            .map(|i| (&self.base, i))
+            .collect::<Vec<_>>();
+
+        for split in self.split.iter() {
+            if let Ok(mut mod_ids) = split.btf.resolve_split_ids_by_regex(re) {
+                mod_ids.drain(..).for_each(|i| ids.push((split, i)));
+            }
+        }
+
+        if ids.is_empty() {
+            bail!("No id linked to regex {re}");
+        }
+
+        Ok(ids)
+    }
+
     /// Find a list of BTF types using their name as a key. Matching types can
     /// be found in multiple underlying BTF, thus this function returns a list
     /// of tuples containing each a reference to `NamedBtf` (representing the
@@ -206,6 +237,36 @@ impl BtfCollection {
         Ok(types)
     }
 
+    /// Find a list of BTF types using a regex describing their name as a key.
+    /// Matching types can be found in multiple underlying BTF, thus this
+    /// function returns a list of tuples containing each a reference to
+    /// `NamedBtf` (representing the BTF where a match was found) and the type.
+    /// Further lookups must be done using the `Btf` object contained in the
+    /// linked `NamedBtf` one.
+    #[cfg(feature = "regex")]
+    pub fn resolve_types_by_regex(&self, re: &regex::Regex) -> Result<Vec<(&NamedBtf, Type)>> {
+        let mut types = self
+            .base
+            .btf
+            .resolve_types_by_regex(re)
+            .unwrap_or_default()
+            .drain(..)
+            .map(|t| (&self.base, t))
+            .collect::<Vec<_>>();
+
+        for split in self.split.iter() {
+            if let Ok(mut mod_types) = split.btf.resolve_split_types_by_regex(re) {
+                mod_types.drain(..).for_each(|t| types.push((split, t)));
+            }
+        }
+
+        if types.is_empty() {
+            bail!("No type linked to regex {re}");
+        }
+
+        Ok(types)
+    }
+
     // Internal helper to extract a file name as a String from a Path.
     fn file_name(path: &Path) -> Result<String> {
         Ok(match path.file_name() {

tests/integration_test.rs

@@ -350,6 +350,38 @@ fn resolve_split_func(btf: Btf) {
     assert_eq!(struct1.members.len(), 28);
 }
 
+#[test_case(split_file())]
+#[test_case(split_bytes())]
+#[cfg_attr(feature = "elf", test_case(split_elf()))]
+#[cfg(feature = "regex")]
+fn resolve_regex(btf: Btf) {
+    use std::collections::HashSet;
+
+    // Look for drop reason enums:
+    // - skb_drop_reason
+    // - mac80211_drop_reason
+    // - ovs_drop_reason
+    let re = regex::Regex::new(r"^[[:alnum:]]+_drop_reason$").unwrap();
+    let ids = btf.resolve_ids_by_regex(&re).unwrap();
+    assert_eq!(ids.len(), 3);
+
+    let types = btf.resolve_types_by_regex(&re).unwrap();
+    assert_eq!(types.len(), 3);
+
+    let mut reasons = HashSet::from(["ovs_drop_reason", "mac80211_drop_reason", "skb_drop_reason"]);
+    let get_enum_name = |r#type: &Type| {
+        let r#enum = match r#type {
+            Type::Enum(r#enum) => r#enum,
+            _ => panic!("Type is not an enum"),
+        };
+        btf.resolve_name(r#enum).unwrap()
+    };
+    types.iter().for_each(|t| {
+        assert!(reasons.remove(get_enum_name(t).as_str()));
+    });
+    assert!(reasons.is_empty());
+}
+
 #[test]
 #[cfg_attr(not(feature = "test_runtime"), ignore)]
 fn test_split_files() {
@@ -458,3 +490,36 @@ fn btfc(btfc: utils::collection::BtfCollection) {
     };
     assert_eq!(ovs.resolve_name(&func).unwrap(), "queue_userspace_packet");
 }
+
+#[test_case(btfc_files())]
+#[test_case(btfc_bytes())]
+#[test_case(btfc_dir())]
+#[cfg_attr(feature = "elf", test_case(btfc_elf()))]
+#[cfg(feature = "regex")]
+fn btfc_resolve_regex(btfc: utils::collection::BtfCollection) {
+    use std::collections::HashSet;
+
+    // Look for drop reason enums:
+    // - skb_drop_reason
+    // - mac80211_drop_reason
+    // - ovs_drop_reason
+    let re = regex::Regex::new(r"^[[:alnum:]]+_drop_reason$").unwrap();
+    let ids = btfc.resolve_ids_by_regex(&re).unwrap();
+    assert_eq!(ids.len(), 3);
+
+    let types = btfc.resolve_types_by_regex(&re).unwrap();
+    assert_eq!(types.len(), 3);
+
+    let mut reasons = HashSet::from(["ovs_drop_reason", "mac80211_drop_reason", "skb_drop_reason"]);
+    let get_enum_name = |r#type: &(&utils::collection::NamedBtf, btf_rs::Type)| {
+        let (nbtf, r#enum) = match r#type {
+            (nbtf, Type::Enum(r#enum)) => (nbtf, r#enum),
+            _ => panic!("Type is not an enum"),
+        };
+        nbtf.resolve_name(r#enum).unwrap()
+    };
+    types.iter().for_each(|t| {
+        assert!(reasons.remove(get_enum_name(t).as_str()));
+    });
+    assert!(reasons.is_empty());
+}