Review comments addressed and cleanup

snmalloc-edp/build.rs

@@ -1,3 +1,5 @@
+use elf::ElfStream;
+use elf::endian::LittleEndian;
 use std::fs::{DirEntry, File};
 use std::path::{Path, PathBuf};
 
@@ -35,6 +37,7 @@ fn main() {
     assert!(ar.status().unwrap().success());
 
     // # Read the symbols from the shim ELF object
+    assert_eq!(files_in_dir(&objs).count(), 1);
     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();
@@ -50,19 +53,19 @@ fn main() {
     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);
+    let sn_alloc_size = get_u64_at_symbol(sn_alloc_size, &mut elf);
+    let sn_alloc_align = get_u64_at_symbol(sn_alloc_align, &mut elf);
 
     // # 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();
 }
+
+fn get_u64_at_symbol(sym: elf::symbol::Symbol, elf: &mut ElfStream<LittleEndian, File>) -> u64 {
+    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)
+}

snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp

@@ -34,7 +34,7 @@
 extern "C" size_t get_tcs_addr();
 
 // from Rust std
-extern "C" void __rust_print_err(const char *m, size_t s);
+extern "C" void __rust_print_err(const char* m, size_t s);
 extern "C" [[noreturn]] void __rust_abort();
 
 /*******************************************************/
@@ -46,7 +46,7 @@ extern "C" [[noreturn]] void abort() __THROW {
     __rust_abort();
 }
 
-// definition needs to match GNU header
+// definition needs to match GNU header and will not return an actual errno
 extern "C" inline int * __attribute_const__ __errno_location (void) __THROW {
     static int errno;
     return &errno;
@@ -65,15 +65,15 @@ extern "C" inline int * __attribute_const__ __errno_location (void) __THROW {
 
 namespace snmalloc {
 void register_clean_up() {
-    // TODO: not sure what this is supposed to do
+    // Unused on SGX
     abort();
 }
 
 class EdpErrorHandler {
   public:
     static void print_stack_trace() {}
 
-    [[noreturn]] static void error(const char *const str) {
+    [[noreturn]] static void error(const char* const str) {
         __rust_print_err(str, strlen(str));
         abort();
     }

snmalloc-edp/tests/global_alloc.rs

@@ -1,4 +1,5 @@
 use std::{alloc::{self, GlobalAlloc}, cell::Cell, ptr};
+use std::ptr::null_mut;
 
 use snmalloc_edp::*;
 
@@ -49,7 +50,7 @@ unsafe fn with_thread_allocator<F: FnOnce() -> R, R>(f: F) -> R {
 
         let r = f();
 
-        THREAD_ALLOC.set(ptr::null_mut());
+        THREAD_ALLOC.set(null_mut());
         sn_thread_cleanup(allocator.as_mut_ptr());
 
         r
@@ -64,13 +65,14 @@ fn test() {
         #[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 _);
+            // allocate a dummy heap
+            let heap = (*Box::into_raw(vec![Page([0; 4096]); 100].into_boxed_slice())).as_mut_ptr_range();
+            let heap_size = heap.end as usize - heap.start as usize;
+            sn_global_init(heap.start as _, heap_size);
     }
 
-    type AllocTestType = [u64; 20];
+    const TEST_ARRAY_SIZE :usize = 20;
+    type AllocTestType = [u64; TEST_ARRAY_SIZE];
 
     let barrier = std::sync::Barrier::new(2);
 
@@ -79,10 +81,17 @@ fn test() {
         let barrier = &barrier;
         s.spawn(move || {
             unsafe {
+                // Initialize thread allocator and perform alloc/alloc_zeroed
                 with_thread_allocator(|| {
                     let p1 = System.alloc(alloc::Layout::new::<AllocTestType>());
                     barrier.wait();
-                    let p2 = System.alloc(alloc::Layout::new::<AllocTestType>());
+                    let p2 = System.alloc_zeroed(alloc::Layout::new::<AllocTestType>());
+
+                    let p2_slice = std::slice::from_raw_parts_mut(p2, TEST_ARRAY_SIZE);
+                    for i in 0..TEST_ARRAY_SIZE {
+                        assert_eq!(p2_slice[i], 0);
+                    }
+
                     tx.send((p1 as usize, p2 as usize)).unwrap();
                 })
             };
@@ -93,16 +102,23 @@ fn test() {
                 let p1 = System.alloc(alloc::Layout::new::<AllocTestType>());
                 barrier.wait();
                 let p2 = System.alloc(alloc::Layout::new::<AllocTestType>());
+
+                // Test realloc
+                let p3 = System.realloc(p1, alloc::Layout::new::<AllocTestType>(), TEST_ARRAY_SIZE * 2);
+                assert_ne!(p3, p2);
+
                 (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);
+
     })
 
 }