From bbc86d8027e3019c5ad81b21f6d8f99a2cb2273a Mon Sep 17 00:00:00 2001
From: Anhad Singh <andypython@protonmail.com>
Date: Sun, 11 Jan 2026 23:06:50 +1100
Subject: [PATCH] fix(libc): use `ld.so`'s allocator

Previously, this was implemented incorrectly because `tcb.mspace` was
set just *before* jumping to the program's entry point. However,
`alloc_init` is called when the init functions are run for `libc.so`, so
`tcb.mspace` must be initialized *before* the init functions are
run.

Signed-off-by: Anhad Singh <andypython@protonmail.com>
---
 src/ld_so/linker.rs           |  2 ++
 src/ld_so/start.rs            |  8 +++-----
 src/platform/allocator/mod.rs | 28 +++++++++++++++++++++-------
 src/start.rs                  |  6 ++++--
 4 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/src/ld_so/linker.rs b/src/ld_so/linker.rs
index 69db3aa1a0..df3542506f 100644
--- a/src/ld_so/linker.rs
+++ b/src/ld_so/linker.rs
@@ -16,6 +16,7 @@ use core::{
 };
 
 use crate::{
+    ALLOCATOR,
     c_str::{CStr, CString},
     error::Errno,
     header::{
@@ -700,6 +701,7 @@ impl Linker {
                     #[cfg(target_os = "redox")]
                     Some(thr_fd),
                 );
+                tcb.mspace = ALLOCATOR.get();
 
                 #[cfg(target_os = "redox")]
                 {
diff --git a/src/ld_so/start.rs b/src/ld_so/start.rs
index eed04d81a4..fc763f8313 100644
--- a/src/ld_so/start.rs
+++ b/src/ld_so/start.rs
@@ -11,7 +11,6 @@ use alloc::{
 };
 
 use crate::{
-    ALLOCATOR,
     c_str::CStr,
     header::{
         sys_auxv::{AT_ENTRY, AT_PHDR},
@@ -249,7 +248,7 @@ pub unsafe extern "C" fn relibc_ld_so_start(sp: &'static mut Stack, ld_entry: us
     let (path, _name) = match resolve_path_name(&name_or_path, &envs) {
         Some((p, n)) => (p, n),
         None => {
-            eprintln!("ld.so: failed to locate '{}'", name_or_path);
+            eprintln!("[ld.so]: failed to locate '{name_or_path}'");
             unistd::_exit(1);
         }
     };
@@ -271,17 +270,16 @@ pub unsafe extern "C" fn relibc_ld_so_start(sp: &'static mut Stack, ld_entry: us
     let entry = match linker.load_program(&path, base_addr) {
         Ok(entry) => entry,
         Err(err) => {
-            eprintln!("[ld.so]: failed to link '{}': {:?}", path, err);
+            eprintln!("[ld.so]: failed to link '{path}': {err:?}");
             eprintln!("[ld.so]: enable debug output with `LD_DEBUG=all` for more information");
             unistd::_exit(1);
         }
     };
     if let Some(tcb) = unsafe { Tcb::current() } {
         tcb.linker_ptr = Box::into_raw(Box::new(Mutex::new(linker)));
-        tcb.mspace = ALLOCATOR.get();
     }
     if is_manual {
-        eprintln!("[ld.so]: entry '{}': {:#x}", path, entry);
+        eprintln!("[ld.so]: entry '{path}': {entry:#x}");
     }
     entry
 }
diff --git a/src/platform/allocator/mod.rs b/src/platform/allocator/mod.rs
index bee1c7628e..f5c4fe6fe7 100644
--- a/src/platform/allocator/mod.rs
+++ b/src/platform/allocator/mod.rs
@@ -6,7 +6,8 @@ use core::{
     cell::SyncUnsafeCell,
     cmp,
     mem::align_of,
-    ptr::{copy_nonoverlapping, write_bytes},
+    ptr::{self, copy_nonoverlapping, write_bytes},
+    sync::atomic::{AtomicPtr, Ordering},
 };
 
 mod sys;
@@ -18,17 +19,30 @@ pub type Dlmalloc = DlmallocCApi<sys::System>;
 
 pub const NEWALLOCATOR: Allocator = Allocator::new();
 
-pub struct Allocator(SyncUnsafeCell<Mutex<Dlmalloc>>);
+pub struct Allocator {
+    inner: SyncUnsafeCell<Mutex<Dlmalloc>>,
+    pub ptr: AtomicPtr<Mutex<Dlmalloc>>,
+}
 
 impl Allocator {
     pub const fn new() -> Self {
-        Allocator(SyncUnsafeCell::new(Mutex::new(Dlmalloc::new(
-            sys::System::new(),
-        ))))
+        Allocator {
+            inner: SyncUnsafeCell::new(Mutex::new(Dlmalloc::new(sys::System::new()))),
+            ptr: AtomicPtr::new(ptr::null_mut()),
+        }
     }
 
-    pub fn get(&self) -> *mut Mutex<Dlmalloc> {
-        self.0.get()
+    pub fn get(&self) -> *const Mutex<Dlmalloc> {
+        let ptr = self.ptr.load(Ordering::Acquire);
+        if !ptr.is_null() {
+            return ptr;
+        }
+
+        self.inner.get()
+    }
+
+    pub fn set(&self, mspace: *const Mutex<Dlmalloc>) {
+        self.ptr.store(mspace.cast_mut(), Ordering::Release);
     }
 }
 
diff --git a/src/start.rs b/src/start.rs
index 4d5257b693..9ea4a10f18 100644
--- a/src/start.rs
+++ b/src/start.rs
@@ -88,7 +88,7 @@ fn alloc_init() {
     unsafe {
         if let Some(tcb) = ld_so::tcb::Tcb::current() {
             if !tcb.mspace.is_null() {
-                ALLOCATOR.get().write(tcb.mspace.read());
+                ALLOCATOR.set(tcb.mspace);
             }
         }
     }
@@ -177,7 +177,9 @@ pub unsafe extern "C" fn relibc_start_v1(
 
     if let Some(tcb) = unsafe { ld_so::tcb::Tcb::current() } {
         // Update TCB mspace
-        tcb.mspace = ALLOCATOR.get();
+        if tcb.mspace.is_null() {
+            tcb.mspace = ALLOCATOR.get();
+        }
 
         // Set linker pointer if necessary
         if tcb.linker_ptr.is_null() {