Glue code for using snmalloc in EDP

.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: 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 test --no-run --verbose --locked -p snmalloc-edp --target=x86_64-fortanix-unknown-sgx
+
     - name: Generate API docs
       run: ./doc/generate-api-docs.sh
 

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "snmalloc-edp/snmalloc"]
+	path = snmalloc-edp/snmalloc
+	url = https://github.com/microsoft/snmalloc

Cargo.toml

@@ -33,6 +33,7 @@ members = [
     "intel-sgx/sgxs",
     "ipc-queue",
     "rs-libc",
+    "snmalloc-edp",
 ]
 exclude = [
     "examples/backtrace_panic",

intel-sgx/ias/src/api.rs

@@ -193,7 +193,7 @@ pub struct VerifiedSig {}
 impl VerificationType for VerifiedSig {}
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
-pub(crate) struct Unverified {}
+pub enum Unverified {}
 impl VerificationType for Unverified {}
 
 trait SafeToDeserializeInto {}

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)

snmalloc-edp/Cargo.toml

@@ -0,0 +1,13 @@
+[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"

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<Item = DirEntry> {
+    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::<elf::endian::LittleEndian, _>::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();
+}

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();
+    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;
+}