rmake: Add helpers for iterating symbols and objects

src/tools/run-make-support/src/lib.rs

@@ -16,11 +16,11 @@ pub mod diff;
 pub mod env;
 pub mod external_deps;
 pub mod linker;
+pub mod obj;
 pub mod path_helpers;
 pub mod run;
 pub mod scoped_run;
 pub mod string;
-pub mod symbols;
 pub mod targets;
 pub mod wasm;
 
@@ -34,6 +34,16 @@ pub mod rfs {
     pub use crate::fs::*;
 }
 
+/// Older version of the [`obj`] module. Prefer importing from there instead.
+pub mod symbols {
+    pub use crate::obj::{
+        ContainsAllSymbolSubstringsOutcome, ContainsAllSymbolsOutcome,
+        exported_dynamic_symbol_names, object_contains_all_symbol_substring,
+        object_contains_all_symbols, object_contains_any_symbol,
+        object_contains_any_symbol_substring,
+    };
+}
+
 // Re-exports of third-party library crates.
 pub use bstr;
 pub use gimli;

src/tools/run-make-support/src/symbols.rs → src/tools/run-make-support/src/obj.rs

@@ -1,7 +1,95 @@
+//! Helpers for working with types from the [`object`] crate.
+
 use std::collections::BTreeSet;
-use std::path::Path;
+#[cfg(unix)]
+use std::os::unix::ffi::OsStrExt;
+use std::path::{Path, PathBuf};
+
+use bstr::BStr;
+pub use object::read::archive::ArchiveFile;
+use object::{self, Object, ObjectSymbol, Symbol};
+pub use object::{File as ObjFile, Result as ObjResult};
+
+use crate::rfs;
+
+/// Thin wrapper for owning data used by `object`, which may be either an archive or a single
+/// object file.
+pub struct BinFile {
+    path: PathBuf,
+    data: Vec<u8>,
+}
+
+impl BinFile {
+    /// Read an archive or object file at `path`.
+    pub fn read_path<P: Into<PathBuf>>(path: P) -> Self {
+        let path = path.into();
+        Self { data: rfs::read(&path), path }
+    }
+
+    /// Access the owned buffer as an archive file.
+    pub fn parse_as_archive_file(&self) -> ObjResult<ArchiveFile<'_>> {
+        ArchiveFile::parse(self.data.as_slice())
+    }
+
+    /// Access the owned buffer as an archive file.
+    pub fn parse_as_obj_file(&self) -> ObjResult<ObjFile<'_>> {
+        ObjFile::parse(self.data.as_slice())
+    }
 
-use object::{self, Object, ObjectSymbol};
+    /// If the file is an archive, run the callback for each object file. If it is an object
+    /// file, the callback will run once.
+    ///
+    /// The callback receives the parsed object file and its name in the archive or on disk.
+    pub fn for_each_object(&self, mut f: impl FnMut(ObjFile<'_>, &BStr)) {
+        // Try as an archive first.
+        let as_archive = self.parse_as_archive_file();
+        if let Ok(archive) = as_archive {
+            for member in archive.members() {
+                let member = member.expect("failed to access member");
+                let obj_data = member.data(self.data.as_slice()).expect("failed to access object");
+                let obj = ObjFile::parse(obj_data).expect("failed to parse object");
+                f(obj, BStr::new(member.name()));
+            }
+
+            return;
+        }
+
+        // Fall back to parsing as an object file.
+        let as_obj = self.parse_as_obj_file();
+        if let Ok(obj) = as_obj {
+            let path_os = self.path.as_os_str();
+            let path = cfg_select! {
+                unix => path_os.as_bytes(),
+                _ => path_os
+                    .to_str()
+                    .unwrap_or_else(|| panic!("non-utf-8 path on non-unix: {:?}", self.path))
+                    .as_bytes(),
+            };
+            f(obj, BStr::new(path));
+            return;
+        }
+
+        panic!(
+            "failed to parse {:?} as either an archive or a object file: {:?}, {:?}",
+            self.path,
+            as_archive.unwrap_err(),
+            as_obj.unwrap_err(),
+        );
+    }
+
+    /// Do something with each symbol in an archive or object file.
+    ///
+    /// The callback receives:
+    ///
+    /// * The symbol.
+    /// * The parsed object file that contains the symbol.
+    /// * The name of the object file in its archive or on disk.
+    pub fn for_each_symbol(&self, mut f: impl FnMut(Symbol<'_, '_>, &ObjFile<'_>, &BStr)) {
+        self.for_each_object(|obj, obj_path| {
+            obj.symbols().for_each(|sym| f(sym, &obj, obj_path));
+        });
+    }
+}
 
 /// Given an [`object::File`], find the exported dynamic symbol names via
 /// [`object::Object::exports`]. This does not distinguish between which section the symbols appear

tests/run-make/libstd-no-protected/rmake.rs

@@ -4,11 +4,9 @@
 
 //@ only-x86_64-unknown-linux-gnu
 
-use run_make_support::object::Endianness;
-use run_make_support::object::read::archive::ArchiveFile;
-use run_make_support::object::read::elf::{FileHeader as _, SectionHeader as _};
-use run_make_support::rfs::read;
-use run_make_support::{has_prefix, has_suffix, object, path, rustc, shallow_find_files, target};
+use run_make_support::obj::BinFile;
+use run_make_support::object::{Endianness, ObjectSymbol, SymbolFlags};
+use run_make_support::{has_prefix, has_suffix, object, rustc, shallow_find_files, target};
 
 type FileHeader = run_make_support::object::elf::FileHeader64<Endianness>;
 type SymbolTable<'data> = run_make_support::object::read::elf::SymbolTable<'data, FileHeader>;
@@ -21,41 +19,25 @@ fn main() {
     });
     assert_eq!(libs.len(), 1);
     let libstd_path = libs.pop().unwrap();
-    let archive_data = read(libstd_path);
+    let archive = BinFile::read_path(libstd_path);
 
     // Parse all the object files within the libstd archive, checking defined symbols.
-    let mut num_protected = 0;
+    let mut protected = Vec::new();
     let mut num_symbols = 0;
 
-    let archive = ArchiveFile::parse(&*archive_data).unwrap();
-    for member in archive.members() {
-        let member = member.unwrap();
-        if member.name() == b"lib.rmeta" {
-            continue;
+    archive.for_each_symbol(|symbol, _obj, obj_name| {
+        if obj_name == "lib.rmeta" {
+            return;
         }
-        let data = member.data(&*archive_data).unwrap();
-
-        let header = FileHeader::parse(data).unwrap();
-        let endian = header.endian().unwrap();
-        let sections = header.sections(endian, data).unwrap();
-
-        for (section_index, section) in sections.enumerate() {
-            if section.sh_type(endian) == object::elf::SHT_SYMTAB {
-                let symbols =
-                    SymbolTable::parse(endian, data, &sections, section_index, section).unwrap();
-                for symbol in symbols.symbols() {
-                    if symbol.st_visibility() == object::elf::STV_PROTECTED {
-                        num_protected += 1;
-                    }
-                    num_symbols += 1;
-                }
-            }
+        if symbol.flags().elf_visibility().expect("non-elf file") == object::elf::STV_PROTECTED {
+            protected.push(String::from_utf8_lossy(symbol.name_bytes().unwrap()).into_owned());
         }
-    }
+        num_symbols += 1;
+    });
 
     // If there were no symbols at all, then something is wrong with the test.
     assert_ne!(num_symbols, 0);
 
     // The purpose of this test - check that no symbols have protected visibility.
-    assert_eq!(num_protected, 0);
+    assert!(protected.is_empty(), "found protected symbols {protected:#?}");
 }