From 39c02c7334300dc43140913eb78486a2b12faa61 Mon Sep 17 00:00:00 2001
From: aditijannu <aditi.jannu@fortanix.com>
Date: Tue, 7 May 2024 15:02:14 +0200
Subject: [PATCH 01/12] Initial updates for sgx allocator

---
 Cargo.lock                             | 41 ++++++++----
 library/std/Cargo.toml                 |  5 +-
 library/std/src/sys/pal/sgx/abi/mem.rs | 12 ----
 library/std/src/sys/pal/sgx/abi/mod.rs | 18 +++++
 library/std/src/sys/pal/sgx/alloc.rs   | 93 +++++++-------------------
 5 files changed, 75 insertions(+), 94 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d5c9a5e4d0d0e..7421a81369600 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -658,6 +658,15 @@ dependencies = [
  "rustc-semver",
 ]
 
+[[package]]
+name = "cmake"
+version = "0.1.50"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a31c789563b815f77f4250caee12365734369f942439b7defd71e18a48197130"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "collect-license-metadata"
 version = "0.1.0"
@@ -1147,19 +1156,6 @@ version = "1.0.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "86e3bdc80eee6e16b2b6b0f87fbc98c04bee3455e35174c0de1a125d0688c632"
 
-[[package]]
-name = "dlmalloc"
-version = "0.2.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3264b043b8e977326c1ee9e723da2c1f8d09a99df52cacf00b4dbce5ac54414d"
-dependencies = [
- "cfg-if",
- "compiler_builtins",
- "libc",
- "rustc-std-workspace-core",
- "windows-sys 0.52.0",
-]
-
 [[package]]
 name = "either"
 version = "1.10.0"
@@ -1178,6 +1174,12 @@ dependencies = [
  "serde_json",
 ]
 
+[[package]]
+name = "elf"
+version = "0.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4445909572dbd556c457c849c4ca58623d84b27c8fff1e74b0b4227d8b90d17b"
+
 [[package]]
 name = "elsa"
 version = "1.7.1"
@@ -5141,6 +5143,17 @@ version = "1.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b"
 
+[[package]]
+name = "snmalloc-edp"
+version = "0.1.0"
+dependencies = [
+ "cc",
+ "cmake",
+ "compiler_builtins",
+ "elf",
+ "rustc-std-workspace-core",
+]
+
 [[package]]
 name = "socket2"
 version = "0.5.6"
@@ -5232,7 +5245,6 @@ dependencies = [
  "cfg-if",
  "compiler_builtins",
  "core",
- "dlmalloc",
  "fortanix-sgx-abi",
  "hashbrown",
  "hermit-abi",
@@ -5247,6 +5259,7 @@ dependencies = [
  "rand",
  "rand_xorshift",
  "rustc-demangle",
+ "snmalloc-edp",
  "std_detect",
  "unwind",
  "wasi",
diff --git a/library/std/Cargo.toml b/library/std/Cargo.toml
index d2804b4d20af3..6456bc0330d1d 100644
--- a/library/std/Cargo.toml
+++ b/library/std/Cargo.toml
@@ -42,7 +42,7 @@ rand = { version = "0.8.5", default-features = false, features = ["alloc"] }
 rand_xorshift = "0.3.0"
 
 [target.'cfg(any(all(target_family = "wasm", target_os = "unknown"), target_os = "xous", all(target_vendor = "fortanix", target_env = "sgx")))'.dependencies]
-dlmalloc = { version = "0.2.4", features = ['rustc-dep-of-std'] }
+snmalloc-edp = { path = "/home/aditi/work/edp/rust-sgx/snmalloc-edp", features = ['rustc-dep-of-std'], public = true }
 
 [target.x86_64-fortanix-unknown-sgx.dependencies]
 fortanix-sgx-abi = { version = "0.5.0", features = ['rustc-dep-of-std'], public = true }
@@ -95,3 +95,6 @@ heap_size = 0x8000000
 name = "stdbenches"
 path = "benches/lib.rs"
 test = true
+
+[profile.dev]
+panic = "abort"
diff --git a/library/std/src/sys/pal/sgx/abi/mem.rs b/library/std/src/sys/pal/sgx/abi/mem.rs
index 18e6d5b3fa24d..80ead7d29d47c 100644
--- a/library/std/src/sys/pal/sgx/abi/mem.rs
+++ b/library/std/src/sys/pal/sgx/abi/mem.rs
@@ -14,18 +14,6 @@ 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
diff --git a/library/std/src/sys/pal/sgx/abi/mod.rs b/library/std/src/sys/pal/sgx/abi/mod.rs
index 9508c38741551..feea28976d0eb 100644
--- a/library/std/src/sys/pal/sgx/abi/mod.rs
+++ b/library/std/src/sys/pal/sgx/abi/mod.rs
@@ -61,6 +61,14 @@ unsafe extern "C" fn tcs_init(secondary: bool) {
 #[cfg(not(test))]
 #[no_mangle]
 extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64) -> EntryReturn {
+
+    let mut allocator = core::mem::MaybeUninit::<snmalloc_edp::Alloc>::uninit();
+    unsafe {
+        // Initialize thread local allocator
+        snmalloc_edp::sn_thread_init(allocator.as_mut_ptr());
+        crate::sys::alloc::THREAD_ALLOC.set(allocator.as_mut_ptr());
+    }
+
     // FIXME: how to support TLS in library mode?
     let tls = Box::new(tls::Tls::new());
     let tls_guard = unsafe { tls.activate() };
@@ -70,6 +78,8 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
         drop(tls_guard);
         drop(join_notifier);
 
+        alloc_thread_cleanup(allocator.as_mut_ptr());
+
         EntryReturn(0, 0)
     } else {
         extern "C" {
@@ -87,11 +97,19 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
             // main function, so we pass these in as the standard pointer-sized
             // values in `argc` and `argv`.
             let ret = main(p2 as _, p1 as _);
+            if ret == 0 {
+                alloc_thread_cleanup(allocator.as_mut_ptr());
+            }
             exit_with_code(ret)
         }
     }
 }
 
+pub(super) fn alloc_thread_cleanup(allocator: *mut snmalloc_edp::Alloc) {
+    crate::sys::alloc::THREAD_ALLOC.set(core::ptr::null_mut());
+    unsafe { snmalloc_edp::sn_thread_cleanup(allocator); }
+}
+
 pub(super) fn exit_with_code(code: isize) -> ! {
     if code != 0 {
         if let Some(mut out) = panic::SgxPanicOutput::new() {
diff --git a/library/std/src/sys/pal/sgx/alloc.rs b/library/std/src/sys/pal/sgx/alloc.rs
index 0c7bf9a9201f2..158934e1d6070 100644
--- a/library/std/src/sys/pal/sgx/alloc.rs
+++ b/library/std/src/sys/pal/sgx/alloc.rs
@@ -1,85 +1,34 @@
-use crate::alloc::{GlobalAlloc, Layout, System};
-use crate::ptr;
+use crate::{alloc::{self, GlobalAlloc, System}, cell::Cell, ptr};
+use snmalloc_edp::*;
 use core::sync::atomic::{AtomicBool, Ordering};
 
-use super::abi::mem as sgx_mem;
-use super::waitqueue::SpinMutex;
-
-// Using a SpinMutex because we never want to exit the enclave waiting for the
-// allocator.
-//
-// The current allocator here is the `dlmalloc` crate which we've got included
-// in the rust-lang/rust repository as a submodule. The crate is a port of
-// dlmalloc.c from C to Rust.
-#[cfg_attr(test, linkage = "available_externally")]
-#[export_name = "_ZN16__rust_internals3std3sys3sgx5alloc8DLMALLOCE"]
-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 {
     #[inline]
-    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        // SAFETY: the caller must uphold the safety contract for `malloc`
-        unsafe { DLMALLOC.lock().malloc(layout.size(), layout.align()) }
+    unsafe fn alloc(&self, layout: alloc::Layout) -> *mut u8 {
+        static INIT: AtomicBool = AtomicBool::new(false);
+        if !INIT.swap(true, Ordering::Relaxed) {
+            unsafe { sn_global_init(); }
+        }
+        unsafe { sn_rust_alloc(layout.align(), layout.size()) }
     }
 
     #[inline]
-    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+    unsafe fn alloc_zeroed(&self, layout: alloc::Layout) -> *mut u8 {
         // SAFETY: the caller must uphold the safety contract for `malloc`
-        unsafe { DLMALLOC.lock().calloc(layout.size(), layout.align()) }
+        unsafe { sn_rust_alloc_zeroed(layout.align(), layout.size()) }
     }
 
     #[inline]
-    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: alloc::Layout) {
         // SAFETY: the caller must uphold the safety contract for `malloc`
-        unsafe { DLMALLOC.lock().free(ptr, layout.size(), layout.align()) }
+        unsafe { sn_rust_dealloc(ptr, layout.align(), layout.size()) }
     }
 
     #[inline]
-    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+    unsafe fn realloc(&self, ptr: *mut u8, layout: alloc::Layout, new_size: usize) -> *mut u8 {
         // SAFETY: the caller must uphold the safety contract for `malloc`
-        unsafe { DLMALLOC.lock().realloc(ptr, layout.size(), layout.align(), new_size) }
+        unsafe { sn_rust_realloc(ptr, layout.align(), layout.size(), new_size) }
     }
 }
 
@@ -88,11 +37,21 @@ unsafe impl GlobalAlloc for System {
 #[cfg(not(test))]
 #[no_mangle]
 pub unsafe extern "C" fn __rust_c_alloc(size: usize, align: usize) -> *mut u8 {
-    unsafe { crate::alloc::alloc(Layout::from_size_align_unchecked(size, align)) }
+    unsafe { crate::alloc::alloc(crate::alloc::Layout::from_size_align_unchecked(size, align)) }
 }
 
 #[cfg(not(test))]
 #[no_mangle]
 pub unsafe extern "C" fn __rust_c_dealloc(ptr: *mut u8, size: usize, align: usize) {
-    unsafe { crate::alloc::dealloc(ptr, Layout::from_size_align_unchecked(size, align)) }
+    unsafe { crate::alloc::dealloc(ptr, crate::alloc::Layout::from_size_align_unchecked(size, align)) }
+}
+
+thread_local! {
+    pub static THREAD_ALLOC: Cell<*mut Alloc> = const { Cell::new(ptr::null_mut()) };
+}
+
+#[cfg(not(test))]
+#[no_mangle]
+pub extern "C" fn __rust_get_thread_allocator() -> *mut Alloc {
+    THREAD_ALLOC.get()
 }

From f6667335b3b9cc794d2db181238edb0ed8a068a2 Mon Sep 17 00:00:00 2001
From: aditijannu <aditi.jannu@fortanix.com>
Date: Tue, 14 May 2024 15:20:46 +0200
Subject: [PATCH 02/12] update call to sn_global_init

---
 library/std/src/sys/pal/sgx/abi/mod.rs | 9 ++++++++-
 library/std/src/sys/pal/sgx/alloc.rs   | 5 -----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/library/std/src/sys/pal/sgx/abi/mod.rs b/library/std/src/sys/pal/sgx/abi/mod.rs
index feea28976d0eb..ec517f6c98dfd 100644
--- a/library/std/src/sys/pal/sgx/abi/mod.rs
+++ b/library/std/src/sys/pal/sgx/abi/mod.rs
@@ -2,7 +2,7 @@
 
 use crate::io::Write;
 use core::arch::global_asm;
-use core::sync::atomic::{AtomicUsize, Ordering};
+use core::sync::atomic::{AtomicUsize, Ordering, AtomicBool};
 
 // runtime features
 pub(super) mod panic;
@@ -24,6 +24,7 @@ struct EntryReturn(u64, u64);
 #[cfg(not(test))]
 #[no_mangle]
 unsafe extern "C" fn tcs_init(secondary: bool) {
+
     // Be very careful when changing this code: it runs before the binary has been
     // relocated. Any indirect accesses to symbols will likely fail.
     const UNINIT: usize = 0;
@@ -41,6 +42,12 @@ unsafe extern "C" fn tcs_init(secondary: bool) {
         // This thread just obtained the lock and other threads will observe BUSY
         Ok(_) => {
             reloc::relocate_elf_rela();
+
+            static INIT: AtomicBool = AtomicBool::new(false);
+            if !INIT.swap(true, Ordering::Relaxed) {
+                unsafe { snmalloc_edp::sn_global_init(); }
+            }
+
             RELOC_STATE.store(DONE, Ordering::Release);
         }
         // We need to wait until the initialization is done.
diff --git a/library/std/src/sys/pal/sgx/alloc.rs b/library/std/src/sys/pal/sgx/alloc.rs
index 158934e1d6070..c1f546516b6de 100644
--- a/library/std/src/sys/pal/sgx/alloc.rs
+++ b/library/std/src/sys/pal/sgx/alloc.rs
@@ -1,15 +1,10 @@
 use crate::{alloc::{self, GlobalAlloc, System}, cell::Cell, ptr};
 use snmalloc_edp::*;
-use core::sync::atomic::{AtomicBool, Ordering};
 
 #[stable(feature = "alloc_system_type", since = "1.28.0")]
 unsafe impl GlobalAlloc for System {
     #[inline]
     unsafe fn alloc(&self, layout: alloc::Layout) -> *mut u8 {
-        static INIT: AtomicBool = AtomicBool::new(false);
-        if !INIT.swap(true, Ordering::Relaxed) {
-            unsafe { sn_global_init(); }
-        }
         unsafe { sn_rust_alloc(layout.align(), layout.size()) }
     }
 

From 94850c4f406878c1df06a0193198e5cf55f533d5 Mon Sep 17 00:00:00 2001
From: aditijannu <aditi.jannu@fortanix.com>
Date: Tue, 14 May 2024 16:49:54 +0200
Subject: [PATCH 03/12] Update std's snmalloc dependency to point to git

---
 Cargo.lock             | 1 +
 library/std/Cargo.toml | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7421a81369600..201c9f739e8a4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5146,6 +5146,7 @@ checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b"
 [[package]]
 name = "snmalloc-edp"
 version = "0.1.0"
+source = "git+https://github.com/fortanix/rust-sgx?branch=aj/update-sgx-alloc#ac0c2a16324f2a74b5c2d9b6dbf204fb3c89b0f7"
 dependencies = [
  "cc",
  "cmake",
diff --git a/library/std/Cargo.toml b/library/std/Cargo.toml
index 6456bc0330d1d..4b3b9fd33f514 100644
--- a/library/std/Cargo.toml
+++ b/library/std/Cargo.toml
@@ -42,7 +42,7 @@ rand = { version = "0.8.5", default-features = false, features = ["alloc"] }
 rand_xorshift = "0.3.0"
 
 [target.'cfg(any(all(target_family = "wasm", target_os = "unknown"), target_os = "xous", all(target_vendor = "fortanix", target_env = "sgx")))'.dependencies]
-snmalloc-edp = { path = "/home/aditi/work/edp/rust-sgx/snmalloc-edp", features = ['rustc-dep-of-std'], public = true }
+snmalloc-edp = { git = "https://github.com/fortanix/rust-sgx", branch = "aj/update-sgx-alloc", features = ['rustc-dep-of-std'], public = true }
 
 [target.x86_64-fortanix-unknown-sgx.dependencies]
 fortanix-sgx-abi = { version = "0.5.0", features = ['rustc-dep-of-std'], public = true }

From 95cce62a2fd27e2e77ff195526365c6b5ac46c53 Mon Sep 17 00:00:00 2001
From: aditijannu <aditi.jannu@fortanix.com>
Date: Thu, 16 May 2024 11:07:57 +0200
Subject: [PATCH 04/12] Fix global init of alloc due to incorrect heap
 size/base params

---
 library/std/src/sys/pal/sgx/abi/mem.rs | 12 ++++++++++++
 library/std/src/sys/pal/sgx/abi/mod.rs | 15 ++++++---------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/library/std/src/sys/pal/sgx/abi/mem.rs b/library/std/src/sys/pal/sgx/abi/mem.rs
index 80ead7d29d47c..18e6d5b3fa24d 100644
--- a/library/std/src/sys/pal/sgx/abi/mem.rs
+++ b/library/std/src/sys/pal/sgx/abi/mem.rs
@@ -14,6 +14,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
diff --git a/library/std/src/sys/pal/sgx/abi/mod.rs b/library/std/src/sys/pal/sgx/abi/mod.rs
index ec517f6c98dfd..41ae49e912850 100644
--- a/library/std/src/sys/pal/sgx/abi/mod.rs
+++ b/library/std/src/sys/pal/sgx/abi/mod.rs
@@ -3,7 +3,6 @@
 use crate::io::Write;
 use core::arch::global_asm;
 use core::sync::atomic::{AtomicUsize, Ordering, AtomicBool};
-
 // runtime features
 pub(super) mod panic;
 mod reloc;
@@ -18,6 +17,8 @@ pub mod usercalls;
 #[cfg(not(test))]
 global_asm!(include_str!("entry.S"), options(att_syntax));
 
+static INIT: AtomicBool = AtomicBool::new(false);
+
 #[repr(C)]
 struct EntryReturn(u64, u64);
 
@@ -42,10 +43,8 @@ unsafe extern "C" fn tcs_init(secondary: bool) {
         // This thread just obtained the lock and other threads will observe BUSY
         Ok(_) => {
             reloc::relocate_elf_rela();
-
-            static INIT: AtomicBool = AtomicBool::new(false);
             if !INIT.swap(true, Ordering::Relaxed) {
-                unsafe { snmalloc_edp::sn_global_init(); }
+                unsafe { snmalloc_edp::sn_global_init(mem::heap_base(), mem::heap_size()); }
             }
 
             RELOC_STATE.store(DONE, Ordering::Release);
@@ -69,12 +68,10 @@ unsafe extern "C" fn tcs_init(secondary: bool) {
 #[no_mangle]
 extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64) -> EntryReturn {
 
+    // Initialize thread local allocator
     let mut allocator = core::mem::MaybeUninit::<snmalloc_edp::Alloc>::uninit();
-    unsafe {
-        // Initialize thread local allocator
-        snmalloc_edp::sn_thread_init(allocator.as_mut_ptr());
-        crate::sys::alloc::THREAD_ALLOC.set(allocator.as_mut_ptr());
-    }
+    unsafe { snmalloc_edp::sn_thread_init(allocator.as_mut_ptr()); }
+    crate::sys::alloc::THREAD_ALLOC.set(allocator.as_mut_ptr());
 
     // FIXME: how to support TLS in library mode?
     let tls = Box::new(tls::Tls::new());

From 2ee0d1483a0524f7aecdd1652287a92f6c75d16d Mon Sep 17 00:00:00 2001
From: aditijannu <aditi.jannu@fortanix.com>
Date: Thu, 16 May 2024 15:01:02 +0200
Subject: [PATCH 05/12] Store thread allocator in TCSLS

---
 library/std/Cargo.toml                     |  3 ---
 library/std/src/sys/pal/sgx/abi/entry.S    | 10 +++++-----
 library/std/src/sys/pal/sgx/abi/mod.rs     | 22 +++++++++++++++++-----
 library/std/src/sys/pal/sgx/abi/tls/mod.rs | 18 +++++++++++++-----
 library/std/src/sys/pal/sgx/alloc.rs       |  8 ++------
 5 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/library/std/Cargo.toml b/library/std/Cargo.toml
index 4b3b9fd33f514..d45b86a017acd 100644
--- a/library/std/Cargo.toml
+++ b/library/std/Cargo.toml
@@ -95,6 +95,3 @@ heap_size = 0x8000000
 name = "stdbenches"
 path = "benches/lib.rs"
 test = true
-
-[profile.dev]
-panic = "abort"
diff --git a/library/std/src/sys/pal/sgx/abi/entry.S b/library/std/src/sys/pal/sgx/abi/entry.S
index 8a063b65dac50..223661e3577ca 100644
--- a/library/std/src/sys/pal/sgx/abi/entry.S
+++ b/library/std/src/sys/pal/sgx/abi/entry.S
@@ -97,8 +97,8 @@ IMAGE_BASE:
 .equ tcsls_user_r13,            0x48
 .equ tcsls_user_r14,            0x50
 .equ tcsls_user_r15,            0x58
-.equ tcsls_tls_ptr,             0x60
-.equ tcsls_tcs_addr,            0x68
+.equ tcsls_tcs_addr,            0x60
+.equ tcsls_tls_ptr,             0x68
 
 .macro load_tcsls_flag_secondary_bool reg:req comments:vararg
     .ifne tcsls_flag_secondary /* to convert to a bool, must be the first bit */
@@ -131,7 +131,7 @@ elf_entry:
     sub %rax,%rdx                     /* all chars written?     */
     jnz .Lelf_entry_call
 
-.Lelf_exit:    
+.Lelf_exit:
     movq $60,%rax                     /* exit() syscall         */
     movq $1,%rdi                      /* exit code 1            */
     syscall
@@ -355,14 +355,14 @@ get_tcs_addr:
 
 .global get_tls_ptr
 get_tls_ptr:
-    mov %gs:tcsls_tls_ptr,%rax
+    mov %gs:tcsls_tls_ptr(,%rdi,8),%rax
     pop %r11
     lfence
     jmp *%r11
 
 .global set_tls_ptr
 set_tls_ptr:
-    mov %rdi,%gs:tcsls_tls_ptr
+    mov %rsi,%gs:tcsls_tls_ptr(,%rdi,8)
     pop %r11
     lfence
     jmp *%r11
diff --git a/library/std/src/sys/pal/sgx/abi/mod.rs b/library/std/src/sys/pal/sgx/abi/mod.rs
index 41ae49e912850..6faef1abde6c3 100644
--- a/library/std/src/sys/pal/sgx/abi/mod.rs
+++ b/library/std/src/sys/pal/sgx/abi/mod.rs
@@ -3,6 +3,8 @@
 use crate::io::Write;
 use core::arch::global_asm;
 use core::sync::atomic::{AtomicUsize, Ordering, AtomicBool};
+use snmalloc_edp::*;
+
 // runtime features
 pub(super) mod panic;
 mod reloc;
@@ -17,6 +19,7 @@ pub mod usercalls;
 #[cfg(not(test))]
 global_asm!(include_str!("entry.S"), options(att_syntax));
 
+// To initialize global allocator once per program
 static INIT: AtomicBool = AtomicBool::new(false);
 
 #[repr(C)]
@@ -43,8 +46,10 @@ unsafe extern "C" fn tcs_init(secondary: bool) {
         // This thread just obtained the lock and other threads will observe BUSY
         Ok(_) => {
             reloc::relocate_elf_rela();
+
+            // Snmalloc global allocator initialization
             if !INIT.swap(true, Ordering::Relaxed) {
-                unsafe { snmalloc_edp::sn_global_init(mem::heap_base(), mem::heap_size()); }
+                unsafe { sn_global_init(mem::heap_base(), mem::heap_size()); }
             }
 
             RELOC_STATE.store(DONE, Ordering::Release);
@@ -70,8 +75,10 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
 
     // Initialize thread local allocator
     let mut allocator = core::mem::MaybeUninit::<snmalloc_edp::Alloc>::uninit();
-    unsafe { snmalloc_edp::sn_thread_init(allocator.as_mut_ptr()); }
-    crate::sys::alloc::THREAD_ALLOC.set(allocator.as_mut_ptr());
+    unsafe {
+        sn_thread_init(allocator.as_mut_ptr());
+        tls::set_tls_ptr(tls::TlsIndex::AllocPtr, allocator.as_mut_ptr() as *const u8);
+    }
 
     // FIXME: how to support TLS in library mode?
     let tls = Box::new(tls::Tls::new());
@@ -81,6 +88,7 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
         let join_notifier = super::thread::Thread::entry();
         drop(tls_guard);
         drop(join_notifier);
+        drop(tls);
 
         alloc_thread_cleanup(allocator.as_mut_ptr());
 
@@ -102,6 +110,8 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
             // values in `argc` and `argv`.
             let ret = main(p2 as _, p1 as _);
             if ret == 0 {
+                drop(tls_guard);
+                drop(tls);
                 alloc_thread_cleanup(allocator.as_mut_ptr());
             }
             exit_with_code(ret)
@@ -110,8 +120,10 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
 }
 
 pub(super) fn alloc_thread_cleanup(allocator: *mut snmalloc_edp::Alloc) {
-    crate::sys::alloc::THREAD_ALLOC.set(core::ptr::null_mut());
-    unsafe { snmalloc_edp::sn_thread_cleanup(allocator); }
+    unsafe {
+        tls::set_tls_ptr(tls::TlsIndex::AllocPtr, core::ptr::null_mut());
+        sn_thread_cleanup(allocator);
+    }
 }
 
 pub(super) fn exit_with_code(code: isize) -> ! {
diff --git a/library/std/src/sys/pal/sgx/abi/tls/mod.rs b/library/std/src/sys/pal/sgx/abi/tls/mod.rs
index 8a9ea4ac00df0..de47e646628dc 100644
--- a/library/std/src/sys/pal/sgx/abi/tls/mod.rs
+++ b/library/std/src/sys/pal/sgx/abi/tls/mod.rs
@@ -23,9 +23,17 @@ macro_rules! dup {
 #[export_name = "_ZN16__rust_internals3std3sys3sgx3abi3tls14TLS_DESTRUCTORE"]
 static TLS_DESTRUCTOR: [AtomicUsize; TLS_KEYS] = dup!((* * * * * * *) (AtomicUsize::new(0)));
 
+// Use TCSLS to store both the thread specific allocator
+// and TLS address at different indices
+#[repr(u8)]
+pub enum TlsIndex {
+    AllocPtr = 0,
+    TlsPtr = 1,
+}
+
 extern "C" {
-    fn get_tls_ptr() -> *const u8;
-    fn set_tls_ptr(tls: *const u8);
+    pub fn get_tls_ptr(index: TlsIndex) -> *const u8;
+    pub fn set_tls_ptr(index: TlsIndex, tls: *const u8);
 }
 
 #[derive(Copy, Clone)]
@@ -88,20 +96,20 @@ impl Tls {
 
     pub unsafe fn activate(&self) -> ActiveTls<'_> {
         // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        unsafe { set_tls_ptr(self as *const Tls as _) };
+        unsafe { set_tls_ptr(TlsIndex::TlsPtr, self as *const Tls as _) };
         ActiveTls { tls: self }
     }
 
     #[allow(unused)]
     pub unsafe fn activate_persistent(self: Box<Self>) {
         // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        unsafe { set_tls_ptr(core::ptr::addr_of!(*self) as _) };
+        unsafe { set_tls_ptr(TlsIndex::TlsPtr, core::ptr::addr_of!(*self) as _) };
         mem::forget(self);
     }
 
     unsafe fn current<'a>() -> &'a Tls {
         // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        unsafe { &*(get_tls_ptr() as *const Tls) }
+        unsafe { &*(get_tls_ptr(TlsIndex::TlsPtr) as *const Tls) }
     }
 
     pub fn create(dtor: Option<unsafe extern "C" fn(*mut u8)>) -> Key {
diff --git a/library/std/src/sys/pal/sgx/alloc.rs b/library/std/src/sys/pal/sgx/alloc.rs
index c1f546516b6de..ba6c16b260a9a 100644
--- a/library/std/src/sys/pal/sgx/alloc.rs
+++ b/library/std/src/sys/pal/sgx/alloc.rs
@@ -1,4 +1,4 @@
-use crate::{alloc::{self, GlobalAlloc, System}, cell::Cell, ptr};
+use crate::{alloc::{self, GlobalAlloc, System}};
 use snmalloc_edp::*;
 
 #[stable(feature = "alloc_system_type", since = "1.28.0")]
@@ -41,12 +41,8 @@ pub unsafe extern "C" fn __rust_c_dealloc(ptr: *mut u8, size: usize, align: usiz
     unsafe { crate::alloc::dealloc(ptr, crate::alloc::Layout::from_size_align_unchecked(size, align)) }
 }
 
-thread_local! {
-    pub static THREAD_ALLOC: Cell<*mut Alloc> = const { Cell::new(ptr::null_mut()) };
-}
-
 #[cfg(not(test))]
 #[no_mangle]
 pub extern "C" fn __rust_get_thread_allocator() -> *mut Alloc {
-    THREAD_ALLOC.get()
+    unsafe{ crate::sys::abi::tls::get_tls_ptr(crate::sys::abi::tls::TlsIndex::AllocPtr) as *mut Alloc }
 }

From df4ce89c3cffc1f90917d28044d6f29dbf53848e Mon Sep 17 00:00:00 2001
From: Yuxiang Cao <yuxiang.cao@fortanix.com>
Date: Tue, 28 May 2024 16:49:56 -0700
Subject: [PATCH 06/12] fix compile

---
 library/std/src/sys/pal/sgx/abi/mem.rs |  2 ++
 library/std/src/sys/pal/sgx/abi/mod.rs |  2 +-
 mybuild.sh                             | 43 ++++++++++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100755 mybuild.sh

diff --git a/library/std/src/sys/pal/sgx/abi/mem.rs b/library/std/src/sys/pal/sgx/abi/mem.rs
index 18e6d5b3fa24d..cd3a064d64ebc 100644
--- a/library/std/src/sys/pal/sgx/abi/mem.rs
+++ b/library/std/src/sys/pal/sgx/abi/mem.rs
@@ -19,11 +19,13 @@ extern "C" {
 }
 
 /// Returns the base memory address of the heap
+#[allow(dead_code)]
 pub(crate) fn heap_base() -> *const u8 {
     unsafe { rel_ptr_mut(HEAP_BASE) }
 }
 
 /// Returns the size of the heap
+#[allow(dead_code)]
 pub(crate) fn heap_size() -> usize {
     unsafe { HEAP_SIZE }
 }
diff --git a/library/std/src/sys/pal/sgx/abi/mod.rs b/library/std/src/sys/pal/sgx/abi/mod.rs
index 6faef1abde6c3..a7162e0a7454e 100644
--- a/library/std/src/sys/pal/sgx/abi/mod.rs
+++ b/library/std/src/sys/pal/sgx/abi/mod.rs
@@ -49,7 +49,7 @@ unsafe extern "C" fn tcs_init(secondary: bool) {
 
             // Snmalloc global allocator initialization
             if !INIT.swap(true, Ordering::Relaxed) {
-                unsafe { sn_global_init(mem::heap_base(), mem::heap_size()); }
+                unsafe { sn_global_init(); }
             }
 
             RELOC_STATE.store(DONE, Ordering::Release);
diff --git a/mybuild.sh b/mybuild.sh
new file mode 100755
index 0000000000000..d42183f621085
--- /dev/null
+++ b/mybuild.sh
@@ -0,0 +1,43 @@
+#!/bin/bash -ex
+export AR_x86_64_fortanix_unknown_sgx=ar
+export CC_x86_64_fortanix_unknown_sgx=clang-12
+export CFLAGS_x86_64_fortanix_unknown_sgx="-D__ELF__ -isystem/usr/include/x86_64-linux-gnu -mlvi-hardening -mllvm -x86-experimental-lvi-inline-asm-hardening"
+export CXX_x86_64_fortanix_unknown_sgx=clang++-12
+export CXXFLAGS_x86_64_fortanix_unknown_sgx="-D__ELF__ -isystem/usr/include/x86_64-linux-gnu -mlvi-hardening -mllvm -x86-experimental-lvi-inline-asm-hardening"
+export CC_x86_64_unknown_linux_gnu=clang-12
+export CXX_x86_64_unknown_linux_gnu=clang++-12
+
+git submodule update --init --recursive
+        detect_cxx_include_path() {
+            for path in $(clang++-12 -print-search-dirs|sed -n 's/^libraries:\s*=//p'|tr : ' '); do
+                num_component="$(basename "$path")"
+                if [[ "$num_component" =~ ^[0-9]+(\.[0-9]+)*$ ]]; then
+                    if [[ "$(basename "$(dirname "$path")")" == 'x86_64-linux-gnu' ]]; then
+                        echo $num_component
+                        return
+                    fi
+                fi
+            done
+            exit 1
+        }
+export CXXFLAGS_x86_64_fortanix_unknown_sgx="-cxx-isystem/usr/include/c++/$(detect_cxx_include_path) -cxx-isystem/usr/include/x86_64-linux-gnu/c++/$(detect_cxx_include_path) $CFLAGS_x86_64_fortanix_unknown_sgx"
+
+if [ ! -f config.toml ]; then
+    rustup default nightly
+    rustup update nightly
+    rustup target add x86_64-fortanix-unknown-sgx
+
+    cargo install fortanix-sgx-tools
+    cargo install sgxs-tools
+
+    git submodule foreach git reset --hard
+    ./configure --enable-lld --set profile=compiler \
+        --set install.prefix="/home/yuxiangcao/main_ws/toolchain/rust/rustup/toolchains/dev_snmalloc-2024-04-03-x86_64-unknown-linux-gnu" \
+        --set install.sysconfdir="etc" \
+        --set build.extended=true \
+        --set build.tools="cargo,rustfmt,src" \
+        --set build.docs=false
+
+fi
+
+./x install --target=x86_64-fortanix-unknown-sgx,x86_64-unknown-linux-gnu --host=x86_64-unknown-linux-gnu

From 6947445d386b2110ab1b0b015d158440be6874f3 Mon Sep 17 00:00:00 2001
From: Yuxiang Cao <yuxiang.cao@fortanix.com>
Date: Mon, 3 Jun 2024 10:58:02 -0700
Subject: [PATCH 07/12] Revert "fix compile"

This reverts commit df4ce89c3cffc1f90917d28044d6f29dbf53848e.
---
 library/std/src/sys/pal/sgx/abi/mem.rs |  2 --
 library/std/src/sys/pal/sgx/abi/mod.rs |  2 +-
 mybuild.sh                             | 43 --------------------------
 3 files changed, 1 insertion(+), 46 deletions(-)
 delete mode 100755 mybuild.sh

diff --git a/library/std/src/sys/pal/sgx/abi/mem.rs b/library/std/src/sys/pal/sgx/abi/mem.rs
index cd3a064d64ebc..18e6d5b3fa24d 100644
--- a/library/std/src/sys/pal/sgx/abi/mem.rs
+++ b/library/std/src/sys/pal/sgx/abi/mem.rs
@@ -19,13 +19,11 @@ extern "C" {
 }
 
 /// Returns the base memory address of the heap
-#[allow(dead_code)]
 pub(crate) fn heap_base() -> *const u8 {
     unsafe { rel_ptr_mut(HEAP_BASE) }
 }
 
 /// Returns the size of the heap
-#[allow(dead_code)]
 pub(crate) fn heap_size() -> usize {
     unsafe { HEAP_SIZE }
 }
diff --git a/library/std/src/sys/pal/sgx/abi/mod.rs b/library/std/src/sys/pal/sgx/abi/mod.rs
index a7162e0a7454e..6faef1abde6c3 100644
--- a/library/std/src/sys/pal/sgx/abi/mod.rs
+++ b/library/std/src/sys/pal/sgx/abi/mod.rs
@@ -49,7 +49,7 @@ unsafe extern "C" fn tcs_init(secondary: bool) {
 
             // Snmalloc global allocator initialization
             if !INIT.swap(true, Ordering::Relaxed) {
-                unsafe { sn_global_init(); }
+                unsafe { sn_global_init(mem::heap_base(), mem::heap_size()); }
             }
 
             RELOC_STATE.store(DONE, Ordering::Release);
diff --git a/mybuild.sh b/mybuild.sh
deleted file mode 100755
index d42183f621085..0000000000000
--- a/mybuild.sh
+++ /dev/null
@@ -1,43 +0,0 @@
-#!/bin/bash -ex
-export AR_x86_64_fortanix_unknown_sgx=ar
-export CC_x86_64_fortanix_unknown_sgx=clang-12
-export CFLAGS_x86_64_fortanix_unknown_sgx="-D__ELF__ -isystem/usr/include/x86_64-linux-gnu -mlvi-hardening -mllvm -x86-experimental-lvi-inline-asm-hardening"
-export CXX_x86_64_fortanix_unknown_sgx=clang++-12
-export CXXFLAGS_x86_64_fortanix_unknown_sgx="-D__ELF__ -isystem/usr/include/x86_64-linux-gnu -mlvi-hardening -mllvm -x86-experimental-lvi-inline-asm-hardening"
-export CC_x86_64_unknown_linux_gnu=clang-12
-export CXX_x86_64_unknown_linux_gnu=clang++-12
-
-git submodule update --init --recursive
-        detect_cxx_include_path() {
-            for path in $(clang++-12 -print-search-dirs|sed -n 's/^libraries:\s*=//p'|tr : ' '); do
-                num_component="$(basename "$path")"
-                if [[ "$num_component" =~ ^[0-9]+(\.[0-9]+)*$ ]]; then
-                    if [[ "$(basename "$(dirname "$path")")" == 'x86_64-linux-gnu' ]]; then
-                        echo $num_component
-                        return
-                    fi
-                fi
-            done
-            exit 1
-        }
-export CXXFLAGS_x86_64_fortanix_unknown_sgx="-cxx-isystem/usr/include/c++/$(detect_cxx_include_path) -cxx-isystem/usr/include/x86_64-linux-gnu/c++/$(detect_cxx_include_path) $CFLAGS_x86_64_fortanix_unknown_sgx"
-
-if [ ! -f config.toml ]; then
-    rustup default nightly
-    rustup update nightly
-    rustup target add x86_64-fortanix-unknown-sgx
-
-    cargo install fortanix-sgx-tools
-    cargo install sgxs-tools
-
-    git submodule foreach git reset --hard
-    ./configure --enable-lld --set profile=compiler \
-        --set install.prefix="/home/yuxiangcao/main_ws/toolchain/rust/rustup/toolchains/dev_snmalloc-2024-04-03-x86_64-unknown-linux-gnu" \
-        --set install.sysconfdir="etc" \
-        --set build.extended=true \
-        --set build.tools="cargo,rustfmt,src" \
-        --set build.docs=false
-
-fi
-
-./x install --target=x86_64-fortanix-unknown-sgx,x86_64-unknown-linux-gnu --host=x86_64-unknown-linux-gnu

From 3df8a1273a9b06b9714c185c9a28c9048ff3b6ec Mon Sep 17 00:00:00 2001
From: Yuxiang Cao <yuxiang.cao@fortanix.com>
Date: Mon, 3 Jun 2024 10:59:55 -0700
Subject: [PATCH 08/12] update snmalloc-edp code

---
 Cargo.lock | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 201c9f739e8a4..6dac6c9b46ac3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5146,7 +5146,7 @@ checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b"
 [[package]]
 name = "snmalloc-edp"
 version = "0.1.0"
-source = "git+https://github.com/fortanix/rust-sgx?branch=aj/update-sgx-alloc#ac0c2a16324f2a74b5c2d9b6dbf204fb3c89b0f7"
+source = "git+https://github.com/fortanix/rust-sgx?branch=aj/update-sgx-alloc#6fc0b33a8c0fe2a36e9511b22c5cbc2df57df19e"
 dependencies = [
  "cc",
  "cmake",

From ff5ee76f3675c1b69e90195216a41f7e112437e3 Mon Sep 17 00:00:00 2001
From: Yuxiang Cao <yuxiang.cao@fortanix.com>
Date: Mon, 3 Jun 2024 11:17:03 -0700
Subject: [PATCH 09/12] add my build script

---
 mybuild.sh | 43 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)
 create mode 100755 mybuild.sh

diff --git a/mybuild.sh b/mybuild.sh
new file mode 100755
index 0000000000000..d42183f621085
--- /dev/null
+++ b/mybuild.sh
@@ -0,0 +1,43 @@
+#!/bin/bash -ex
+export AR_x86_64_fortanix_unknown_sgx=ar
+export CC_x86_64_fortanix_unknown_sgx=clang-12
+export CFLAGS_x86_64_fortanix_unknown_sgx="-D__ELF__ -isystem/usr/include/x86_64-linux-gnu -mlvi-hardening -mllvm -x86-experimental-lvi-inline-asm-hardening"
+export CXX_x86_64_fortanix_unknown_sgx=clang++-12
+export CXXFLAGS_x86_64_fortanix_unknown_sgx="-D__ELF__ -isystem/usr/include/x86_64-linux-gnu -mlvi-hardening -mllvm -x86-experimental-lvi-inline-asm-hardening"
+export CC_x86_64_unknown_linux_gnu=clang-12
+export CXX_x86_64_unknown_linux_gnu=clang++-12
+
+git submodule update --init --recursive
+        detect_cxx_include_path() {
+            for path in $(clang++-12 -print-search-dirs|sed -n 's/^libraries:\s*=//p'|tr : ' '); do
+                num_component="$(basename "$path")"
+                if [[ "$num_component" =~ ^[0-9]+(\.[0-9]+)*$ ]]; then
+                    if [[ "$(basename "$(dirname "$path")")" == 'x86_64-linux-gnu' ]]; then
+                        echo $num_component
+                        return
+                    fi
+                fi
+            done
+            exit 1
+        }
+export CXXFLAGS_x86_64_fortanix_unknown_sgx="-cxx-isystem/usr/include/c++/$(detect_cxx_include_path) -cxx-isystem/usr/include/x86_64-linux-gnu/c++/$(detect_cxx_include_path) $CFLAGS_x86_64_fortanix_unknown_sgx"
+
+if [ ! -f config.toml ]; then
+    rustup default nightly
+    rustup update nightly
+    rustup target add x86_64-fortanix-unknown-sgx
+
+    cargo install fortanix-sgx-tools
+    cargo install sgxs-tools
+
+    git submodule foreach git reset --hard
+    ./configure --enable-lld --set profile=compiler \
+        --set install.prefix="/home/yuxiangcao/main_ws/toolchain/rust/rustup/toolchains/dev_snmalloc-2024-04-03-x86_64-unknown-linux-gnu" \
+        --set install.sysconfdir="etc" \
+        --set build.extended=true \
+        --set build.tools="cargo,rustfmt,src" \
+        --set build.docs=false
+
+fi
+
+./x install --target=x86_64-fortanix-unknown-sgx,x86_64-unknown-linux-gnu --host=x86_64-unknown-linux-gnu

From f35b0aa0ea91cf0aa5a4c2a1661d79d426ca11ef Mon Sep 17 00:00:00 2001
From: aditijannu <aditi.jannu@fortanix.com>
Date: Tue, 25 Jun 2024 15:06:37 +0200
Subject: [PATCH 10/12] Review comments addressed

---
 library/std/src/sys/pal/sgx/abi/entry.S    | 18 +++++----
 library/std/src/sys/pal/sgx/abi/mod.rs     | 21 ++++-------
 library/std/src/sys/pal/sgx/abi/tls/mod.rs | 27 ++++++++------
 library/std/src/sys/pal/sgx/alloc.rs       |  2 +-
 mybuild.sh                                 | 43 ----------------------
 5 files changed, 35 insertions(+), 76 deletions(-)
 delete mode 100755 mybuild.sh

diff --git a/library/std/src/sys/pal/sgx/abi/entry.S b/library/std/src/sys/pal/sgx/abi/entry.S
index 223661e3577ca..c1da196d9b352 100644
--- a/library/std/src/sys/pal/sgx/abi/entry.S
+++ b/library/std/src/sys/pal/sgx/abi/entry.S
@@ -98,7 +98,7 @@ IMAGE_BASE:
 .equ tcsls_user_r14,            0x50
 .equ tcsls_user_r15,            0x58
 .equ tcsls_tcs_addr,            0x60
-.equ tcsls_tls_ptr,             0x68
+.equ tcsls_tls_data,            0x68
 
 .macro load_tcsls_flag_secondary_bool reg:req comments:vararg
     .ifne tcsls_flag_secondary /* to convert to a bool, must be the first bit */
@@ -353,16 +353,20 @@ get_tcs_addr:
     lfence
     jmp *%r11
 
-.global get_tls_ptr
-get_tls_ptr:
-    mov %gs:tcsls_tls_ptr(,%rdi,8),%rax
+/* Fetches the data available at an offset from tcsls_tls_data. Refer to TlsIndex enum */
+/* for details on the offsets that can be used as an input to this function. */
+.global get_tls_data
+get_tls_data:
+    mov %gs:tcsls_tls_data(,%rdi,8),%rax
     pop %r11
     lfence
     jmp *%r11
 
-.global set_tls_ptr
-set_tls_ptr:
-    mov %rsi,%gs:tcsls_tls_ptr(,%rdi,8)
+/* Sets input data at an offset from tcsls_tls_data. Refer to TlsIndex enum for details on */
+/* the offsets that can be used as an input to this function. */
+.global set_tls_data
+set_tls_data:
+    mov %rsi,%gs:tcsls_tls_data(,%rdi,8)
     pop %r11
     lfence
     jmp *%r11
diff --git a/library/std/src/sys/pal/sgx/abi/mod.rs b/library/std/src/sys/pal/sgx/abi/mod.rs
index 6faef1abde6c3..c54b7f7cddaf4 100644
--- a/library/std/src/sys/pal/sgx/abi/mod.rs
+++ b/library/std/src/sys/pal/sgx/abi/mod.rs
@@ -2,7 +2,7 @@
 
 use crate::io::Write;
 use core::arch::global_asm;
-use core::sync::atomic::{AtomicUsize, Ordering, AtomicBool};
+use core::sync::atomic::{AtomicUsize, Ordering};
 use snmalloc_edp::*;
 
 // runtime features
@@ -19,16 +19,12 @@ pub mod usercalls;
 #[cfg(not(test))]
 global_asm!(include_str!("entry.S"), options(att_syntax));
 
-// To initialize global allocator once per program
-static INIT: AtomicBool = AtomicBool::new(false);
-
 #[repr(C)]
 struct EntryReturn(u64, u64);
 
 #[cfg(not(test))]
 #[no_mangle]
 unsafe extern "C" fn tcs_init(secondary: bool) {
-
     // Be very careful when changing this code: it runs before the binary has been
     // relocated. Any indirect accesses to symbols will likely fail.
     const UNINIT: usize = 0;
@@ -48,9 +44,7 @@ unsafe extern "C" fn tcs_init(secondary: bool) {
             reloc::relocate_elf_rela();
 
             // Snmalloc global allocator initialization
-            if !INIT.swap(true, Ordering::Relaxed) {
-                unsafe { sn_global_init(mem::heap_base(), mem::heap_size()); }
-            }
+            unsafe { sn_global_init(mem::heap_base(), mem::heap_size()); }
 
             RELOC_STATE.store(DONE, Ordering::Release);
         }
@@ -77,7 +71,7 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
     let mut allocator = core::mem::MaybeUninit::<snmalloc_edp::Alloc>::uninit();
     unsafe {
         sn_thread_init(allocator.as_mut_ptr());
-        tls::set_tls_ptr(tls::TlsIndex::AllocPtr, allocator.as_mut_ptr() as *const u8);
+        tls::set_tls_data(tls::TlsIndex::AllocPtr, allocator.as_mut_ptr() as *const u8);
     }
 
     // FIXME: how to support TLS in library mode?
@@ -90,7 +84,7 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
         drop(join_notifier);
         drop(tls);
 
-        alloc_thread_cleanup(allocator.as_mut_ptr());
+        alloc_thread_cleanup();
 
         EntryReturn(0, 0)
     } else {
@@ -112,16 +106,17 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
             if ret == 0 {
                 drop(tls_guard);
                 drop(tls);
-                alloc_thread_cleanup(allocator.as_mut_ptr());
+                alloc_thread_cleanup();
             }
             exit_with_code(ret)
         }
     }
 }
 
-pub(super) fn alloc_thread_cleanup(allocator: *mut snmalloc_edp::Alloc) {
+pub(super) fn alloc_thread_cleanup() {
     unsafe {
-        tls::set_tls_ptr(tls::TlsIndex::AllocPtr, core::ptr::null_mut());
+        let allocator = tls::get_tls_data(crate::sys::abi::tls::TlsIndex::AllocPtr) as *mut Alloc;
+        tls::set_tls_data(tls::TlsIndex::AllocPtr, core::ptr::null_mut());
         sn_thread_cleanup(allocator);
     }
 }
diff --git a/library/std/src/sys/pal/sgx/abi/tls/mod.rs b/library/std/src/sys/pal/sgx/abi/tls/mod.rs
index de47e646628dc..6d00e03766309 100644
--- a/library/std/src/sys/pal/sgx/abi/tls/mod.rs
+++ b/library/std/src/sys/pal/sgx/abi/tls/mod.rs
@@ -23,17 +23,20 @@ macro_rules! dup {
 #[export_name = "_ZN16__rust_internals3std3sys3sgx3abi3tls14TLS_DESTRUCTORE"]
 static TLS_DESTRUCTOR: [AtomicUsize; TLS_KEYS] = dup!((* * * * * * *) (AtomicUsize::new(0)));
 
-// Use TCSLS to store both the thread specific allocator
-// and TLS address at different indices
+// Use TCS Local Storage to store various thread specific data at
+// indices definied by the TlsIndex enum below
 #[repr(u8)]
 pub enum TlsIndex {
-    AllocPtr = 0,
-    TlsPtr = 1,
+    AllocPtr = 0, /* Thread specific allocator address */
+    TlsPtr = 1,   /* TLS address */
 }
 
+// Functions to access TLS data from the TCSLS. Use TlsIndex
+// enum to reference the desired field.
+// Function definition available in entry.S
 extern "C" {
-    pub fn get_tls_ptr(index: TlsIndex) -> *const u8;
-    pub fn set_tls_ptr(index: TlsIndex, tls: *const u8);
+    pub fn get_tls_data(index: TlsIndex) -> *const u8;
+    pub fn set_tls_data(index: TlsIndex, tls: *const u8);
 }
 
 #[derive(Copy, Clone)]
@@ -95,21 +98,21 @@ impl Tls {
     }
 
     pub unsafe fn activate(&self) -> ActiveTls<'_> {
-        // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        unsafe { set_tls_ptr(TlsIndex::TlsPtr, self as *const Tls as _) };
+        // FIXME: Needs safety information. See entry.S for `set_tls_data` definition.
+        unsafe { set_tls_data(TlsIndex::TlsPtr, self as *const Tls as _) };
         ActiveTls { tls: self }
     }
 
     #[allow(unused)]
     pub unsafe fn activate_persistent(self: Box<Self>) {
-        // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        unsafe { set_tls_ptr(TlsIndex::TlsPtr, core::ptr::addr_of!(*self) as _) };
+        // FIXME: Needs safety information. See entry.S for `set_tls_data` definition.
+        unsafe { set_tls_data(TlsIndex::TlsPtr, core::ptr::addr_of!(*self) as _) };
         mem::forget(self);
     }
 
     unsafe fn current<'a>() -> &'a Tls {
-        // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        unsafe { &*(get_tls_ptr(TlsIndex::TlsPtr) as *const Tls) }
+        // FIXME: Needs safety information. See entry.S for `set_tls_data` definition.
+        unsafe { &*(get_tls_data(TlsIndex::TlsPtr) as *const Tls) }
     }
 
     pub fn create(dtor: Option<unsafe extern "C" fn(*mut u8)>) -> Key {
diff --git a/library/std/src/sys/pal/sgx/alloc.rs b/library/std/src/sys/pal/sgx/alloc.rs
index ba6c16b260a9a..b50c58ecd8318 100644
--- a/library/std/src/sys/pal/sgx/alloc.rs
+++ b/library/std/src/sys/pal/sgx/alloc.rs
@@ -44,5 +44,5 @@ pub unsafe extern "C" fn __rust_c_dealloc(ptr: *mut u8, size: usize, align: usiz
 #[cfg(not(test))]
 #[no_mangle]
 pub extern "C" fn __rust_get_thread_allocator() -> *mut Alloc {
-    unsafe{ crate::sys::abi::tls::get_tls_ptr(crate::sys::abi::tls::TlsIndex::AllocPtr) as *mut Alloc }
+    unsafe{ crate::sys::abi::tls::get_tls_data(crate::sys::abi::tls::TlsIndex::AllocPtr) as *mut Alloc }
 }
diff --git a/mybuild.sh b/mybuild.sh
deleted file mode 100755
index d42183f621085..0000000000000
--- a/mybuild.sh
+++ /dev/null
@@ -1,43 +0,0 @@
-#!/bin/bash -ex
-export AR_x86_64_fortanix_unknown_sgx=ar
-export CC_x86_64_fortanix_unknown_sgx=clang-12
-export CFLAGS_x86_64_fortanix_unknown_sgx="-D__ELF__ -isystem/usr/include/x86_64-linux-gnu -mlvi-hardening -mllvm -x86-experimental-lvi-inline-asm-hardening"
-export CXX_x86_64_fortanix_unknown_sgx=clang++-12
-export CXXFLAGS_x86_64_fortanix_unknown_sgx="-D__ELF__ -isystem/usr/include/x86_64-linux-gnu -mlvi-hardening -mllvm -x86-experimental-lvi-inline-asm-hardening"
-export CC_x86_64_unknown_linux_gnu=clang-12
-export CXX_x86_64_unknown_linux_gnu=clang++-12
-
-git submodule update --init --recursive
-        detect_cxx_include_path() {
-            for path in $(clang++-12 -print-search-dirs|sed -n 's/^libraries:\s*=//p'|tr : ' '); do
-                num_component="$(basename "$path")"
-                if [[ "$num_component" =~ ^[0-9]+(\.[0-9]+)*$ ]]; then
-                    if [[ "$(basename "$(dirname "$path")")" == 'x86_64-linux-gnu' ]]; then
-                        echo $num_component
-                        return
-                    fi
-                fi
-            done
-            exit 1
-        }
-export CXXFLAGS_x86_64_fortanix_unknown_sgx="-cxx-isystem/usr/include/c++/$(detect_cxx_include_path) -cxx-isystem/usr/include/x86_64-linux-gnu/c++/$(detect_cxx_include_path) $CFLAGS_x86_64_fortanix_unknown_sgx"
-
-if [ ! -f config.toml ]; then
-    rustup default nightly
-    rustup update nightly
-    rustup target add x86_64-fortanix-unknown-sgx
-
-    cargo install fortanix-sgx-tools
-    cargo install sgxs-tools
-
-    git submodule foreach git reset --hard
-    ./configure --enable-lld --set profile=compiler \
-        --set install.prefix="/home/yuxiangcao/main_ws/toolchain/rust/rustup/toolchains/dev_snmalloc-2024-04-03-x86_64-unknown-linux-gnu" \
-        --set install.sysconfdir="etc" \
-        --set build.extended=true \
-        --set build.tools="cargo,rustfmt,src" \
-        --set build.docs=false
-
-fi
-
-./x install --target=x86_64-fortanix-unknown-sgx,x86_64-unknown-linux-gnu --host=x86_64-unknown-linux-gnu

From e9e1270591de2bf11daceb4e726e271b0b0d9add Mon Sep 17 00:00:00 2001
From: aditijannu <aditi.jannu@fortanix.com>
Date: Thu, 27 Jun 2024 00:57:35 +0200
Subject: [PATCH 11/12] Move snmalloc-edp from rust-sgx to std library

---
 .gitmodules                                   |   4 +
 Cargo.lock                                    |  15 +-
 library/snmalloc-edp/CMakeLists.txt           |  10 +
 library/snmalloc-edp/Cargo.toml               |  21 ++
 library/snmalloc-edp/build.rs                 |  71 +++++++
 library/snmalloc-edp/snmalloc                 |   1 +
 library/snmalloc-edp/src/lib.rs               |  15 ++
 .../src/rust-sgx-snmalloc-shim.cpp            | 181 ++++++++++++++++++
 library/snmalloc-edp/tests/global_alloc.rs    | 126 ++++++++++++
 library/std/Cargo.toml                        |   7 +-
 library/std/src/sys/pal/sgx/abi/tls/mod.rs    |   4 +-
 library/std/src/sys/pal/sgx/alloc.rs          |   1 +
 ...d-x86_64-fortanix-unknown-sgx-toolchain.sh |  16 ++
 13 files changed, 467 insertions(+), 5 deletions(-)
 create mode 100644 library/snmalloc-edp/CMakeLists.txt
 create mode 100644 library/snmalloc-edp/Cargo.toml
 create mode 100644 library/snmalloc-edp/build.rs
 create mode 160000 library/snmalloc-edp/snmalloc
 create mode 100644 library/snmalloc-edp/src/lib.rs
 create mode 100644 library/snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp
 create mode 100644 library/snmalloc-edp/tests/global_alloc.rs

diff --git a/.gitmodules b/.gitmodules
index 802d61eea293b..a5d571a6e635d 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -43,3 +43,7 @@
 	path = library/backtrace
 	url = https://github.com/rust-lang/backtrace-rs.git
 	shallow = true
+[submodule "library/snmalloc-edp/snmalloc"]
+	path = library/snmalloc-edp/snmalloc
+	url = https://github.com/microsoft/snmalloc.git
+        shallow = true
diff --git a/Cargo.lock b/Cargo.lock
index 6dac6c9b46ac3..8ac846cb2ac95 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1156,6 +1156,19 @@ version = "1.0.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "86e3bdc80eee6e16b2b6b0f87fbc98c04bee3455e35174c0de1a125d0688c632"
 
+[[package]]
+name = "dlmalloc"
+version = "0.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3264b043b8e977326c1ee9e723da2c1f8d09a99df52cacf00b4dbce5ac54414d"
+dependencies = [
+ "cfg-if",
+ "compiler_builtins",
+ "libc",
+ "rustc-std-workspace-core",
+ "windows-sys 0.52.0",
+]
+
 [[package]]
 name = "either"
 version = "1.10.0"
@@ -5146,7 +5159,6 @@ checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b"
 [[package]]
 name = "snmalloc-edp"
 version = "0.1.0"
-source = "git+https://github.com/fortanix/rust-sgx?branch=aj/update-sgx-alloc#6fc0b33a8c0fe2a36e9511b22c5cbc2df57df19e"
 dependencies = [
  "cc",
  "cmake",
@@ -5246,6 +5258,7 @@ dependencies = [
  "cfg-if",
  "compiler_builtins",
  "core",
+ "dlmalloc",
  "fortanix-sgx-abi",
  "hashbrown",
  "hermit-abi",
diff --git a/library/snmalloc-edp/CMakeLists.txt b/library/snmalloc-edp/CMakeLists.txt
new file mode 100644
index 0000000000000..b87eee490d6e8
--- /dev/null
+++ b/library/snmalloc-edp/CMakeLists.txt
@@ -0,0 +1,10 @@
+cmake_minimum_required(VERSION 3.14)
+set(CMAKE_TRY_COMPILE_TARGET_TYPE STATIC_LIBRARY)
+project(snmalloc-edp CXX)
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED True)
+set(SNMALLOC_HEADER_ONLY_LIBRARY ON)
+add_subdirectory(snmalloc EXCLUDE_FROM_ALL)
+add_library(snmalloc-edp src/rust-sgx-snmalloc-shim.cpp)
+target_link_libraries(snmalloc-edp PRIVATE snmalloc_lib)
+target_compile_options(snmalloc-edp PRIVATE -nostdlib -ffreestanding -fno-exceptions -mrdrnd -fPIC)
diff --git a/library/snmalloc-edp/Cargo.toml b/library/snmalloc-edp/Cargo.toml
new file mode 100644
index 0000000000000..a5ef6f8ad2a64
--- /dev/null
+++ b/library/snmalloc-edp/Cargo.toml
@@ -0,0 +1,21 @@
+[package]
+name = "snmalloc-edp"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+build = "build.rs"
+
+[build-dependencies]
+cc = "1.0.86"
+cmake = "0.1.50"
+elf = "0.7"
+
+[dependencies]
+core = { version = "1.0.0", optional = true, package = "rustc-std-workspace-core" }
+compiler_builtins = { version = "0.1.0", optional = true }
+
+[features]
+docs = []
+rustc-dep-of-std = ["core", "compiler_builtins/rustc-dep-of-std"]
diff --git a/library/snmalloc-edp/build.rs b/library/snmalloc-edp/build.rs
new file mode 100644
index 0000000000000..cf49d451a525d
--- /dev/null
+++ b/library/snmalloc-edp/build.rs
@@ -0,0 +1,71 @@
+use elf::ElfStream;
+use elf::endian::LittleEndian;
+use std::fs::{DirEntry, File};
+use std::path::{Path, PathBuf};
+
+fn files_in_dir(p: &Path) -> impl Iterator<Item = DirEntry> {
+    p.read_dir().unwrap().map(|e| e.unwrap()).filter(|e| e.file_type().unwrap().is_file())
+}
+
+fn main() {
+    let out_dir = PathBuf::from(std::env::var_os("OUT_DIR").unwrap());
+
+    // # Use CMake to build the shim
+    let mut dst = cmake::build(".");
+    dst.push("build");
+    println!("cargo:rustc-link-search=native={}", dst.display());
+
+    // ideally, the cmake crate would have a way to output this
+    println!("cargo:rerun-if-changed=CMakeLists.txt");
+    println!("cargo:rerun-if-changed=src/rust-sgx-snmalloc-shim.cpp");
+
+    // # Extract the static library archive into a temporary directory
+    let mut objs = out_dir.clone();
+    objs.push("objs");
+    std::fs::create_dir_all(&objs).unwrap();
+    // clear existing files in the temp dir
+    for file in files_in_dir(&objs) {
+        std::fs::remove_file(file.path()).unwrap();
+    }
+
+    dst.push("libsnmalloc-edp.a");
+
+    let mut ar = cc::Build::new().get_archiver();
+    ar.args(&["x", "--output"]);
+    ar.arg(&objs);
+    ar.arg(dst);
+    assert!(ar.status().unwrap().success());
+
+    // # Read the symbols from the shim ELF object
+    assert_eq!(files_in_dir(&objs).count(), 1);
+    let f = files_in_dir(&objs).next().unwrap();
+    let mut elf = elf::ElfStream::<elf::endian::LittleEndian, _>::open_stream(File::open(f.path()).unwrap()).unwrap();
+    let (symtab, strtab) = elf.symbol_table().unwrap().unwrap();
+    let mut sn_alloc_size = None;
+    let mut sn_alloc_align = None;
+    for sym in symtab {
+        match strtab.get(sym.st_name as _).unwrap() {
+            "sn_alloc_size" => assert!(sn_alloc_size.replace(sym).is_none()),
+            "sn_alloc_align" => assert!(sn_alloc_align.replace(sym).is_none()),
+            _ => {}
+        }
+    }
+    let sn_alloc_size = sn_alloc_size.expect("sn_alloc_size");
+    let sn_alloc_align = sn_alloc_align.expect("sn_alloc_align");
+
+    let sn_alloc_size = get_u64_at_symbol(sn_alloc_size, &mut elf);
+    let sn_alloc_align = get_u64_at_symbol(sn_alloc_align, &mut elf);
+
+    // # Write the type
+    let contents = format!("#[repr(align({}), C)] pub struct Alloc {{ _0: [u8; {}] }}", sn_alloc_align, sn_alloc_size);
+    let mut alloc_type_rs = out_dir.clone();
+    alloc_type_rs.push("alloc-type.rs");
+    std::fs::write(alloc_type_rs, contents).unwrap();
+}
+
+fn get_u64_at_symbol(sym: elf::symbol::Symbol, elf: &mut ElfStream<LittleEndian, File>) -> u64 {
+    assert_eq!(sym.st_size, 8);
+    let (data, _) = elf.section_data(&elf.section_headers()[sym.st_shndx as usize].clone()).unwrap();
+    let data: &[u8; 8] = data.split_at(8).0.try_into().unwrap();
+    u64::from_le_bytes(*data)
+}
\ No newline at end of file
diff --git a/library/snmalloc-edp/snmalloc b/library/snmalloc-edp/snmalloc
new file mode 160000
index 0000000000000..4620220080d25
--- /dev/null
+++ b/library/snmalloc-edp/snmalloc
@@ -0,0 +1 @@
+Subproject commit 4620220080d2582c829954624c5a84175c05fd93
diff --git a/library/snmalloc-edp/src/lib.rs b/library/snmalloc-edp/src/lib.rs
new file mode 100644
index 0000000000000..6669e65361c2d
--- /dev/null
+++ b/library/snmalloc-edp/src/lib.rs
@@ -0,0 +1,15 @@
+#![no_std]
+
+include!(concat!(env!("OUT_DIR"), "/alloc-type.rs"));
+
+#[link(name = "snmalloc-edp", kind = "static")]
+extern {
+    pub fn sn_global_init(heap_base: *const u8, heap_size: usize);
+    pub fn sn_thread_init(allocator: *mut Alloc);
+    pub fn sn_thread_cleanup(allocator: *mut Alloc);
+
+    pub fn sn_rust_alloc(alignment: usize, size: usize) -> *mut u8;
+    pub fn sn_rust_alloc_zeroed(alignment: usize, size: usize) -> *mut u8;
+    pub fn sn_rust_dealloc(ptr: *mut u8, alignment: usize, size: usize);
+    pub fn sn_rust_realloc(ptr: *mut u8, alignment: usize, old_size: usize, new_size: usize) -> *mut u8;
+}
diff --git a/library/snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp b/library/snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp
new file mode 100644
index 0000000000000..ee6850a2cb7b2
--- /dev/null
+++ b/library/snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp
@@ -0,0 +1,181 @@
+// Copyright (c) Microsoft Corporation.
+// Copyright (c) Open Enclave SDK contributors.
+// Copyright (c) 2020 SchrodingerZhu
+// Copyright (c) Fortanix, Inc.
+//
+// MIT License
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE
+
+#include <immintrin.h>
+#include <string.h>
+
+/***************************************************/
+/*** Imported symbols needed by snmalloc SGX PAL ***/
+/***************************************************/
+
+// from entry.S
+extern "C" size_t get_tcs_addr();
+
+// from Rust std
+extern "C" void __rust_print_err(const char* m, size_t s);
+extern "C" [[noreturn]] void __rust_abort();
+
+/*******************************************************/
+/*** Standard C functions needed by snmalloc SGX PAL ***/
+/*******************************************************/
+
+// definition needs to match GNU header
+extern "C" [[noreturn]] void abort() __THROW {
+    __rust_abort();
+}
+
+// definition needs to match GNU header and will not return an actual errno
+extern "C" inline int * __attribute_const__ __errno_location (void) __THROW {
+    static int errno;
+    return &errno;
+}
+
+/***********************************/
+/*** snmalloc SGX PAL definition ***/
+/***********************************/
+
+#define SNMALLOC_PROVIDE_OWN_CONFIG
+#define SNMALLOC_SGX
+#define SNMALLOC_USE_SMALL_CHUNKS
+#define SNMALLOC_MEMORY_PROVIDER PALEdpSgx
+
+#include "../snmalloc/src/snmalloc/pal/pal_noalloc.h"
+
+namespace snmalloc {
+void register_clean_up() {
+    // Unused on SGX
+    abort();
+}
+
+class EdpErrorHandler {
+  public:
+    static void print_stack_trace() {}
+
+    [[noreturn]] static void error(const char* const str) {
+        __rust_print_err(str, strlen(str));
+        abort();
+    }
+    static constexpr size_t address_bits = Aal::address_bits;
+    static constexpr size_t page_size = Aal::smallest_page_size;
+};
+
+using EdpBasePAL = PALNoAlloc<EdpErrorHandler>;
+
+class PALEdpSgx : public EdpBasePAL {
+   public:
+    const static size_t RAND_NUM_GEN_MAX_RETRIES = 64;
+    using ThreadIdentity = size_t;
+    static constexpr uint64_t pal_features = EdpBasePAL::pal_features | Entropy;
+
+    template <bool page_aligned = false>
+    static void zero(void *p, size_t size) noexcept {
+        memset(p, 0, size);
+    }
+
+    static inline uint64_t get_entropy64() {
+        for (size_t retry_count = 0; retry_count < RAND_NUM_GEN_MAX_RETRIES; retry_count++) {
+            long long unsigned int result;
+            if (_rdrand64_step(&result) == 1) {
+                return result;
+            }
+        }
+        EdpErrorHandler::error("no entropy available");
+    }
+
+    static inline ThreadIdentity get_tid() noexcept {
+        return (size_t)get_tcs_addr();
+    }
+};
+} // namespace snmalloc
+
+/**************************************/
+/*** Instantiation of the allocator ***/
+/**************************************/
+
+#include "../snmalloc/src/snmalloc/backend/fixedglobalconfig.h"
+#include "../snmalloc/src/snmalloc/snmalloc_core.h"
+
+using namespace snmalloc;
+
+using Globals = FixedRangeConfig<PALEdpSgx>;
+using Alloc = LocalAllocator<Globals>;
+
+/// Do global initialization for snmalloc. Should be called exactly once prior
+/// to any other snmalloc function calls.
+extern "C" void sn_global_init(void* heap_base, size_t heap_size) {
+    Globals::init(nullptr, heap_base, heap_size);
+}
+
+/// Construct a thread-local allocator object in place
+extern "C" void sn_thread_init(Alloc* allocator) {
+    new(allocator) Alloc();
+    allocator->init();
+}
+
+/// Destruct a thread-local allocator object in place
+extern "C" void sn_thread_cleanup(Alloc* allocator) {
+    allocator->teardown();
+    allocator->~Alloc();
+}
+
+extern "C" size_t sn_alloc_size = sizeof(Alloc);
+extern "C" size_t sn_alloc_align = alignof(Alloc);
+
+/// Return a pointer to a thread-local allocator object of size
+/// `sn_alloc_size` and alignment `sn_alloc_align`.
+extern "C" Alloc* __rust_get_thread_allocator();
+
+/******************************************************/
+/*** Rust-compatible shims for the global allocator ***/
+/******************************************************/
+
+extern "C" void *sn_rust_alloc(size_t alignment, size_t size) {
+    return __rust_get_thread_allocator()->alloc(aligned_size(alignment, size));
+}
+
+extern "C" void *sn_rust_alloc_zeroed(size_t alignment, size_t size) {
+    return __rust_get_thread_allocator()->alloc<YesZero>(
+        aligned_size(alignment, size));
+}
+
+extern "C" void sn_rust_dealloc(void *ptr, size_t alignment, size_t size) {
+    __rust_get_thread_allocator()->dealloc(ptr, aligned_size(alignment, size));
+}
+
+extern "C" void *sn_rust_realloc(void *ptr, size_t alignment, size_t old_size,
+                                 size_t new_size) {
+    size_t aligned_old_size = aligned_size(alignment, old_size),
+           aligned_new_size = aligned_size(alignment, new_size);
+    if (size_to_sizeclass_full(aligned_old_size).raw() ==
+        size_to_sizeclass_full(aligned_new_size).raw())
+        return ptr;
+    Alloc* allocator = __rust_get_thread_allocator();
+    void *p = allocator->alloc(aligned_new_size);
+    if (p) {
+        std::memcpy(p, ptr, old_size < new_size ? old_size : new_size);
+        allocator->dealloc(ptr, aligned_old_size);
+    }
+    return p;
+}
diff --git a/library/snmalloc-edp/tests/global_alloc.rs b/library/snmalloc-edp/tests/global_alloc.rs
new file mode 100644
index 0000000000000..4b5a8be82d241
--- /dev/null
+++ b/library/snmalloc-edp/tests/global_alloc.rs
@@ -0,0 +1,126 @@
+use std::{alloc::{self, GlobalAlloc}, cell::Cell, ptr};
+use std::ptr::null_mut;
+
+use snmalloc_edp::*;
+
+thread_local! {
+    static THREAD_ALLOC: Cell<*mut Alloc> = const { Cell::new(ptr::null_mut()) };
+}
+
+#[no_mangle]
+pub fn __rust_get_thread_allocator() -> *mut Alloc {
+    THREAD_ALLOC.get()
+}
+
+struct System;
+
+unsafe impl alloc::GlobalAlloc for System {
+    #[inline]
+    unsafe fn alloc(&self, layout: alloc::Layout) -> *mut u8 {
+        // SAFETY: the caller must uphold the safety contract for `malloc`
+        sn_rust_alloc(layout.align(), layout.size())
+    }
+
+    #[inline]
+    unsafe fn alloc_zeroed(&self, layout: alloc::Layout) -> *mut u8 {
+        // SAFETY: the caller must uphold the safety contract for `malloc`
+        sn_rust_alloc_zeroed(layout.align(), layout.size())
+    }
+
+    #[inline]
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: alloc::Layout) {
+        // SAFETY: the caller must uphold the safety contract for `malloc`
+        sn_rust_dealloc(ptr, layout.align(), layout.size())
+    }
+
+    #[inline]
+    unsafe fn realloc(&self, ptr: *mut u8, layout: alloc::Layout, new_size: usize) -> *mut u8 {
+        // SAFETY: the caller must uphold the safety contract for `malloc`
+        sn_rust_realloc(ptr, layout.align(), layout.size(), new_size)
+    }
+}
+
+// SAFETY: this should only be called once per thread, and the global
+// allocator shouldn't be used outside of this function
+unsafe fn with_thread_allocator<F: FnOnce() -> R, R>(f: F) -> R {
+    unsafe {
+        let mut allocator = std::mem::MaybeUninit::<Alloc>::uninit();
+        let thread_alloc = __rust_get_thread_allocator();
+        assert_eq!(thread_alloc, null_mut());
+        sn_thread_init(allocator.as_mut_ptr());
+        THREAD_ALLOC.set(allocator.as_mut_ptr());
+
+        let r = f();
+
+        THREAD_ALLOC.set(null_mut());
+        sn_thread_cleanup(allocator.as_mut_ptr());
+
+        r
+    }
+}
+
+#[test]
+fn test() {
+    unsafe {
+        #[allow(dead_code)]
+        #[derive(Copy, Clone)]
+        #[repr(align(0x1000))]
+        struct Page([u8; 0x1000]);
+
+            // allocate a dummy heap
+            let heap = (*Box::into_raw(vec![Page([0; 4096]); 100].into_boxed_slice())).as_mut_ptr_range();
+            let heap_size = heap.end as usize - heap.start as usize;
+            sn_global_init(heap.start as _, heap_size);
+    }
+
+    const TEST_ARRAY_SIZE :usize = 20;
+    type AllocTestType = [u64; TEST_ARRAY_SIZE];
+
+    let barrier = std::sync::Barrier::new(2);
+
+    std::thread::scope(|s| {
+        let (tx, rx) = std::sync::mpsc::sync_channel(0);
+        let barrier = &barrier;
+        s.spawn(move || {
+            unsafe {
+                // Initialize thread allocator and perform alloc/alloc_zeroed
+                with_thread_allocator(|| {
+                    let p1 = System.alloc(alloc::Layout::new::<AllocTestType>());
+                    barrier.wait();
+                    let p2 = System.alloc_zeroed(alloc::Layout::new::<AllocTestType>());
+
+                    let p2_slice = std::slice::from_raw_parts_mut(p2, TEST_ARRAY_SIZE);
+                    for i in 0..TEST_ARRAY_SIZE {
+                        assert_eq!(p2_slice[i], 0);
+                    }
+
+                    tx.send((p1 as usize, p2 as usize)).unwrap();
+                })
+            };
+        });
+
+        let (p1, p2) = unsafe {
+            with_thread_allocator(|| {
+                let p1 = System.alloc(alloc::Layout::new::<AllocTestType>());
+                barrier.wait();
+                let p2 = System.alloc(alloc::Layout::new::<AllocTestType>());
+
+                // Test realloc
+                let p3 = System.realloc(p1, alloc::Layout::new::<AllocTestType>(), TEST_ARRAY_SIZE * 2);
+                assert_ne!(p3, p2);
+
+                (p1 as usize, p2 as usize)
+            })
+        };
+
+        let (p3, p4) = rx.recv().unwrap();
+        assert_ne!(p1, p2);
+        assert_ne!(p1, p3);
+        assert_ne!(p1, p4);
+        assert_ne!(p2, p3);
+        assert_ne!(p2, p4);
+        assert_ne!(p3, p4);
+
+    })
+
+}
diff --git a/library/std/Cargo.toml b/library/std/Cargo.toml
index d45b86a017acd..cd228d20ff130 100644
--- a/library/std/Cargo.toml
+++ b/library/std/Cargo.toml
@@ -41,8 +41,11 @@ object = { version = "0.32.0", default-features = false, optional = true, featur
 rand = { version = "0.8.5", default-features = false, features = ["alloc"] }
 rand_xorshift = "0.3.0"
 
-[target.'cfg(any(all(target_family = "wasm", target_os = "unknown"), target_os = "xous", all(target_vendor = "fortanix", target_env = "sgx")))'.dependencies]
-snmalloc-edp = { git = "https://github.com/fortanix/rust-sgx", branch = "aj/update-sgx-alloc", features = ['rustc-dep-of-std'], public = true }
+[target.'cfg(any(all(target_family = "wasm", target_os = "unknown"), target_os = "xous"))'.dependencies]
+dlmalloc = { version = "0.2.4", features = ['rustc-dep-of-std'] }
+
+[target.'cfg(all(target_vendor = "fortanix", target_env = "sgx"))'.dependencies]
+snmalloc-edp = { path = "../snmalloc-edp" , features = ['rustc-dep-of-std'], public = true }
 
 [target.x86_64-fortanix-unknown-sgx.dependencies]
 fortanix-sgx-abi = { version = "0.5.0", features = ['rustc-dep-of-std'], public = true }
diff --git a/library/std/src/sys/pal/sgx/abi/tls/mod.rs b/library/std/src/sys/pal/sgx/abi/tls/mod.rs
index 6d00e03766309..f5e3746d37d41 100644
--- a/library/std/src/sys/pal/sgx/abi/tls/mod.rs
+++ b/library/std/src/sys/pal/sgx/abi/tls/mod.rs
@@ -24,7 +24,7 @@ macro_rules! dup {
 static TLS_DESTRUCTOR: [AtomicUsize; TLS_KEYS] = dup!((* * * * * * *) (AtomicUsize::new(0)));
 
 // Use TCS Local Storage to store various thread specific data at
-// indices definied by the TlsIndex enum below
+// indices definied by the TlsIndex enum below.
 #[repr(u8)]
 pub enum TlsIndex {
     AllocPtr = 0, /* Thread specific allocator address */
@@ -111,7 +111,7 @@ impl Tls {
     }
 
     unsafe fn current<'a>() -> &'a Tls {
-        // FIXME: Needs safety information. See entry.S for `set_tls_data` definition.
+        // FIXME: Needs safety information. See entry.S for `get_tls_data` definition.
         unsafe { &*(get_tls_data(TlsIndex::TlsPtr) as *const Tls) }
     }
 
diff --git a/library/std/src/sys/pal/sgx/alloc.rs b/library/std/src/sys/pal/sgx/alloc.rs
index b50c58ecd8318..fa18b7ee5419c 100644
--- a/library/std/src/sys/pal/sgx/alloc.rs
+++ b/library/std/src/sys/pal/sgx/alloc.rs
@@ -5,6 +5,7 @@ use snmalloc_edp::*;
 unsafe impl GlobalAlloc for System {
     #[inline]
     unsafe fn alloc(&self, layout: alloc::Layout) -> *mut u8 {
+        // SAFETY: the caller must uphold the safety contract for `malloc`
         unsafe { sn_rust_alloc(layout.align(), layout.size()) }
     }
 
diff --git a/src/ci/docker/host-x86_64/dist-various-2/build-x86_64-fortanix-unknown-sgx-toolchain.sh b/src/ci/docker/host-x86_64/dist-various-2/build-x86_64-fortanix-unknown-sgx-toolchain.sh
index eabff87284dc1..5766fcf6b14aa 100755
--- a/src/ci/docker/host-x86_64/dist-various-2/build-x86_64-fortanix-unknown-sgx-toolchain.sh
+++ b/src/ci/docker/host-x86_64/dist-various-2/build-x86_64-fortanix-unknown-sgx-toolchain.sh
@@ -17,4 +17,20 @@ install_prereq() {
             clang-11
 }
 
+detect_cxx_include_path() {
+    for path in $(clang++-11 -print-search-dirs|sed -n 's/^libraries:\s*=//p'|tr : ' '); do
+        num_component="$(basename "$path")"
+        if [[ "$num_component" =~ ^[0-9]+(\.[0-9]+)*$ ]]; then
+            if [[ "$(basename "$(dirname "$path")")" == 'x86_64-linux-gnu' ]]; then
+                echo $num_component
+                return
+            fi
+        fi
+    done
+    exit 1
+}
+
 hide_output install_prereq
+
+# Note - this overwrites the environment variable set in the Dockerfile
+export CXXFLAGS_x86_64_fortanix_unknown_sgx="-cxx-isystem/usr/include/c++/$(detect_cxx_include_path) -cxx-isystem/usr/include/x86_64-linux-gnu/c++/$(detect_cxx_include_path) $CFLAGS_x86_64_fortanix_unknown_sgx"

From 87264c69cca10fe434b2174f7795cb86ad813ac2 Mon Sep 17 00:00:00 2001
From: aditijannu <aditi.jannu@fortanix.com>
Date: Tue, 16 Jul 2024 11:37:13 +0200
Subject: [PATCH 12/12] Update formatting in .gitmodules

---
 .gitmodules | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.gitmodules b/.gitmodules
index a5d571a6e635d..727053d94e9fd 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -46,4 +46,4 @@
 [submodule "library/snmalloc-edp/snmalloc"]
 	path = library/snmalloc-edp/snmalloc
 	url = https://github.com/microsoft/snmalloc.git
-        shallow = true
+	shallow = true