SDK: Only define global alloc in code output by entrypoint macro (#4733)

* Permit removal of the global alloc declaration in the SDK

We shouldn't declare a global alloc in unit tests. (Ours is currently interopable with linux, which is why nothing breaks right now)

* Revert "Permit removal of the global alloc declaration in the SDK"

This reverts commit d4fdf80.

* Move alloc and logger delcarations into entrypoint macro

* make change non breaking

* re-export the alloc struct from the sdk, so the entrypoint macro can use it

* fix typo specifying dep

* Add workspace wide setting to disable default features for oak_enclave_runtime_support, as otherwise it is ignored

* Make oak_enclave_runtime_support/src/heap.rs clippy compliant since it's now public and parsed

* Remove empty allocator init function

Author: jul-sh

Cargo.toml

@@ -103,7 +103,7 @@ oak_containers_sdk = { path = "./oak_containers_sdk" }
 oak_core = { path = "./oak_core" }
 oak_crypto = { path = "./oak_crypto" }
 oak_dice = { path = "./oak_dice" }
-oak_enclave_runtime_support = { path = "./oak_enclave_runtime_support" }
+oak_enclave_runtime_support = { path = "./oak_enclave_runtime_support", default-features = false }
 oak_functions_abi = { path = "./oak_functions_abi" }
 oak_functions_client = { path = "./oak_functions_client" }
 oak_functions_launcher = { path = "./oak_functions_launcher" }

oak_enclave_runtime_support/Cargo.toml

@@ -5,6 +5,10 @@ authors = ["Andri Saar <andrisaar@google.com>"]
 edition = "2021"
 license = "Apache-2.0"
 
+[features]
+default = ["global_allocator"]
+global_allocator = []
+
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 [dependencies]
 libm = "*"

oak_enclave_runtime_support/README.md

@@ -6,4 +6,4 @@
 
 Runtime support library for applications built for Oak Restricted Kernel.
 
-For now, the runtime support library only sets up the global heap allocator.
+For now, the runtime support library provides a global heap allocator.

oak_enclave_runtime_support/src/heap.rs

@@ -79,14 +79,18 @@ impl GrowableHeap {
         }
     }
 
-    pub unsafe fn init(&mut self) {}
-
+    #[allow(clippy::result_unit_err)]
     pub fn allocate(&mut self, layout: Layout) -> Result<NonNull<u8>, ()> {
         self.heap
             .allocate(layout)
             .ok_or_else(|| log::error!("failed to allocate memory with layout: {:?}", layout))
     }
 
+    /// # Safety
+    ///
+    ///  - `ptr` must denote a memory block previously allocated via `self`.
+    ///  - The memory block must have been allocated with the same alignment ([`Layout::align`]) as
+    ///    `align`.
     pub unsafe fn deallocate(&mut self, ptr: NonNull<u8>, align: usize) {
         self.heap.deallocate(ptr, align)
     }

oak_enclave_runtime_support/src/lib.rs

@@ -18,14 +18,12 @@
 
 use heap::LockedGrowableHeap;
 
-mod heap;
+pub mod heap;
 mod libm;
 
-#[cfg_attr(not(test), global_allocator)]
+#[cfg(feature = "global_allocator")]
+#[global_allocator]
 static ALLOCATOR: LockedGrowableHeap = LockedGrowableHeap::empty();
 
-pub fn init() {
-    unsafe {
-        ALLOCATOR.lock().init();
-    }
-}
+#[deprecated(note = "please make use of `oak_restricted_kernel_sdk::entrypoint` instead.")]
+pub fn init() {}

oak_restricted_kernel_sdk/Cargo.toml

@@ -17,7 +17,7 @@ oak_crypto = { workspace = true }
 oak_core = { workspace = true }
 oak_dice = { workspace = true }
 oak_restricted_kernel_interface = { workspace = true }
-oak_enclave_runtime_support = { workspace = true }
+oak_enclave_runtime_support = { default-features = false, workspace = true }
 oak_restricted_kernel_sdk_proc_macro = { workspace = true }
 oak_restricted_kernel_dice = { workspace = true, optional = true }
 oak_stage0_dice = { workspace = true, optional = true }

oak_restricted_kernel_sdk/src/lib.rs

@@ -22,20 +22,16 @@ mod channel;
 mod dice;
 mod logging;
 
+pub mod utils {
+    pub use oak_core::*;
+    pub use oak_enclave_runtime_support::heap;
+}
+
 pub use channel::*;
 pub use dice::*;
 pub use logging::StderrLogger;
-use logging::STDERR_LOGGER;
-pub use oak_core as utils;
 pub use oak_restricted_kernel_sdk_proc_macro::entrypoint;
 
-/// Initialization function that sets up the allocator and logger.
-pub fn init(log_level: log::LevelFilter) {
-    log::set_logger(&STDERR_LOGGER).expect("failed to set logger");
-    log::set_max_level(log_level);
-    oak_enclave_runtime_support::init();
-}
-
 pub fn alloc_error_handler(layout: ::core::alloc::Layout) -> ! {
     panic!("error allocating memory: {:#?}", layout);
 }

oak_restricted_kernel_sdk/src/logging.rs

@@ -18,8 +18,6 @@ use core::fmt::Write;
 
 use oak_restricted_kernel_interface::syscall::{fsync, write};
 
-pub static STDERR_LOGGER: StderrLogger = StderrLogger {};
-
 struct Stderr {}
 
 impl Stderr {

oak_restricted_kernel_sdk_proc_macro/src/lib.rs

@@ -62,11 +62,15 @@ fn process_entry_fn(entry_fn: ItemFn) -> TokenStream {
     let generated = quote! {
         #entry_fn
 
+        #[global_allocator]
+        static ALLOCATOR: oak_restricted_kernel_sdk::utils::heap::LockedGrowableHeap = oak_restricted_kernel_sdk::utils::heap::LockedGrowableHeap::empty();
+
         static LOGGER: oak_restricted_kernel_sdk::StderrLogger = oak_restricted_kernel_sdk::StderrLogger {};
 
         #[no_mangle]
         fn _start() -> ! {
-            oak_restricted_kernel_sdk::init(log::LevelFilter::Debug);
+            log::set_logger(&LOGGER).expect("failed to set logger");
+            log::set_max_level(log::LevelFilter::Debug);
             log::info!("In main!");
             #entry_fn_name();
         }