Shuffle some more platform-intialization stuff into the HAL.

I feel like there must be a better abstraction to trigger these
initialization phases at the correct time, but at this point my goal is
just to get all the SEV-specific stuff out of `lib.rs` into the HAL to
unblock TDX enlightenment.

Bug: 350496083
Change-Id: I4d5b62644d0276d684b1aaab44f5422d89983c13

stage0/src/hal/mod.rs

@@ -189,13 +189,29 @@ impl<T: PortWrite> PortWriter<T> for Port<T> {
     }
 }
 
-/// Performs platform-specific initializations (for example, accepts guest
-/// memory)
+/// Platform-specific early initialization.
 ///
-/// What exactly is done depends on the platform (SEV or TDX).
+/// This sets up the bare minimum to get things going; for example, under
+/// SEV-ES and above, we set up the GHCB here, but nothing more.
 ///
-/// This code is run early in the startup process. You will have access to the
-/// boot allocator and logging, but _not_ the global heap allocator.
+/// This gets executed very soon after stage0 starts and comes with many
+/// restrictions:
+///   - You do not have access to logging.
+///   - You do not have access to heap allocator (BOOT_ALLOCATOR will still
+///     work).
+pub fn early_initialize_platform() {
+    #[cfg(feature = "sev")]
+    sev::early_initialize_platform();
+}
+
+/// Platform-specific intialization.
+///
+/// This gets executed after `early_initalize_platform()` and some other
+/// auxiliary services, such as logging, have been set up; the main purpose is
+/// to accept all guest memory so that we can set up a heap allocator.
+///
+/// This does mean you do not have access to the heap allocator (BOOT_ALLOCATOR
+/// will still work).
 pub fn initialize_platform(e820_table: &[BootE820Entry]) {
     #[cfg(feature = "sev")]
     sev::initialize_platform(e820_table)

stage0/src/hal/sev/mod.rs

@@ -29,7 +29,32 @@ use oak_linux_boot_params::BootE820Entry;
 use oak_sev_guest::msr::SevStatus;
 pub use port::*;
 
+use crate::BOOT_ALLOC;
+
+pub fn early_initialize_platform() {
+    // If we're under SEV-ES or SNP, we need a GHCB block for communication (SNP
+    // implies SEV-ES).
+    if crate::sev_status().contains(SevStatus::SEV_ES_ENABLED) {
+        crate::sev::GHCB_WRAPPER.init(&BOOT_ALLOC);
+    }
+    if crate::sev_status().contains(SevStatus::SEV_ENABLED) {
+        // Safety: This is safe for SEV-ES and SNP because we're using an originally
+        // supported mode of the Pentium 6: Write-protect, with MTRR enabled.
+        // If we get CPUID reads working, we may want to check that MTRR is
+        // supported, but only if we want to support very old processors.
+        // However, note that, this branch is only executed if
+        // we have encryption, and this wouldn't be true for very old processors.
+        unsafe {
+            crate::msr::MTRRDefType::write(
+                crate::msr::MTRRDefTypeFlags::MTRR_ENABLE,
+                crate::msr::MemoryType::WP,
+            );
+        }
+    }
+}
+
 pub fn initialize_platform(e820_table: &[BootE820Entry]) {
+    log::info!("Enabled SEV features: {:?}", crate::sev_status());
     if crate::sev_status().contains(SevStatus::SNP_ACTIVE) {
         dice_attestation::init_guest_message_encryptor()
             .expect("couldn't initialize guest message encryptor");

stage0/src/lib.rs

@@ -145,29 +145,9 @@ pub fn sev_status() -> SevStatus {
 ///   tables.
 pub fn rust64_start() -> ! {
     paging::init_page_table_refs();
-
-    // If we're under SEV-ES or SNP, we need a GHCB block for communication (SNP
-    // implies SEV-ES).
-    if sev_status().contains(SevStatus::SEV_ES_ENABLED) {
-        sev::GHCB_WRAPPER.init(&BOOT_ALLOC);
-    }
-
+    hal::early_initialize_platform();
     logging::init_logging();
     log::info!("starting...");
-    log::info!("Enabled SEV features: {:?}", sev_status());
-
-    if sev_status().contains(SevStatus::SEV_ENABLED) {
-        // Safety: This is safe for SEV-ES and SNP because we're using an originally
-        // supported mode of the Pentium 6: Write-protect, with MTRR enabled.
-        // If we get CPUID reads working, we may want to check that MTRR is
-        // supported, but only if we want to support very old processors.
-        // However, note that, this branch is only executed if
-        // we have encryption, and this wouldn't be true for very old processors.
-        unsafe {
-            msr::MTRRDefType::write(msr::MTRRDefTypeFlags::MTRR_ENABLE, msr::MemoryType::WP);
-        }
-    }
-
     // Safety: we assume there won't be any other hardware devices using the fw_cfg
     // IO ports.
     let mut fwcfg = unsafe { fw_cfg::FwCfg::new(&BOOT_ALLOC) }.expect("fw_cfg device not found!");