From 106b8c3a645a9808d7a7e57fdd3b180c9da0db23 Mon Sep 17 00:00:00 2001
From: aditijannu <aditi.jannu@fortanix.com>
Date: Tue, 25 Jun 2024 16:37:44 +0200
Subject: [PATCH] Review comments addressed and cleanup

---
 snmalloc-edp/build.rs                       | 21 ++++++++++++---------
 snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp |  8 ++++----
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/snmalloc-edp/build.rs b/snmalloc-edp/build.rs
index 6e891726..ff263de2 100644
--- a/snmalloc-edp/build.rs
+++ b/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,15 +53,8 @@ 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);
@@ -66,3 +62,10 @@ fn main() {
     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)
+}
diff --git a/snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp b/snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp
index 06a59f76..ee6850a2 100644
--- a/snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp
+++ b/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,7 +65,7 @@ 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();
 }
 
@@ -73,7 +73,7 @@ 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();
     }