Initialize the mspaces of allocator and keep track of it

src/ld_so/start.rs

@@ -3,7 +3,12 @@
 use alloc::{borrow::ToOwned, boxed::Box, collections::BTreeMap, string::String, vec::Vec};
 
 use crate::{
-    c_str::CStr, header::unistd, platform::types::c_char, start::Stack, sync::mutex::Mutex,
+    c_str::CStr,
+    header::unistd,
+    platform::{new_mspace, types::c_char},
+    start::Stack,
+    sync::mutex::Mutex,
+    ALLOCATOR,
 };
 
 use super::{
@@ -116,7 +121,9 @@ unsafe fn adjust_stack(sp: &'static mut Stack) {
 }
 #[no_mangle]
 pub extern "C" fn relibc_ld_so_start(sp: &'static mut Stack, ld_entry: usize) -> usize {
-    // first we get the arguments, the environment, and the auxilary vector
+    // First thing we initialize the mspace
+    ALLOCATOR.set_book_keeper(new_mspace());
+    // next we get the arguments, the environment, and the auxilary vector
     let (argv, envs, auxv) = unsafe {
         let argv_start = sp.argv() as *mut usize;
         let (argv, argv_end) = get_argv(argv_start);
@@ -210,6 +217,7 @@ pub extern "C" fn relibc_ld_so_start(sp: &'static mut Stack, ld_entry: usize) ->
     }
     if let Some(tcb) = unsafe { Tcb::current() } {
         tcb.linker_ptr = Box::into_raw(Box::new(Mutex::new(linker)));
+        tcb.mspace = ALLOCATOR.get_book_keeper();
     }
     if is_manual {
         eprintln!("ld.so: entry '{}': {:#x}", path, entry);

src/ld_so/tcb.rs

@@ -45,6 +45,8 @@ pub struct Tcb {
     pub masters_len: usize,
     /// Pointer to dynamic linker
     pub linker_ptr: *const Mutex<Linker>,
+    /// pointer to rust memory allocator structure
+    pub mspace: usize,
 }
 
 impl Tcb {
@@ -64,6 +66,7 @@ impl Tcb {
                 masters_ptr: ptr::null_mut(),
                 masters_len: 0,
                 linker_ptr: ptr::null(),
+                mspace: 0,
             },
         );
 

src/platform/allocator/dlmalloc.rs

@@ -6,6 +6,7 @@ use core::{
 use super::types::*;
 
 extern "C" {
+    fn create_mspace(capacity: size_t, locked: c_int) -> usize;
     fn dlmalloc(bytes: size_t) -> *mut c_void;
     fn dlmemalign(alignment: size_t, bytes: size_t) -> *mut c_void;
     fn dlrealloc(oldmem: *mut c_void, bytes: size_t) -> *mut c_void;
@@ -16,15 +17,15 @@ pub struct Allocator {
     mstate: AtomicUsize,
 }
 
-pub const NEWALLOCATOR:Allocator = Allocator{
+pub const NEWALLOCATOR: Allocator = Allocator {
     mstate: AtomicUsize::new(0),
 };
 
 impl Allocator {
-    pub fn set_bookkeeper(&self, mstate: usize) {
+    pub fn set_book_keeper(&self, mstate: usize) {
         self.mstate.store(mstate, Ordering::Relaxed);
     }
-    fn get_bookKeeper(&self) ->usize {
+    pub fn get_book_keeper(&self) -> usize {
         self.mstate.load(Ordering::Relaxed)
     }
 }
@@ -53,3 +54,7 @@ pub unsafe fn realloc(ptr: *mut c_void, size: size_t) -> *mut c_void {
 pub unsafe fn free(ptr: *mut c_void) {
     dlfree(ptr)
 }
+
+pub fn new_mspace() -> usize {
+    unsafe { create_mspace(0, 0) }
+}

src/platform/pte.rs

@@ -17,6 +17,7 @@ use crate::{
         Pal, Sys,
     },
     sync::Mutex,
+    ALLOCATOR,
 };
 
 pub struct Semaphore {
@@ -76,6 +77,7 @@ unsafe extern "C" fn pte_osThreadShim(
     tls_masters_ptr: *mut Master,
     tls_masters_len: usize,
     tls_linker_ptr: *const Mutex<Linker>,
+    tls_mspace: usize,
 ) {
     // The kernel allocated TLS does not have masters set, so do not attempt to copy it.
     // It will be copied by the kernel.
@@ -84,6 +86,7 @@ unsafe extern "C" fn pte_osThreadShim(
         tcb.masters_ptr = tls_masters_ptr;
         tcb.masters_len = tls_masters_len;
         tcb.linker_ptr = tls_linker_ptr;
+        tcb.mspace = tls_mspace;
         tcb.copy_masters().unwrap();
         tcb.activate();
     }
@@ -134,11 +137,13 @@ pub unsafe extern "C" fn pte_osThreadCreate(
         push(0);
 
         if let Some(tcb) = Tcb::current() {
+            push(tcb.mspace as usize);
             push(tcb.linker_ptr as usize);
             push(tcb.masters_len);
             push(tcb.masters_ptr as usize);
             push(tcb.tls_len);
         } else {
+            push(ALLOCATOR.get_book_keeper());
             push(0);
             push(0);
             push(0);

src/start.rs

@@ -4,7 +4,8 @@ use core::{intrinsics, ptr};
 use crate::{
     header::{stdio, stdlib},
     ld_so,
-    platform::{self, types::*, Pal, Sys},
+    platform::{self, new_mspace, types::*, Pal, Sys},
+    ALLOCATOR,
 };
 
 #[repr(C)]
@@ -92,6 +93,13 @@ pub unsafe extern "C" fn relibc_start(sp: &'static Stack) -> ! {
         fn _init();
         fn main(argc: isize, argv: *mut *mut c_char, envp: *mut *mut c_char) -> c_int;
     }
+    // Step 1 setup the right allocator...
+    // if any memory rust based memory allocation happen before this step .. we are doomed.
+    if let Some(tcb) = ld_so::tcb::Tcb::current() {
+        ALLOCATOR.set_book_keeper(tcb.mspace);
+    } else {
+        ALLOCATOR.set_book_keeper(new_mspace());
+    }
 
     // Ensure correct host system before executing more system calls
     relibc_verify_host();