Move page state and encryptedness operations to stage0 HAL

Bug: 350496083
Change-Id: I58bd9caa16ec03a3d92e1d6ffea8a7335f96fb8d

stage0/src/hal/mod.rs

@@ -21,11 +21,14 @@ mod sev;
 use core::{arch::x86_64::CpuidResult, marker::PhantomData, mem::size_of};
 
 use oak_linux_boot_params::BootE820Entry;
-use oak_sev_guest::io::{IoPortFactory, PortReader, PortWriter};
+use oak_sev_guest::{
+    io::{IoPortFactory, PortReader, PortWriter},
+    msr::PageAssignment,
+};
 use oak_sev_snp_attestation_report::{AttestationReport, REPORT_DATA_SIZE};
 use oak_stage0_dice::DerivedKey;
 use x86_64::{
-    structures::paging::{PageSize, Size4KiB},
+    structures::paging::{Page, PageSize, Size4KiB},
     PhysAddr,
 };
 
@@ -143,6 +146,8 @@ pub use x86_64::structures::port::PortRead;
 #[cfg(not(feature = "sev"))]
 pub use x86_64::structures::port::PortWrite;
 
+use crate::paging::PageEncryption;
+
 impl<'a, T> IoPortFactory<'a, T, Port<T>, Port<T>> for PortFactory
 where
     T: PortRead + PortWrite + 'a,
@@ -251,3 +256,37 @@ pub fn get_derived_key() -> Result<DerivedKey, &'static str> {
     #[cfg(not(feature = "sev"))]
     return base::get_derived_key();
 }
+
+/// Ask for the page state to be changed by the hypervisor.
+pub fn change_page_state(page: Page<Size4KiB>, state: PageAssignment) {
+    #[cfg(feature = "sev")]
+    sev::change_page_state(page, state).unwrap();
+}
+
+/// Validate one page of memory.
+///
+/// This operation is required for SEV after going from a SHARED state to a
+/// PRIVATE state.
+pub fn revalidate_page(page: Page<Size4KiB>) {
+    #[cfg(feature = "sev")]
+    sev::revalidate_page(page).unwrap();
+}
+
+/// Mask to use in the page tables for the given encrypion state.
+///
+/// SEV and TDX have opposite behaviours: for SEV, encrypted pages are marked;
+/// for TDX, unencrypted pages are marked.
+pub fn page_table_mask(encryption_state: PageEncryption) -> u64 {
+    #[cfg(feature = "sev")]
+    return sev::page_table_mask(encryption_state);
+    #[cfg(not(feature = "sev"))]
+    return 0;
+}
+
+/// Encrypted/shared bit mask irrespective of its semantics.
+pub fn encrypted() -> u64 {
+    #[cfg(feature = "sev")]
+    return sev::encrypted();
+    #[cfg(not(feature = "sev"))]
+    return 0;
+}

stage0/src/hal/sev/accept_memory.rs

@@ -18,7 +18,7 @@ use core::sync::atomic::{AtomicUsize, Ordering};
 use oak_linux_boot_params::{BootE820Entry, E820EntryType};
 use oak_sev_guest::{
     instructions::{pvalidate, InstructionError, PageSize as SevPageSize, Validation},
-    msr::PageAssignment,
+    msr::{change_snp_page_state, PageAssignment, SevStatus, SnpPageStateChangeRequest},
 };
 use x86_64::{
     instructions::tlb,
@@ -408,3 +408,24 @@ pub fn validate_memory(e820_table: &[BootE820Entry]) {
         counters::ERROR_FAIL_SIZE_MISMATCH.load(Ordering::SeqCst)
     );
 }
+
+pub fn change_page_state(page: Page<Size4KiB>, state: PageAssignment) -> Result<(), &'static str> {
+    if crate::sev_status().contains(SevStatus::SNP_ACTIVE) {
+        let request = SnpPageStateChangeRequest::new(page.start_address().as_u64() as usize, state)
+            .expect("invalid address for page location");
+        change_snp_page_state(request)?;
+    }
+    Ok(())
+}
+
+pub fn revalidate_page(page: Page<Size4KiB>) -> Result<(), &'static str> {
+    if crate::sev_status().contains(SevStatus::SEV_ENABLED) {
+        let counter = AtomicUsize::new(0);
+        if let Err(err) = page.pvalidate(&counter) {
+            if err != InstructionError::ValidationStatusNotUpdated {
+                return Err("shared page revalidation failed");
+            }
+        }
+    }
+    Ok(())
+}

stage0/src/hal/sev/mod.rs

@@ -21,6 +21,7 @@ mod mmio;
 mod msr;
 mod port;
 
+pub use accept_memory::*;
 pub use cpuid::*;
 pub use dice_attestation::*;
 pub use mmio::*;
@@ -29,7 +30,7 @@ use oak_linux_boot_params::BootE820Entry;
 use oak_sev_guest::msr::SevStatus;
 pub use port::*;
 
-use crate::BOOT_ALLOC;
+use crate::{paging::PageEncryption, BOOT_ALLOC};
 
 pub fn early_initialize_platform() {
     // If we're under SEV-ES or SNP, we need a GHCB block for communication (SNP
@@ -61,3 +62,24 @@ pub fn initialize_platform(e820_table: &[BootE820Entry]) {
         accept_memory::validate_memory(e820_table)
     }
 }
+
+/// Returns the location of the ENCRYPTED bit when running under AMD SEV.
+pub(crate) fn encrypted() -> u64 {
+    #[no_mangle]
+    static mut ENCRYPTED: u64 = 0;
+
+    // Safety: we don't allow mutation and this is initialized in the bootstrap
+    // assembly.
+    unsafe { ENCRYPTED }
+}
+
+pub fn page_table_mask(encryption_state: PageEncryption) -> u64 {
+    if crate::sev_status().contains(SevStatus::SEV_ENABLED) {
+        match encryption_state {
+            PageEncryption::Encrypted => encrypted(),
+            PageEncryption::Unencrypted => 0,
+        }
+    } else {
+        0
+    }
+}

stage0/src/paging.rs

@@ -21,12 +21,13 @@ use core::{
 };
 
 use oak_core::sync::OnceCell;
+use oak_sev_guest::msr::PageAssignment;
 use spinning_top::Spinlock;
 use x86_64::{
     instructions::tlb::flush_all,
     structures::paging::{
         page_table::{PageTableEntry as BasePageTableEntry, PageTableFlags},
-        PageSize, PageTable as BasePageTable, PageTableIndex, Size2MiB, Size4KiB,
+        Page, PageSize, PageTable as BasePageTable, PageTableIndex, Size2MiB, Size4KiB,
     },
     PhysAddr,
 };
@@ -123,16 +124,13 @@ impl PageTableEntry {
     /// Map the entry to the specified address with the specified flags and
     /// encryption state.
     pub fn set_address(&mut self, addr: PhysAddr, flags: PageTableFlags, state: PageEncryption) {
-        let addr = match state {
-            PageEncryption::Encrypted => PhysAddr::new(addr.as_u64() | encrypted()),
-            PageEncryption::Unencrypted => addr,
-        };
+        let addr = PhysAddr::new(addr.as_u64() | crate::hal::page_table_mask(state));
         self.0.set_addr(addr, flags);
     }
 
     /// Returns the physical address mapped by this entry. May be zero.
     pub fn address(&self) -> PhysAddr {
-        PhysAddr::new(self.0.addr().as_u64() & !encrypted())
+        PhysAddr::new(self.0.addr().as_u64() & !crate::hal::encrypted())
     }
 
     /// Returns whether the entry is zero.
@@ -173,16 +171,6 @@ pub enum PageEncryption {
     Unencrypted,
 }
 
-/// Returns the location of the ENCRYPTED bit when running under AMD SEV.
-fn encrypted() -> u64 {
-    #[no_mangle]
-    static mut ENCRYPTED: u64 = 0;
-
-    // Safety: we don't allow mutation and this is initialized in the bootstrap
-    // assembly.
-    unsafe { ENCRYPTED }
-}
-
 /// Initialises the page table references.
 pub fn init_page_table_refs() {
     // Safety: accessing the mutable statics here is safe since we only do it once
@@ -253,3 +241,48 @@ pub fn remap_first_huge_page() {
 
     flush_all();
 }
+
+/// Shares a single 4KiB page with the hypervisor.
+pub fn share_page(page: Page<Size4KiB>) {
+    let page_start = page.start_address().as_u64();
+    // Only the first 2MiB is mapped as 4KiB pages, so make sure we fall in that
+    // range.
+    assert!(page_start < Size2MiB::SIZE);
+    // Remove the ENCRYPTED bit from the entry that maps the page.
+    {
+        let mut page_tables = crate::paging::PAGE_TABLE_REFS.get().unwrap().lock();
+        let pt = &mut page_tables.pt_0;
+        pt[page.p1_index()].set_address(
+            PhysAddr::new(page_start),
+            PageTableFlags::PRESENT | PageTableFlags::WRITABLE,
+            PageEncryption::Unencrypted,
+        );
+    }
+    flush_all();
+
+    crate::hal::change_page_state(page, PageAssignment::Shared);
+}
+
+/// Stops sharing a single 4KiB page with the hypervisor when running with AMD
+/// SEV-SNP enabled.
+pub fn unshare_page(page: Page<Size4KiB>) {
+    let page_start = page.start_address().as_u64();
+    // Only the first 2MiB is mapped as 4KiB pages, so make sure we fall in that
+    // range.
+    assert!(page_start < Size2MiB::SIZE);
+    crate::hal::change_page_state(page, PageAssignment::Private);
+    // Mark the page as encrypted.
+    {
+        let mut page_tables = crate::paging::PAGE_TABLE_REFS.get().unwrap().lock();
+        let pt = &mut page_tables.pt_0;
+        pt[page.p1_index()].set_address(
+            PhysAddr::new(page_start),
+            PageTableFlags::PRESENT | PageTableFlags::WRITABLE,
+            PageEncryption::Encrypted,
+        );
+    }
+    flush_all();
+
+    // We have to revalidate the page again after un-sharing it.
+    crate::hal::revalidate_page(page);
+}

stage0/src/sev.rs

@@ -22,19 +22,17 @@ use core::{
 };
 
 use oak_core::sync::OnceCell;
-use oak_sev_guest::{
-    ghcb::GhcbProtocol,
-    instructions::{pvalidate, InstructionError, PageSize as SevPageSize, Validation},
-    msr::{change_snp_page_state, PageAssignment, SevStatus, SnpPageStateChangeRequest},
-};
+use oak_sev_guest::{ghcb::GhcbProtocol, msr::SevStatus};
 use spinning_top::{lock_api::MutexGuard, RawSpinlock, Spinlock};
 use x86_64::{
-    instructions::tlb,
-    structures::paging::{Page, PageSize, PageTableFlags, Size2MiB, Size4KiB},
+    structures::paging::{Page, PageSize, Size4KiB},
     PhysAddr, VirtAddr,
 };
 
-use crate::{paging::PageEncryption, sev_status, BootAllocator};
+use crate::{
+    paging::{share_page, unshare_page},
+    sev_status, BootAllocator,
+};
 
 pub static GHCB_WRAPPER: Ghcb = Ghcb::new();
 
@@ -113,14 +111,12 @@ unsafe impl<A: Allocator> Allocator for SharedAllocator<A> {
         let layout =
             layout.align_to(Size4KiB::SIZE as usize).map_err(|_| AllocError)?.pad_to_align();
         let allocation = self.inner.allocate(layout)?;
-        if sev_status().contains(SevStatus::SEV_ENABLED) {
-            for offset in (0..allocation.len()).step_by(Size4KiB::SIZE as usize) {
-                // Safety: the allocation has succeeded and the offset won't exceed the size of
-                // the allocation.
-                share_page(Page::containing_address(VirtAddr::from_ptr(unsafe {
-                    allocation.as_non_null_ptr().as_ptr().add(offset)
-                })))
-            }
+        for offset in (0..allocation.len()).step_by(Size4KiB::SIZE as usize) {
+            // Safety: the allocation has succeeded and the offset won't exceed the size of
+            // the allocation.
+            share_page(Page::containing_address(VirtAddr::from_ptr(unsafe {
+                allocation.as_non_null_ptr().as_ptr().add(offset)
+            })))
         }
         Ok(allocation)
     }
@@ -131,14 +127,12 @@ unsafe impl<A: Allocator> Allocator for SharedAllocator<A> {
             .map_err(|_| AllocError)
             .unwrap()
             .pad_to_align();
-        if sev_status().contains(SevStatus::SEV_ENABLED) {
-            for offset in (0..layout.size()).step_by(Size4KiB::SIZE as usize) {
-                // Safety: the allocation has succeeded and the offset won't exceed the size of
-                // the allocation.
-                unshare_page(Page::containing_address(VirtAddr::from_ptr(unsafe {
-                    ptr.as_ptr().add(offset)
-                })))
-            }
+        for offset in (0..layout.size()).step_by(Size4KiB::SIZE as usize) {
+            // Safety: the allocation has succeeded and the offset won't exceed the size of
+            // the allocation.
+            unshare_page(Page::containing_address(VirtAddr::from_ptr(unsafe {
+                ptr.as_ptr().add(offset)
+            })))
         }
         self.inner.deallocate(ptr, layout)
     }
@@ -204,60 +198,3 @@ impl<T, A: Allocator> AsMut<T> for Shared<T, A> {
         &mut self.inner
     }
 }
-
-/// Shares a single 4KiB page with the hypervisor.
-fn share_page(page: Page<Size4KiB>) {
-    let page_start = page.start_address().as_u64();
-    // Only the first 2MiB is mapped as 4KiB pages, so make sure we fall in that
-    // range.
-    assert!(page_start < Size2MiB::SIZE);
-    // Remove the ENCRYPTED bit from the entry that maps the page.
-    {
-        let mut page_tables = crate::paging::PAGE_TABLE_REFS.get().unwrap().lock();
-        let pt = &mut page_tables.pt_0;
-        pt[page.p1_index()].set_address(
-            PhysAddr::new(page_start),
-            PageTableFlags::PRESENT | PageTableFlags::WRITABLE,
-            PageEncryption::Unencrypted,
-        );
-    }
-    tlb::flush_all();
-
-    // SNP requires extra handling beyond just removing the encrypted bit.
-    if sev_status().contains(SevStatus::SNP_ACTIVE) {
-        let request = SnpPageStateChangeRequest::new(page_start as usize, PageAssignment::Shared)
-            .expect("invalid address for page location");
-        change_snp_page_state(request).expect("couldn't change SNP state for page");
-    }
-}
-
-/// Stops sharing a single 4KiB page with the hypervisor when running with AMD
-/// SEV-SNP enabled.
-fn unshare_page(page: Page<Size4KiB>) {
-    let page_start = page.start_address().as_u64();
-    // Only the first 2MiB is mapped as 4KiB pages, so make sure we fall in that
-    // range.
-    assert!(page_start < Size2MiB::SIZE);
-    if sev_status().contains(SevStatus::SNP_ACTIVE) {
-        let request = SnpPageStateChangeRequest::new(page_start as usize, PageAssignment::Private)
-            .expect("invalid address for page location");
-        change_snp_page_state(request).expect("couldn't change SNP state for page");
-    }
-    // Mark the page as encrypted.
-    {
-        let mut page_tables = crate::paging::PAGE_TABLE_REFS.get().unwrap().lock();
-        let pt = &mut page_tables.pt_0;
-        pt[page.p1_index()].set_address(
-            PhysAddr::new(page_start),
-            PageTableFlags::PRESENT | PageTableFlags::WRITABLE,
-            PageEncryption::Encrypted,
-        );
-    }
-    tlb::flush_all();
-    // We have to revalidate the page again after un-sharing it.
-    if let Err(err) = pvalidate(page_start as usize, SevPageSize::Page4KiB, Validation::Validated) {
-        if err != InstructionError::ValidationStatusNotUpdated {
-            panic!("shared page revalidation failed");
-        }
-    }
-}