From 29174fb2ee7d846c858fe56134703c17e37f61b2 Mon Sep 17 00:00:00 2001
From: jul-sh <julsh@google.com>
Date: Mon, 25 Mar 2024 14:34:21 -0400
Subject: [PATCH] Differentiate orcherstrator logs (#4953)

Now orchestrator logs will be prefixed, just as stage0 and kernel logs are
---
 enclave_apps/oak_orchestrator/src/main.rs | 58 +++++++++++++++++++++--
 oak_restricted_kernel_sdk/src/utils.rs    |  4 +-
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/enclave_apps/oak_orchestrator/src/main.rs b/enclave_apps/oak_orchestrator/src/main.rs
index c087dd44000..b9874fa993e 100644
--- a/enclave_apps/oak_orchestrator/src/main.rs
+++ b/enclave_apps/oak_orchestrator/src/main.rs
@@ -20,13 +20,66 @@
 
 extern crate alloc;
 
+use core::fmt::Write;
+
 use oak_dice::evidence::Stage0DiceData;
 use oak_restricted_kernel_interface::{syscall, DERIVED_KEY_FD, DICE_DATA_FD};
 use oak_restricted_kernel_orchestrator::AttestedApp;
-use oak_restricted_kernel_sdk::{channel::FileDescriptorChannel, entrypoint};
+use oak_restricted_kernel_sdk::channel::FileDescriptorChannel;
 use zerocopy::{AsBytes, FromZeroes};
 use zeroize::Zeroize;
 
+struct OrchestratorLogger {}
+
+impl log::Log for OrchestratorLogger {
+    fn enabled(&self, _metadata: &log::Metadata) -> bool {
+        true
+    }
+
+    fn log(&self, record: &log::Record) {
+        writeln!(
+            oak_restricted_kernel_sdk::utils::Stderr {},
+            "orchestrator {}: {}",
+            record.level(),
+            record.args()
+        )
+        .unwrap();
+    }
+
+    fn flush(&self) {
+        oak_restricted_kernel_sdk::utils::Stderr::flush();
+    }
+}
+
+#[global_allocator]
+static ALLOCATOR: oak_restricted_kernel_sdk::utils::heap::LockedGrowableHeap =
+    oak_restricted_kernel_sdk::utils::heap::LockedGrowableHeap::empty();
+
+static LOGGER: OrchestratorLogger = OrchestratorLogger {};
+
+// The orchestrator uses a custom logging implementation, hence the
+// #[oak_restricted_kernel_sdk::entrypoint] is not used. The allocator,
+// handlers, etc are declared explicitly.
+#[no_mangle]
+fn _start() -> ! {
+    oak_restricted_kernel_sdk::utils::log::set_logger(&LOGGER).expect("failed to set logger");
+    oak_restricted_kernel_sdk::utils::log::set_max_level(
+        oak_restricted_kernel_sdk::utils::log::LevelFilter::Debug,
+    );
+    entrypoint()
+}
+
+#[alloc_error_handler]
+fn out_of_memory(layout: ::core::alloc::Layout) -> ! {
+    panic!("error allocating memory in orchestrator: {:#?}", layout);
+}
+
+#[panic_handler]
+fn panic(info: &core::panic::PanicInfo) -> ! {
+    log::error!("orchestrator PANIC: {}", info);
+    oak_restricted_kernel_interface::syscall::exit(-1);
+}
+
 fn read_stage0_dice_data() -> Stage0DiceData {
     let mut result = Stage0DiceData::new_zeroed();
     let buffer = result.as_bytes_mut();
@@ -35,8 +88,7 @@ fn read_stage0_dice_data() -> Stage0DiceData {
     result
 }
 
-#[entrypoint]
-fn start() -> ! {
+fn entrypoint() -> ! {
     let mut attested_app = {
         let stage0_dice_data = read_stage0_dice_data();
         let channel = FileDescriptorChannel::default();
diff --git a/oak_restricted_kernel_sdk/src/utils.rs b/oak_restricted_kernel_sdk/src/utils.rs
index 7ee2b16a231..90a7667b307 100644
--- a/oak_restricted_kernel_sdk/src/utils.rs
+++ b/oak_restricted_kernel_sdk/src/utils.rs
@@ -23,12 +23,12 @@ pub use oak_core::*;
 pub use oak_enclave_runtime_support::heap;
 use oak_restricted_kernel_interface::syscall::{fsync, write};
 
-struct Stderr {}
+pub struct Stderr {}
 
 impl Stderr {
     const STDERR_FD: i32 = 2;
 
-    fn flush() {
+    pub fn flush() {
         fsync(Self::STDERR_FD).unwrap();
     }
 }