From f3a33cee3abdd4b2e788faa613c82900dcda9328 Mon Sep 17 00:00:00 2001
From: Thalia Archibald <thalia@archibald.dev>
Date: Mon, 14 Apr 2025 05:01:11 -0700
Subject: [PATCH] SGX: Use OnceLock for env::Args and LazyLock for env::Vars

Fixes fuzzy provenance casts with `AtomicUsize`.
---
 library/std/src/sys/args/sgx.rs   | 20 +++++++--------
 library/std/src/sys/pal/sgx/os.rs | 41 +++++++------------------------
 2 files changed, 18 insertions(+), 43 deletions(-)

diff --git a/library/std/src/sys/args/sgx.rs b/library/std/src/sys/args/sgx.rs
index 0185a8a600094..90b1391be014b 100644
--- a/library/std/src/sys/args/sgx.rs
+++ b/library/std/src/sys/args/sgx.rs
@@ -1,7 +1,5 @@
-#![allow(fuzzy_provenance_casts)] // FIXME: this module systematically confuses pointers and integers
-
 use crate::ffi::OsString;
-use crate::sync::atomic::{AtomicUsize, Ordering};
+use crate::sync::OnceLock;
 use crate::sys::os_str::Buf;
 use crate::sys::pal::abi::usercalls::alloc;
 use crate::sys::pal::abi::usercalls::raw::ByteBuffer;
@@ -11,23 +9,23 @@ use crate::{fmt, slice};
 // Specifying linkage/symbol name is solely to ensure a single instance between this crate and its unit tests
 #[cfg_attr(test, linkage = "available_externally")]
 #[unsafe(export_name = "_ZN16__rust_internals3std3sys3sgx4args4ARGSE")]
-static ARGS: AtomicUsize = AtomicUsize::new(0);
+static ARGS: OnceLock<ArgsStore> = OnceLock::new();
 type ArgsStore = Vec<OsString>;
 
 #[cfg_attr(test, allow(dead_code))]
 pub unsafe fn init(argc: isize, argv: *const *const u8) {
     if argc != 0 {
-        let args = unsafe { alloc::User::<[ByteBuffer]>::from_raw_parts(argv as _, argc as _) };
-        let args = args
-            .iter()
-            .map(|a| OsString::from_inner(Buf { inner: a.copy_user_buffer() }))
-            .collect::<ArgsStore>();
-        ARGS.store(Box::into_raw(Box::new(args)) as _, Ordering::Relaxed);
+        ARGS.get_or_init(|| {
+            let args = unsafe { alloc::User::<[ByteBuffer]>::from_raw_parts(argv as _, argc as _) };
+            args.iter()
+                .map(|a| OsString::from_inner(Buf { inner: a.copy_user_buffer() }))
+                .collect::<ArgsStore>()
+        });
     }
 }
 
 pub fn args() -> Args {
-    let args = unsafe { (ARGS.load(Ordering::Relaxed) as *const ArgsStore).as_ref() };
+    let args = ARGS.get();
     if let Some(args) = args { Args(args.iter()) } else { Args([].iter()) }
 }
 
diff --git a/library/std/src/sys/pal/sgx/os.rs b/library/std/src/sys/pal/sgx/os.rs
index 010634cf31063..4723c3923ade4 100644
--- a/library/std/src/sys/pal/sgx/os.rs
+++ b/library/std/src/sys/pal/sgx/os.rs
@@ -1,3 +1,5 @@
+#![forbid(fuzzy_provenance_casts)]
+
 use fortanix_sgx_abi::{Error, RESULT_SUCCESS};
 
 use crate::collections::HashMap;
@@ -5,8 +7,7 @@ use crate::error::Error as StdError;
 use crate::ffi::{OsStr, OsString};
 use crate::marker::PhantomData;
 use crate::path::{self, PathBuf};
-use crate::sync::atomic::{AtomicUsize, Ordering};
-use crate::sync::{Mutex, Once};
+use crate::sync::{LazyLock, Mutex};
 use crate::sys::{decode_error_kind, sgx_ineffective, unsupported};
 use crate::{fmt, io, str, vec};
 
@@ -76,24 +77,9 @@ pub fn current_exe() -> io::Result<PathBuf> {
 // Specifying linkage/symbol name is solely to ensure a single instance between this crate and its unit tests
 #[cfg_attr(test, linkage = "available_externally")]
 #[unsafe(export_name = "_ZN16__rust_internals3std3sys3pal3sgx2os3ENVE")]
-static ENV: AtomicUsize = AtomicUsize::new(0);
-// Specifying linkage/symbol name is solely to ensure a single instance between this crate and its unit tests
-#[cfg_attr(test, linkage = "available_externally")]
-#[unsafe(export_name = "_ZN16__rust_internals3std3sys3pal3sgx2os8ENV_INITE")]
-static ENV_INIT: Once = Once::new();
+static ENV: LazyLock<EnvStore> = LazyLock::new(|| EnvStore::default());
 type EnvStore = Mutex<HashMap<OsString, OsString>>;
 
-fn get_env_store() -> Option<&'static EnvStore> {
-    unsafe { (ENV.load(Ordering::Relaxed) as *const EnvStore).as_ref() }
-}
-
-fn create_env_store() -> &'static EnvStore {
-    ENV_INIT.call_once(|| {
-        ENV.store(Box::into_raw(Box::new(EnvStore::default())) as _, Ordering::Relaxed)
-    });
-    unsafe { &*(ENV.load(Ordering::Relaxed) as *const EnvStore) }
-}
-
 pub struct Env {
     iter: vec::IntoIter<(OsString, OsString)>,
 }
@@ -140,31 +126,22 @@ impl Iterator for Env {
 }
 
 pub fn env() -> Env {
-    let clone_to_vec = |map: &HashMap<OsString, OsString>| -> Vec<_> {
-        map.iter().map(|(k, v)| (k.clone(), v.clone())).collect()
-    };
-
-    let iter = get_env_store()
-        .map(|env| clone_to_vec(&env.lock().unwrap()))
-        .unwrap_or_default()
-        .into_iter();
-    Env { iter }
+    let env = ENV.lock().unwrap().iter().map(|(k, v)| (k.clone(), v.clone())).collect::<Vec<_>>();
+    Env { iter: env.into_iter() }
 }
 
 pub fn getenv(k: &OsStr) -> Option<OsString> {
-    get_env_store().and_then(|s| s.lock().unwrap().get(k).cloned())
+    ENV.lock().unwrap().get(k).cloned()
 }
 
 pub unsafe fn setenv(k: &OsStr, v: &OsStr) -> io::Result<()> {
     let (k, v) = (k.to_owned(), v.to_owned());
-    create_env_store().lock().unwrap().insert(k, v);
+    ENV.lock().unwrap().insert(k, v);
     Ok(())
 }
 
 pub unsafe fn unsetenv(k: &OsStr) -> io::Result<()> {
-    if let Some(env) = get_env_store() {
-        env.lock().unwrap().remove(k);
-    }
+    ENV.lock().unwrap().remove(k);
     Ok(())
 }