From e3e8e9ce9d814dd9c502576a6616ceceaeee4ac3 Mon Sep 17 00:00:00 2001
From: Francisco Javier Honduvilla Coto <javierhonduco@gmail.com>
Date: Sun, 3 Nov 2024 16:55:06 +0000
Subject: [PATCH] object: Improve correctness

Due to the fact that `object::File` takes a reference to the mmap-ed
file, we needed to do some awful lifetime hacks to fit it in the API we
wanted. While getting rid of the raw pointer and drop impl, Ricardo
noticed that there was some lurking UB too, so this fixes it as well.

Co-authored-by: Ricardo Delfin <me@rdelfin.com>

Test Plan
=========

CI
---
 lightswitch-object/src/object.rs | 36 +++++++++++++++-----------------
 src/profiler.rs                  |  4 ++--
 2 files changed, 19 insertions(+), 21 deletions(-)

diff --git a/lightswitch-object/src/object.rs b/lightswitch-object/src/object.rs
index d8519d6..2774aec 100644
--- a/lightswitch-object/src/object.rs
+++ b/lightswitch-object/src/object.rs
@@ -34,34 +34,32 @@ pub struct ElfLoad {
 }
 
 #[derive(Debug)]
-pub struct ObjectFile<'a> {
-    leaked_mmap_ptr: *const Mmap,
-    object: object::File<'a>,
+pub struct ObjectFile {
+    /// Warning! `object` must always go above `mmap` to ensure it will be dropped
+    /// before. Rust guarantees that fields are dropped in the order they are defined.
+    object: object::File<'static>, // Its lifetime is tied to the `mmap` below.
+    mmap: Box<Mmap>,
     code_hash: Digest,
 }
 
-impl Drop for ObjectFile<'_> {
-    fn drop(&mut self) {
-        unsafe {
-            let _to_free = Box::from_raw(self.leaked_mmap_ptr as *mut Mmap);
-        }
-    }
-}
-
-impl ObjectFile<'_> {
+impl ObjectFile {
     pub fn new(path: &Path) -> Result<Self> {
         let file = fs::File::open(path)?;
-        let mmap = unsafe { Mmap::map(&file) }?;
-        let mmap = Box::new(mmap);
-        let leaked = Box::leak(mmap);
-        let object = object::File::parse(&**leaked)?;
+        // Rust offers no guarantees on whether a "move" is done virtually or by memcpying,
+        // so to ensure that the memory value is valid we store it in the heap.
+        // Safety: Memory mapping files can cause issues if the file is modified or unmapped.
+        let mmap = Box::new(unsafe { Mmap::map(&file) }?);
+        let object = object::File::parse(&**mmap)?;
+        // Safety: The lifetime of `object` will outlive `mmap`'s. We ensure `mmap` lives as long as
+        // `object` by defining `object` before.
+        let object =
+            unsafe { std::mem::transmute::<object::File<'_>, object::File<'static>>(object) };
         let Some(code_hash) = code_hash(&object) else {
             return Err(anyhow!("code hash is None"));
         };
-
         Ok(ObjectFile {
-            leaked_mmap_ptr: leaked as *const Mmap,
             object,
+            mmap,
             code_hash,
         })
     }
@@ -115,7 +113,7 @@ impl ObjectFile<'_> {
     }
 
     pub fn elf_load_segments(&self) -> Result<Vec<ElfLoad>> {
-        let mmap = unsafe { &**self.leaked_mmap_ptr };
+        let mmap = &**self.mmap;
 
         match FileKind::parse(mmap) {
             Ok(FileKind::Elf32) => {
diff --git a/src/profiler.rs b/src/profiler.rs
index 957be81..3661ca3 100644
--- a/src/profiler.rs
+++ b/src/profiler.rs
@@ -140,12 +140,12 @@ impl Default for Profiler<'_> {
 }
 
 /// Extract the vdso object file loaded in the address space of each process.
-fn fetch_vdso_info<'a>(
+fn fetch_vdso_info(
     pid: Pid,
     start_addr: u64,
     end_addr: u64,
     offset: u64,
-) -> Result<(PathBuf, ObjectFile<'a>)> {
+) -> Result<(PathBuf, ObjectFile)> {
     // Read raw memory
     let file = fs::File::open(format!("/proc/{}/mem", pid))?;
     let size = end_addr - start_addr;