From b93bb28e1b39004fc2e3a194a131aacf95c774ae Mon Sep 17 00:00:00 2001
From: John Schock <joschock@microsoft.com>
Date: Wed, 25 Oct 2023 11:16:37 -0700
Subject: [PATCH] Update HelloWorldRustDxe to permit compiling for unit tests
 (#341)

## Description

Update HelloWorldRustDxe to allow compiling tests.

- [ ] Impacts functionality?
- [ ] Impacts security?
- [ ] Breaking change?
- [x] Includes tests?
  - adds a sample unit test for the HelloWorldRustDxe driver.
- [ ] Includes documentation?

## How This Was Tested

Cargo test executes sample unit test as expected.
Build for UEFI target and execution QEMU work as expected.

## Integration Instructions
N/A
---
 MsCorePkg/HelloWorldRustDxe/Cargo.toml  |  1 -
 MsCorePkg/HelloWorldRustDxe/src/main.rs | 66 ++++++++++++++++---------
 2 files changed, 43 insertions(+), 24 deletions(-)

diff --git a/MsCorePkg/HelloWorldRustDxe/Cargo.toml b/MsCorePkg/HelloWorldRustDxe/Cargo.toml
index 005c87f6ee..501c791ec3 100644
--- a/MsCorePkg/HelloWorldRustDxe/Cargo.toml
+++ b/MsCorePkg/HelloWorldRustDxe/Cargo.toml
@@ -6,7 +6,6 @@ edition = "2021"
 [[bin]]
 name = "HelloWorldRustDxe"
 path = "src/main.rs"
-test = false
 
 [dependencies]
 r-efi = {workspace=true}
diff --git a/MsCorePkg/HelloWorldRustDxe/src/main.rs b/MsCorePkg/HelloWorldRustDxe/src/main.rs
index 4cff82374c..e925e2c0e8 100644
--- a/MsCorePkg/HelloWorldRustDxe/src/main.rs
+++ b/MsCorePkg/HelloWorldRustDxe/src/main.rs
@@ -8,39 +8,59 @@
 //!
 //! SPDX-License-Identifier: BSD-2-Clause-Patent
 //!
+
 #![no_std]
-#![no_main]
 #![allow(non_snake_case)]
+// UEFI requires "no_main" and a panic handler. To facilitate unit tests, gate "no_main", UEFI entry point, and panic
+// handler behind target_os configuration.
+#![cfg_attr(target_os = "uefi", no_main)]
+#[cfg(target_os = "uefi")]
+mod uefi_entry {
+  extern crate alloc;
+
+  use alloc::vec;
+  use core::panic::PanicInfo;
+  use r_efi::efi::Status;
+  use rust_advanced_logger_dxe::{debugln, init_debug, DEBUG_INFO};
 
-extern crate alloc;
+  #[no_mangle]
+  pub extern "efiapi" fn efi_main(
+    _image_handle: *const core::ffi::c_void,
+    _system_table: *const r_efi::system::SystemTable,
+  ) -> u64 {
+    rust_boot_services_allocator_dxe::GLOBAL_ALLOCATOR.init(unsafe { (*_system_table).boot_services });
+    init_debug(unsafe { (*_system_table).boot_services });
 
-use alloc::vec;
-use core::panic::PanicInfo;
-use r_efi::efi::Status;
-use rust_advanced_logger_dxe::{debugln, init_debug, DEBUG_INFO};
+    debugln!(DEBUG_INFO, "Hello, World. This is Rust in UEFI.");
 
-#[no_mangle]
-pub extern "efiapi" fn efi_main(
-  _image_handle: *const core::ffi::c_void,
-  _system_table: *const r_efi::system::SystemTable,
-) -> u64 {
-  rust_boot_services_allocator_dxe::GLOBAL_ALLOCATOR.init(unsafe { (*_system_table).boot_services });
-  init_debug(unsafe { (*_system_table).boot_services });
+    debugln!(DEBUG_INFO, "file: {:} line: {:} as hex: {:x}", file!(), line!(), line!());
 
-  debugln!(DEBUG_INFO, "Hello, World. This is Rust in UEFI.");
+    let mut foo = vec!["asdf", "xyzpdq", "abcdefg", "thxyzb"];
 
-  debugln!(DEBUG_INFO, "file: {:} line: {:} as hex: {:x}", file!(), line!(), line!());
+    debugln!(DEBUG_INFO, "Unsorted vec: {:?}", foo);
+    foo.sort();
+    debugln!(DEBUG_INFO, "Sorted vec: {:?}", foo);
 
-  let mut foo = vec!["asdf", "xyzpdq", "abcdefg", "thxyzb"];
+    Status::SUCCESS.as_usize() as u64
+  }
 
-  debugln!(DEBUG_INFO, "Unsorted vec: {:?}", foo);
-  foo.sort();
-  debugln!(DEBUG_INFO, "Sorted vec: {:?}", foo);
+  #[panic_handler]
+  fn panic(_info: &PanicInfo) -> ! {
+    loop {}
+  }
+}
 
-  Status::SUCCESS.as_usize() as u64
+// For non-UEFI targets (e.g. compiling for unit test or clippy), supply a "main" function.
+#[cfg(not(target_os = "uefi"))]
+fn main() {
+  //do nothing.
 }
 
-#[panic_handler]
-fn panic(_info: &PanicInfo) -> ! {
-  loop {}
+#[cfg(test)]
+mod test {
+
+  #[test]
+  fn sample_test() {
+    assert_eq!(1 + 1, 2);
+  }
 }