From 33390b09658af4bd7190b60388a7429d2cd1bdb9 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 | 29 +++++++++++++----------------
 1 file changed, 13 insertions(+), 16 deletions(-)

diff --git a/lightswitch-object/src/object.rs b/lightswitch-object/src/object.rs
index d8519d6..3e26ec3 100644
--- a/lightswitch-object/src/object.rs
+++ b/lightswitch-object/src/object.rs
@@ -35,33 +35,30 @@ pub struct ElfLoad {
 
 #[derive(Debug)]
 pub struct ObjectFile<'a> {
-    leaked_mmap_ptr: *const Mmap,
+    /// 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<'a>,
+    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<'_> {
     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.
+        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 +112,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) => {