SimpleJIT - Add a Drop impl to PtrLen (#1001)

* Implement SimpleJIT deallocation for non-windows targets.

* Remove custom Drop from feature(selinux-fix).

* Fix typo in unprotect error message.

* Drop SimpleJIT memory by default and add free_memory method instead.

* Move free_memory to a handle returned by Module::finish.

* Reduce memory handle content to necessary fields only.

* Use lower overhead method for leaking.

cranelift/simplejit/src/backend.rs

@@ -118,9 +118,7 @@ pub struct SimpleJITBackend {
     isa: Box<dyn TargetIsa>,
     symbols: HashMap<String, *const u8>,
     libcall_names: Box<dyn Fn(ir::LibCall) -> String>,
-    code_memory: Memory,
+    memory: SimpleJITMemoryHandle,
-    readonly_memory: Memory,
-    writable_memory: Memory,
 }
 
 /// A record of a relocation to perform.
@@ -150,6 +148,13 @@ pub struct SimpleJITCompiledData {
     relocs: Vec<RelocRecord>,
 }
 
+/// A handle to allow freeing memory allocated by the `Backend`.
+pub struct SimpleJITMemoryHandle {
+    code: Memory,
+    readonly: Memory,
+    writable: Memory,
+}
+
 impl SimpleJITBackend {
     fn lookup_symbol(&self, name: &str) -> *const u8 {
         match self.symbols.get(name) {
@@ -199,23 +204,32 @@ impl<'simple_jit_backend> Backend for SimpleJITBackend {
     type CompiledData = SimpleJITCompiledData;
 
     /// SimpleJIT emits code and data into memory, and provides raw pointers
-    /// to them.
+    /// to them. They are valid for the remainder of the program's life, unless
+    /// [`free_memory`] is used.
+    ///
+    /// [`free_memory`]: #method.free_memory
     type FinalizedFunction = *const u8;
     type FinalizedData = (*mut u8, usize);
 
     /// SimpleJIT emits code and data into memory as it processes them, so it
     /// doesn't need to provide anything after the `Module` is complete.
-    type Product = ();
+    /// The handle object that is returned can optionally be used to free
+    /// allocated memory if required.
+    type Product = SimpleJITMemoryHandle;
 
     /// Create a new `SimpleJITBackend`.
     fn new(builder: SimpleJITBuilder) -> Self {
+        let memory = SimpleJITMemoryHandle {
+            code: Memory::new(),
+            readonly: Memory::new(),
+            writable: Memory::new(),
+        };
+
         Self {
             isa: builder.isa,
             symbols: builder.symbols,
             libcall_names: builder.libcall_names,
-            code_memory: Memory::new(),
+            memory,
-            readonly_memory: Memory::new(),
-            writable_memory: Memory::new(),
         }
     }
 
@@ -248,7 +262,8 @@ impl<'simple_jit_backend> Backend for SimpleJITBackend {
     ) -> ModuleResult<Self::CompiledFunction> {
         let size = code_size as usize;
         let ptr = self
-            .code_memory
+            .memory
+            .code
             .allocate(size, EXECUTABLE_DATA_ALIGNMENT)
             .expect("TODO: handle OOM etc.");
 
@@ -303,11 +318,13 @@ impl<'simple_jit_backend> Backend for SimpleJITBackend {
 
         let size = init.size();
         let storage = if writable {
-            self.writable_memory
+            self.memory
+                .writable
                 .allocate(size, align.unwrap_or(WRITABLE_DATA_ALIGNMENT))
                 .expect("TODO: handle OOM etc.")
         } else {
-            self.readonly_memory
+            self.memory
+                .readonly
                 .allocate(size, align.unwrap_or(READONLY_DATA_ALIGNMENT))
                 .expect("TODO: handle OOM etc.")
         };
@@ -479,13 +496,20 @@ impl<'simple_jit_backend> Backend for SimpleJITBackend {
 
     fn publish(&mut self) {
         // Now that we're done patching, prepare the memory for execution!
-        self.readonly_memory.set_readonly();
+        self.memory.readonly.set_readonly();
-        self.code_memory.set_readable_and_executable();
+        self.memory.code.set_readable_and_executable();
     }
 
-    /// SimpleJIT emits code and data into memory as it processes them, so it
+    /// SimpleJIT emits code and data into memory as it processes them. This
-    /// doesn't need to provide anything after the `Module` is complete.
+    /// method performs no additional processing, but returns a handle which
-    fn finish(self) {}
+    /// allows freeing the allocated memory. Otherwise said memory is leaked
+    /// to enable safe handling of the resulting pointers.
+    ///
+    /// This method does not need to be called when access to the memory
+    /// handle is not required.
+    fn finish(self) -> Self::Product {
+        self.memory
+    }
 }
 
 #[cfg(not(windows))]
@@ -531,6 +555,22 @@ fn lookup_with_dlsym(name: &str) -> *const u8 {
     }
 }
 
+impl SimpleJITMemoryHandle {
+    /// Free memory allocated for code and data segments of compiled functions.
+    ///
+    /// # Safety
+    ///
+    /// Because this function invalidates any pointers retrived from the
+    /// corresponding module, it should only be used when none of the functions
+    /// from that module are currently executing and none of the`fn` pointers
+    /// are called afterwards.
+    pub unsafe fn free_memory(&mut self) {
+        self.code.free_memory();
+        self.readonly.free_memory();
+        self.writable.free_memory();
+    }
+}
+
 struct SimpleJITRelocSink {
     pub relocs: Vec<RelocRecord>,
 }

cranelift/simplejit/src/memory.rs

@@ -105,8 +105,26 @@ impl PtrLen {
     }
 }
 
+// `MMapMut` from `cfg(feature = "selinux-fix")` already deallocates properly.
+#[cfg(all(not(target_os = "windows"), not(feature = "selinux-fix")))]
+impl Drop for PtrLen {
+    fn drop(&mut self) {
+        if !self.ptr.is_null() {
+            unsafe {
+                region::protect(self.ptr, self.len, region::Protection::ReadWrite)
+                    .expect("unable to unprotect memory");
+                libc::free(self.ptr as _);
+            }
+        }
+    }
+}
+
+// TODO: add a `Drop` impl for `cfg(target_os = "windows")`
+
 /// JIT memory manager. This manages pages of suitably aligned and
-/// accessible memory.
+/// accessible memory. Memory will be leaked by default to have
+/// function pointers remain valid for the remainder of the
+/// program's life.
 pub struct Memory {
     allocations: Vec<PtrLen>,
     executable: usize,
@@ -209,9 +227,22 @@ impl Memory {
             }
         }
     }
+
+    /// Frees all allocated memory regions that would be leaked otherwise.
+    /// Likely to invalidate existing function pointers, causing unsafety.
+    pub unsafe fn free_memory(&mut self) {
+        self.allocations.clear();
+    }
 }
 
-// TODO: Implement Drop to unprotect and deallocate the memory?
+impl Drop for Memory {
+    fn drop(&mut self) {
+        // leak memory to guarantee validity of function pointers
+        mem::replace(&mut self.allocations, Vec::new())
+            .into_iter()
+            .for_each(mem::forget);
+    }
+}
 
 #[cfg(test)]
 mod tests {