Object file improvements

Do not memory map and unmap the same object file 3 times, and make the
code a bit cleaner.

In order to make this code work I had to create a static lifetime, and
then manually drop it, this is not ideal, but it's the best I could do
after many hours.

Also, handle errors a bit more gracefully, for example, before this
commit, 32 bit binaries can cause panics.

Test Plan
==========

Ran for a while and also under valgrind.

src/main.rs

@@ -16,7 +16,7 @@ use tracing_subscriber::fmt::format::FmtSpan;
 use tracing_subscriber::FmtSubscriber;
 
 use lightswitch::collector::Collector;
-use lightswitch::object::build_id;
+use lightswitch::object::ObjectFile;
 use lightswitch::profiler::Profiler;
 use lightswitch::unwind_info::{compact_printing_callback, UnwindInfoBuilder};
 use primal::is_prime;
@@ -141,7 +141,8 @@ fn main() -> Result<(), Box<dyn Error>> {
     }
 
     if let Some(path) = args.show_info {
-        println!("build id {:?}", build_id(&PathBuf::from(path.clone())));
+        let objet_file = ObjectFile::new(&PathBuf::from(path.clone())).unwrap();
+        println!("build id {:?}", objet_file.build_id());
         let unwind_info: Result<UnwindInfoBuilder<'_>, anyhow::Error> =
             UnwindInfoBuilder::with_callback(&path, |_| {});
         println!("unwind info {:?}", unwind_info.unwrap().process());

src/object.rs

@@ -2,6 +2,7 @@ use std::fs;
 use std::io::Read;
 use std::path::PathBuf;
 
+use anyhow::{anyhow, Result};
 use data_encoding::HEXUPPER;
 use memmap2;
 use ring::digest::{Context, Digest, SHA256};
@@ -21,109 +22,128 @@ pub enum BuildId {
     Sha256(String),
 }
 
-fn sha256_digest<R: Read>(mut reader: R) -> Digest {
-    let mut context = Context::new(&SHA256);
-    let mut buffer = [0; 1024];
+pub struct ElfLoad {
+    pub offset: u64,
+    pub vaddr: u64,
+}
 
-    loop {
-        let count = reader.read(&mut buffer).unwrap();
-        if count == 0 {
-            break;
+#[derive(Debug)]
+pub struct ObjectFile<'a> {
+    leaked_mmap_ptr: *const memmap2::Mmap,
+    object: object::File<'a>,
+}
+
+impl Drop for ObjectFile<'_> {
+    fn drop(&mut self) {
+        unsafe {
+            let _to_free = Box::from_raw(self.leaked_mmap_ptr as *mut memmap2::Mmap);
         }
-        context.update(&buffer[..count]);
     }
-
-    context.finish()
 }
 
-pub fn build_id(path: &PathBuf) -> anyhow::Result<BuildId> {
-    let file = fs::File::open(path).unwrap();
-    let mmap = unsafe { memmap2::Mmap::map(&file) }.unwrap();
-    let object = object::File::parse(&*mmap).unwrap();
-
-    let build_id = object.build_id()?;
-
-    if let Some(bytes) = build_id {
-        return Ok(BuildId::Gnu(
-            bytes
-                .iter()
-                .map(|b| format!("{:02x}", b))
-                .collect::<Vec<_>>()
-                .join(""),
-        ));
+impl ObjectFile<'_> {
+    pub fn new(path: &PathBuf) -> Result<Self> {
+        let file = fs::File::open(path)?;
+        let mmap = unsafe { memmap2::Mmap::map(&file) }?;
+        let mmap = Box::new(mmap);
+        let leaked = Box::leak(mmap);
+        let object = object::File::parse(&**leaked)?;
+
+        Ok(ObjectFile {
+            leaked_mmap_ptr: leaked as *const memmap2::Mmap,
+            object,
+        })
     }
 
-    // Golang (the Go toolchain does not interpret these bytes as we do).
-    for section in object.sections() {
-        if section.name().unwrap() == ".note.go.buildid" {
-            return Ok(BuildId::Go(
-                section
-                    .data()
-                    .unwrap()
+    pub fn build_id(&self) -> anyhow::Result<BuildId> {
+        let object = &self.object;
+        let build_id = object.build_id()?;
+
+        if let Some(bytes) = build_id {
+            return Ok(BuildId::Gnu(
+                bytes
                     .iter()
                     .map(|b| format!("{:02x}", b))
                     .collect::<Vec<_>>()
                     .join(""),
             ));
         }
-    }
 
-    // No build id (rust, some other libraries).
-    for section in object.sections() {
-        if section.name().unwrap() == ".text" {
-            if let Ok(section) = section.data() {
-                return Ok(BuildId::Sha256(
-                    HEXUPPER.encode(sha256_digest(section).as_ref()),
+        // Golang (the Go toolchain does not interpret these bytes as we do).
+        for section in object.sections() {
+            if section.name().unwrap() == ".note.go.buildid" {
+                return Ok(BuildId::Go(
+                    section
+                        .data()
+                        .unwrap()
+                        .iter()
+                        .map(|b| format!("{:02x}", b))
+                        .collect::<Vec<_>>()
+                        .join(""),
                 ));
             }
         }
-    }
 
-    unreachable!("A build id should always be returned");
-}
+        // No build id (rust, some other libraries).
+        for section in object.sections() {
+            if section.name().unwrap() == ".text" {
+                if let Ok(section) = section.data() {
+                    return Ok(BuildId::Sha256(
+                        HEXUPPER.encode(sha256_digest(section).as_ref()),
+                    ));
+                }
+            }
+        }
 
-pub fn is_dynamic(path: &PathBuf) -> bool {
-    let file = fs::File::open(path).unwrap();
-    let mmap = unsafe { memmap2::Mmap::map(&file) }.unwrap();
-    let object = object::File::parse(&*mmap).unwrap();
+        unreachable!("A build id should always be returned");
+    }
 
-    object.kind() == ObjectKind::Dynamic
-}
+    pub fn is_dynamic(&self) -> bool {
+        self.object.kind() == ObjectKind::Dynamic
+    }
 
-pub fn is_go(path: &PathBuf) -> bool {
-    let file = fs::File::open(path).unwrap();
-    let mmap = unsafe { memmap2::Mmap::map(&file) }.unwrap();
-    let object = object::File::parse(&*mmap).unwrap();
-
-    for section in object.sections() {
-        if let Ok(section_name) = section.name() {
-            if section_name == ".gosymtab"
-                || section_name == ".gopclntab"
-                || section_name == ".note.go.buildid"
-            {
-                return true;
+    pub fn is_go(&self) -> bool {
+        for section in self.object.sections() {
+            if let Ok(section_name) = section.name() {
+                if section_name == ".gosymtab"
+                    || section_name == ".gopclntab"
+                    || section_name == ".note.go.buildid"
+                {
+                    return true;
+                }
             }
         }
+        false
     }
-    false
-}
 
-pub struct ElfLoad {
-    pub offset: u64,
-    pub vaddr: u64,
+    pub fn elf_load(&self) -> Result<ElfLoad> {
+        let mmap = unsafe { &**self.leaked_mmap_ptr };
+        let header: &FileHeader64<Endianness> = FileHeader64::<Endianness>::parse(mmap)?;
+        let endian = header.endian()?;
+        let segments = header.program_headers(endian, mmap)?;
+
+        if let Some(segment) = segments.iter().next() {
+            return Ok(ElfLoad {
+                offset: segment.p_offset(endian),
+                vaddr: segment.p_vaddr(endian),
+            });
+        }
+
+        Err(anyhow!("no segments found"))
+    }
 }
 
-pub fn elf_load(path: &PathBuf) -> ElfLoad {
-    let file = fs::File::open(path).unwrap();
-    let mmap = unsafe { memmap2::Mmap::map(&file) }.unwrap();
-    object::File::parse(&*mmap).unwrap();
-    let header: &FileHeader64<Endianness> = FileHeader64::<Endianness>::parse(&*mmap).unwrap();
-    let endian = header.endian().unwrap();
-    let segments = header.program_headers(endian, &*mmap).unwrap();
-    let s = segments.iter().next().unwrap();
-
-    ElfLoad {
-        offset: s.p_offset(endian),
-        vaddr: s.p_vaddr(endian),
+fn sha256_digest<R: Read>(mut reader: R) -> Digest {
+    let mut context = Context::new(&SHA256);
+    let mut buffer = [0; 1024];
+
+    loop {
+        let count = reader.read(&mut buffer).unwrap();
+        if count == 0 {
+            break;
+        }
+        context.update(&buffer[..count]);
     }
+
+    context.finish()
 }

src/profiler.rs

@@ -17,7 +17,7 @@ use tracing::{debug, error, info, span, warn, Level};
 use crate::bpf::profiler_bindings::*;
 use crate::bpf::profiler_skel::{ProfilerSkel, ProfilerSkelBuilder};
 use crate::collector::*;
-use crate::object::{build_id, elf_load, is_dynamic, is_go, BuildId};
+use crate::object::{BuildId, ObjectFile};
 use crate::perf_events::setup_perf_event;
 use crate::unwind_info::{in_memory_unwind_info, remove_redundant, remove_unnecesary_markers};
 use crate::util::summarize_address_range;
@@ -839,17 +839,22 @@ impl Profiler<'_> {
                     abs_path.push(path);
 
                     let file = fs::File::open(path)?;
-
+                    let object_file = ObjectFile::new(path);
+                    if object_file.is_err() {
+                        warn!("object_file failed with {:?}", object_file);
+                        continue;
+                    }
+                    let object_file = object_file.unwrap();
                     let unwinder = Unwinder::NativeDwarf;
 
                     // Disable profiling Go applications as they are not properly supported yet.
                     // Among other things, blazesym doesn't support symbolizing Go binaries.
-                    if is_go(&abs_path) {
+                    if object_file.is_go() {
                         // todo: deal with CGO and friends
                         return Err(anyhow!("Go applications are not supported yet"));
                     }
 
-                    let Ok(build_id) = build_id(path) else {
+                    let Ok(build_id) = object_file.build_id() else {
                         continue;
                     };
 
@@ -877,19 +882,24 @@ impl Profiler<'_> {
 
                     let mut my_lock = object_files_clone.lock().expect("lock");
 
-                    let elf_load = elf_load(path);
-
-                    my_lock.insert(
-                        build_id,
-                        ObjectFileInfo {
-                            path: abs_path,
-                            file,
-                            load_offset: elf_load.offset,
-                            load_vaddr: elf_load.vaddr,
-                            is_dyn: is_dynamic(path),
-                            main_bin: i < 3,
-                        },
-                    );
+                    match object_file.elf_load() {
+                        Ok(elf_load) => {
+                            my_lock.insert(
+                                build_id,
+                                ObjectFileInfo {
+                                    path: abs_path,
+                                    file,
+                                    load_offset: elf_load.offset,
+                                    load_vaddr: elf_load.vaddr,
+                                    is_dyn: object_file.is_dynamic(),
+                                    main_bin: i < 3,
+                                },
+                            );
+                        }
+                        Err(e) => {
+                            warn!("elf_load() failed with {:?}", e);
+                        }
+                    }
                 }
                 procfs::process::MMapPath::Anonymous => {
                     mappings.push(ExecutableMapping {