From 8841b4d243047ab6464fda50dd54049a480fc4fa Mon Sep 17 00:00:00 2001 From: Jethro Beekman Date: Thu, 2 May 2024 23:21:19 +0200 Subject: [PATCH] Glue code for using snmalloc in EDP --- .github/workflows/build.yml | 22 ++- .gitmodules | 3 + Cargo.lock | 11 +- Cargo.toml | 1 + snmalloc-edp/CMakeLists.txt | 10 ++ snmalloc-edp/Cargo.toml | 11 ++ snmalloc-edp/build.rs | 5 + snmalloc-edp/snmalloc | 1 + snmalloc-edp/src/lib.rs | 24 +++ snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp | 187 ++++++++++++++++++++ snmalloc-edp/tests/global_alloc.rs | 111 ++++++++++++ 11 files changed, 383 insertions(+), 3 deletions(-) create mode 100644 .gitmodules create mode 100644 snmalloc-edp/CMakeLists.txt create mode 100644 snmalloc-edp/Cargo.toml create mode 100644 snmalloc-edp/build.rs create mode 160000 snmalloc-edp/snmalloc create mode 100644 snmalloc-edp/src/lib.rs create mode 100644 snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp create mode 100644 snmalloc-edp/tests/global_alloc.rs diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index d9e63fc9..b90e7ac1 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -16,7 +16,9 @@ env: CARGO_INCREMENTAL: 0 CARGO_NET_RETRY: 10 CFLAGS_x86_64_fortanix_unknown_sgx: "-isystem/usr/include/x86_64-linux-gnu -mlvi-hardening -mllvm -x86-experimental-lvi-inline-asm-hardening" + # CXXFLAGS is set below CC_x86_64_fortanix_unknown_sgx: clang-11 + CXX_x86_64_fortanix_unknown_sgx: clang++-11 jobs: test: @@ -48,7 +50,7 @@ jobs: rustup update - name: Cargo test --all --exclude sgxs-loaders - run: cargo test --verbose --locked --all --exclude sgxs-loaders --exclude async-usercalls && [ "$(echo $(nm -D target/debug/sgx-detect|grep __vdso_sgx_enter_enclave))" = "w __vdso_sgx_enter_enclave" ] + run: true || ( cargo test --verbose --locked --all --exclude sgxs-loaders --exclude async-usercalls --exclude snmalloc-edp && [ "$(echo $(nm -D target/debug/sgx-detect|grep __vdso_sgx_enter_enclave))" = "w __vdso_sgx_enter_enclave" ] ) - name: cargo test -p async-usercalls --target x86_64-fortanix-unknown-sgx --no-run run: cargo +nightly test --verbose --locked -p async-usercalls --target x86_64-fortanix-unknown-sgx --no-run @@ -96,6 +98,24 @@ jobs: - name: Build em-app, get-certificate for x86_64-fortanix-unknown-sgx run: cargo build --verbose --locked -p em-app -p get-certificate --target=x86_64-fortanix-unknown-sgx + - name: Build snmalloc-edp + run: | + 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" + cargo build --verbose --locked -p snmalloc-edp --target=x86_64-fortanix-unknown-sgx + - name: Generate API docs run: ./doc/generate-api-docs.sh diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 00000000..905063cd --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "snmalloc-edp/snmalloc"] + path = snmalloc-edp/snmalloc + url = https://github.com/microsoft/snmalloc diff --git a/Cargo.lock b/Cargo.lock index e33c816c..369ad6e7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -420,9 +420,9 @@ dependencies = [ [[package]] name = "cmake" -version = "0.1.44" +version = "0.1.50" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e56268c17a6248366d66d4a47a3381369d068cce8409bb1716ed77ea32163bb" +checksum = "a31c789563b815f77f4250caee12365734369f942439b7defd71e18a48197130" dependencies = [ "cc", ] @@ -3366,6 +3366,13 @@ version = "1.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ecab6c735a6bb4139c0caafd0cc3635748bbb3acf4550e8138122099251f309" +[[package]] +name = "snmalloc-edp" +version = "0.1.0" +dependencies = [ + "cmake", +] + [[package]] name = "socket2" version = "0.4.2" diff --git a/Cargo.toml b/Cargo.toml index 069a108c..0d1e379f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,7 @@ members = [ "intel-sgx/sgxs", "ipc-queue", "rs-libc", + "snmalloc-edp", ] exclude = [ "examples/backtrace_panic", diff --git a/snmalloc-edp/CMakeLists.txt b/snmalloc-edp/CMakeLists.txt new file mode 100644 index 00000000..b87eee49 --- /dev/null +++ b/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/snmalloc-edp/Cargo.toml b/snmalloc-edp/Cargo.toml new file mode 100644 index 00000000..053d5adb --- /dev/null +++ b/snmalloc-edp/Cargo.toml @@ -0,0 +1,11 @@ +[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] +cmake = "0.1.50" diff --git a/snmalloc-edp/build.rs b/snmalloc-edp/build.rs new file mode 100644 index 00000000..2aeacf02 --- /dev/null +++ b/snmalloc-edp/build.rs @@ -0,0 +1,5 @@ +fn main() { + let mut dst = cmake::build("."); + dst.push("build"); + println!("cargo:rustc-link-search=native={}", dst.display()); +} diff --git a/snmalloc-edp/snmalloc b/snmalloc-edp/snmalloc new file mode 160000 index 00000000..dc126888 --- /dev/null +++ b/snmalloc-edp/snmalloc @@ -0,0 +1 @@ +Subproject commit dc1268886a5d49d38a54e5d1402b5924a71fee0b diff --git a/snmalloc-edp/src/lib.rs b/snmalloc-edp/src/lib.rs new file mode 100644 index 00000000..4840519f --- /dev/null +++ b/snmalloc-edp/src/lib.rs @@ -0,0 +1,24 @@ +#![no_std] + +use core::ffi::c_void; + +#[repr(C)] +pub struct Alloc { + _data: [u8; 0], + _marker: + core::marker::PhantomData<(*mut u8, core::marker::PhantomPinned)>, +} + +#[link(name = "snmalloc-edp", kind = "static")] +extern { + pub fn sn_global_init(heap_start_address: *mut c_void, heap_end_address: *mut c_void); + pub fn sn_thread_init(allocator: *mut Alloc); + pub fn sn_thread_cleanup(allocator: *mut Alloc); + pub static sn_alloc_size: usize; + pub static sn_alloc_align: usize; + + 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/snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp b/snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp new file mode 100644 index 00000000..31324fc0 --- /dev/null +++ b/snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp @@ -0,0 +1,187 @@ +// 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 +#define OPEN_ENCLAVE +// needed for openenclave header: +#define OE_OK 0 + +#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: + 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() { + long long unsigned int result = 0; + while (_rdrand64_step(&result) != 1) + ; + return result; + } + + 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. +// TODO: this function shouldn't need the addresses passed in, these can be +// obtained from the HEAP_* symbols +extern "C" void sn_global_init(void *heap_start_address, + void *heap_end_address) { + size_t _max_heap_size = + static_cast(static_cast(heap_end_address) - + static_cast(heap_start_address)); + + Globals::init(nullptr, heap_start_address, _max_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/snmalloc-edp/tests/global_alloc.rs b/snmalloc-edp/tests/global_alloc.rs new file mode 100644 index 00000000..f09e5a7b --- /dev/null +++ b/snmalloc-edp/tests/global_alloc.rs @@ -0,0 +1,111 @@ +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 layout = alloc::Layout::from_size_align(sn_alloc_size, sn_alloc_align).unwrap(); + // TODO: bootstrap the thread-local allocator allocation a different way + let allocator = alloc::alloc(layout) as *mut Alloc; + sn_thread_init(allocator); + THREAD_ALLOC.set(allocator); + + let r = f(); + + THREAD_ALLOC.set(ptr::null_mut()); + sn_thread_cleanup(allocator); + alloc::dealloc(allocator as _, layout); + + 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); + }) + +}