Upgrading dlmalloc to 0.2.1

Cargo.lock

@@ -952,9 +952,9 @@ dependencies = [
 
 [[package]]
 name = "dlmalloc"
-version = "0.1.4"
+version = "0.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "35055b1021724f4eb5262eb49130eebff23fc59fc5a14160e05faad8eeb36673"
+checksum = "332570860c2edf2d57914987bf9e24835425f75825086b6ba7d1e6a3e4f1f254"
 dependencies = [
  "compiler_builtins",
  "libc",

library/std/Cargo.toml

@@ -36,7 +36,7 @@ features = ['read_core', 'elf', 'macho', 'pe']
 rand = "0.7"
 
 [target.'cfg(any(all(target_arch = "wasm32", not(target_os = "emscripten")), all(target_vendor = "fortanix", target_env = "sgx")))'.dependencies]
-dlmalloc = { version = "0.1", features = ['rustc-dep-of-std'] }
+dlmalloc = { version = "0.2.1", features = ['rustc-dep-of-std'] }
 
 [target.x86_64-fortanix-unknown-sgx.dependencies]
 fortanix-sgx-abi = { version = "0.3.2", features = ['rustc-dep-of-std'] }

library/std/src/sys/sgx/abi/mem.rs

@@ -12,6 +12,18 @@ pub(crate) unsafe fn rel_ptr_mut<T>(offset: u64) -> *mut T {
 
 extern "C" {
     static ENCLAVE_SIZE: usize;
+    static HEAP_BASE: u64;
+    static HEAP_SIZE: usize;
+}
+
+/// Returns the base memory address of the heap
+pub(crate) fn heap_base() -> *const u8 {
+    unsafe { rel_ptr_mut(HEAP_BASE) }
+}
+
+/// Returns the size of the heap
+pub(crate) fn heap_size() -> usize {
+    unsafe { HEAP_SIZE }
 }
 
 // Do not remove inline: will result in relocation failure

library/std/src/sys/sgx/alloc.rs

@@ -1,4 +1,7 @@
 use crate::alloc::{GlobalAlloc, Layout, System};
+use crate::ptr;
+use crate::sys::sgx::abi::mem as sgx_mem;
+use core::sync::atomic::{AtomicBool, Ordering};
 
 use super::waitqueue::SpinMutex;
 
@@ -10,7 +13,48 @@ use super::waitqueue::SpinMutex;
 // dlmalloc.c from C to Rust.
 #[cfg_attr(test, linkage = "available_externally")]
 #[export_name = "_ZN16__rust_internals3std3sys3sgx5alloc8DLMALLOCE"]
-static DLMALLOC: SpinMutex<dlmalloc::Dlmalloc> = SpinMutex::new(dlmalloc::DLMALLOC_INIT);
+static DLMALLOC: SpinMutex<dlmalloc::Dlmalloc<Sgx>> =
+    SpinMutex::new(dlmalloc::Dlmalloc::new_with_allocator(Sgx {}));
+
+struct Sgx;
+
+unsafe impl dlmalloc::Allocator for Sgx {
+    /// Allocs system resources
+    fn alloc(&self, _size: usize) -> (*mut u8, usize, u32) {
+        static INIT: AtomicBool = AtomicBool::new(false);
+
+        // No ordering requirement since this function is protected by the global lock.
+        if !INIT.swap(true, Ordering::Relaxed) {
+            (sgx_mem::heap_base() as _, sgx_mem::heap_size(), 0)
+        } else {
+            (ptr::null_mut(), 0, 0)
+        }
+    }
+
+    fn remap(&self, _ptr: *mut u8, _oldsize: usize, _newsize: usize, _can_move: bool) -> *mut u8 {
+        ptr::null_mut()
+    }
+
+    fn free_part(&self, _ptr: *mut u8, _oldsize: usize, _newsize: usize) -> bool {
+        false
+    }
+
+    fn free(&self, _ptr: *mut u8, _size: usize) -> bool {
+        return false;
+    }
+
+    fn can_release_part(&self, _flags: u32) -> bool {
+        false
+    }
+
+    fn allocates_zeros(&self) -> bool {
+        false
+    }
+
+    fn page_size(&self) -> usize {
+        0x1000
+    }
+}
 
 #[stable(feature = "alloc_system_type", since = "1.28.0")]
 unsafe impl GlobalAlloc for System {