From d7d2a3bc7129aa66c5fb8fdcc93984579425a9fb Mon Sep 17 00:00:00 2001 From: aditijannu Date: Thu, 27 Jun 2024 00:57:35 +0200 Subject: [PATCH] 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 | 68 +++++++ 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 | 108 +++++++++++ library/std/Cargo.toml | 7 +- ...d-x86_64-fortanix-unknown-sgx-toolchain.sh | 16 ++ 11 files changed, 443 insertions(+), 3 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..6e8917264af14 --- /dev/null +++ b/library/snmalloc-edp/build.rs @@ -0,0 +1,68 @@ +use std::fs::{DirEntry, File}; +use std::path::{Path, PathBuf}; + +fn files_in_dir(p: &Path) -> impl Iterator { + 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 + let f = files_in_dir(&objs).next().unwrap(); + let mut elf = elf::ElfStream::::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 mut get_u64_at_symbol = |sym: elf::symbol::Symbol| { + 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) + }; + + let sn_alloc_size = get_u64_at_symbol(sn_alloc_size); + let sn_alloc_align = get_u64_at_symbol(sn_alloc_align); + + // # 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(); +} 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..06a59f76bf753 --- /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 +#include + +/***************************************************/ +/*** 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 +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() { + // TODO: not sure what this is supposed to do + 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; + +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 + 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; +using Alloc = LocalAllocator; + +/// 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( + 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..dfdeeafecefd1 --- /dev/null +++ b/library/snmalloc-edp/tests/global_alloc.rs @@ -0,0 +1,108 @@ +use std::{alloc::{self, GlobalAlloc}, cell::Cell, ptr}; + +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 R, R>(f: F) -> R { + unsafe { + let mut allocator = std::mem::MaybeUninit::::uninit(); + sn_thread_init(allocator.as_mut_ptr()); + THREAD_ALLOC.set(allocator.as_mut_ptr()); + + let r = f(); + + THREAD_ALLOC.set(ptr::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(); + + sn_global_init(heap.start as _, heap.end as _); + } + + type AllocTestType = [u64; 20]; + + 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 { + with_thread_allocator(|| { + let p1 = System.alloc(alloc::Layout::new::()); + barrier.wait(); + let p2 = System.alloc(alloc::Layout::new::()); + tx.send((p1 as usize, p2 as usize)).unwrap(); + }) + }; + }); + + let (p1, p2) = unsafe { + with_thread_allocator(|| { + let p1 = System.alloc(alloc::Layout::new::()); + barrier.wait(); + let p2 = System.alloc(alloc::Layout::new::()); + (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/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"