mm: always use this_cpu().get_pgtable() instead of get_init_pgtable_locked()

kernel/src/cpu/percpu.rs

@@ -15,7 +15,7 @@ use crate::cpu::vmsa::{init_guest_vmsa, init_svsm_vmsa};
 use crate::cpu::{IrqState, LocalApic};
 use crate::error::{ApicError, SvsmError};
 use crate::locking::{LockGuard, RWLock, RWLockIrqSafe, SpinLock};
-use crate::mm::pagetable::{get_init_pgtable_locked, PTEntryFlags, PageTableRef};
+use crate::mm::pagetable::{PTEntryFlags, PageTable};
 use crate::mm::virtualrange::VirtualRange;
 use crate::mm::vm::{Mapping, VMKernelStack, VMPhysMem, VMRMapping, VMReserved, VMR};
 use crate::mm::{
@@ -288,7 +288,7 @@ pub struct PerCpu {
     /// PerCpu IRQ state tracking
     irq_state: IrqState,
 
-    pgtbl: RefCell<PageTableRef>,
+    pgtbl: RefCell<Option<&'static mut PageTable>>,
     tss: Cell<X86Tss>,
     svsm_vmsa: OnceCell<VmsaPage>,
     reset_ip: Cell<u64>,
@@ -322,7 +322,7 @@ impl PerCpu {
     /// Creates a new default [`PerCpu`] struct.
     fn new(apic_id: u32) -> Self {
         Self {
-            pgtbl: RefCell::new(PageTableRef::unset()),
+            pgtbl: RefCell::new(None),
             irq_state: IrqState::new(),
             tss: Cell::new(X86Tss::new()),
             svsm_vmsa: OnceCell::new(),
@@ -437,16 +437,15 @@ impl PerCpu {
         self.shared().apic_id()
     }
 
-    fn allocate_page_table(&self) -> Result<(), SvsmError> {
+    pub fn init_page_table(&self, pgtable: PageBox<PageTable>) -> Result<(), SvsmError> {
         self.vm_range.initialize()?;
-        let pgtable_ref = get_init_pgtable_locked().clone_shared()?;
-        self.set_pgtable(pgtable_ref);
+        self.set_pgtable(PageBox::leak(pgtable));
 
         Ok(())
     }
 
-    pub fn set_pgtable(&self, pgtable: PageTableRef) {
-        *self.get_pgtable() = pgtable;
+    pub fn set_pgtable(&self, pgtable: &'static mut PageTable) {
+        *self.pgtbl.borrow_mut() = Some(pgtable);
     }
 
     fn allocate_stack(&self, base: VirtAddr) -> Result<VirtAddr, SvsmError> {
@@ -471,8 +470,10 @@ impl PerCpu {
         Ok(())
     }
 
-    pub fn get_pgtable(&self) -> RefMut<'_, PageTableRef> {
-        self.pgtbl.borrow_mut()
+    pub fn get_pgtable(&self) -> RefMut<'_, PageTable> {
+        RefMut::map(self.pgtbl.borrow_mut(), |pgtbl| {
+            &mut **pgtbl.as_mut().unwrap()
+        })
     }
 
     /// Registers an already set up GHCB page for this CPU.
@@ -554,9 +555,12 @@ impl PerCpu {
         self.vm_range.dump_ranges();
     }
 
-    pub fn setup(&self, platform: &dyn SvsmPlatform) -> Result<(), SvsmError> {
-        // Allocate page-table
-        self.allocate_page_table()?;
+    pub fn setup(
+        &self,
+        platform: &dyn SvsmPlatform,
+        pgtable: PageBox<PageTable>,
+    ) -> Result<(), SvsmError> {
+        self.init_page_table(pgtable)?;
 
         // Map PerCpu data in own page-table
         self.map_self()?;
@@ -822,7 +826,7 @@ impl PerCpu {
     /// # Arguments
     ///
     /// * `pt` - The page table to populate the the PerCpu range into
-    pub fn populate_page_table(&self, pt: &mut PageTableRef) {
+    pub fn populate_page_table(&self, pt: &mut PageTable) {
         self.vm_range.populate(pt);
     }
 

kernel/src/mm/pagetable.rs

@@ -8,17 +8,16 @@ use crate::address::{Address, PhysAddr, VirtAddr};
 use crate::cpu::control_regs::write_cr3;
 use crate::cpu::flush_tlb_global_sync;
 use crate::error::SvsmError;
-use crate::locking::{LockGuard, SpinLock};
 use crate::mm::PageBox;
 use crate::mm::{phys_to_virt, virt_to_phys, PGTABLE_LVL3_IDX_SHARED};
 use crate::platform::SvsmPlatform;
 use crate::types::{PageSize, PAGE_SIZE, PAGE_SIZE_2M};
 use crate::utils::immut_after_init::{ImmutAfterInitCell, ImmutAfterInitResult};
 use crate::utils::MemoryRegion;
 use bitflags::bitflags;
-use core::ops::{Deref, DerefMut, Index, IndexMut};
+use core::cmp;
+use core::ops::{Index, IndexMut};
 use core::ptr::NonNull;
-use core::{cmp, ptr};
 
 extern crate alloc;
 use alloc::boxed::Box;
@@ -305,8 +304,8 @@ impl PageTable {
     ///
     /// # Errors
     /// Returns [`SvsmError`] if the page cannot be allocated.
-    pub fn clone_shared(&self) -> Result<PageTableRef, SvsmError> {
-        let mut pgtable = PageTableRef::alloc()?;
+    pub fn clone_shared(&self) -> Result<PageBox<PageTable>, SvsmError> {
+        let mut pgtable = PageBox::try_new(PageTable::default())?;
         pgtable.root.entries[PGTABLE_LVL3_IDX_SHARED] = self.root.entries[PGTABLE_LVL3_IDX_SHARED];
         Ok(pgtable)
     }
@@ -881,93 +880,6 @@ impl PageTable {
     }
 }
 
-static INIT_PGTABLE: SpinLock<PageTableRef> = SpinLock::new(PageTableRef::unset());
-
-/// Sets the initial page table unless it is already set.
-///
-/// # Parameters
-/// - `pgtable`: The page table reference to set as the initial page table.
-///
-/// # Panics
-/// Panics if the initial page table is already set.
-pub fn set_init_pgtable(pgtable: PageTableRef) {
-    let mut init_pgtable = INIT_PGTABLE.lock();
-    assert!(!init_pgtable.is_set());
-    *init_pgtable = pgtable;
-}
-
-/// Acquires a lock and returns a guard for the initial page table, which
-/// is locked for the duration of the guard's scope.
-///
-/// # Returns
-/// A `LockGuard` for the initial page table.
-pub fn get_init_pgtable_locked<'a>() -> LockGuard<'a, PageTableRef> {
-    INIT_PGTABLE.lock()
-}
-
-/// A reference wrapper for a [`PageTable`].
-#[derive(Debug)]
-pub enum PageTableRef {
-    Owned(PageBox<PageTable>),
-    Shared(*mut PageTable),
-}
-
-impl PageTableRef {
-    /// Creates a new shared [`PageTableRef`] from a raw pointer.
-    #[inline]
-    pub const fn shared(ptr: *mut PageTable) -> Self {
-        Self::Shared(ptr)
-    }
-
-    /// Allocates an empty owned [`PageTableRef`].
-    pub fn alloc() -> Result<Self, SvsmError> {
-        let table = PageBox::try_new(PageTable::default())?;
-        Ok(Self::Owned(table))
-    }
-
-    /// Creates a new shared and unset (i.e. NULL) [`PageTableRef`].
-    #[inline]
-    pub const fn unset() -> Self {
-        Self::shared(ptr::null_mut())
-    }
-
-    /// Checks if the [`PageTableRef`] is set, i.e. not NULL.
-    #[inline]
-    fn is_set(&self) -> bool {
-        match self {
-            Self::Owned(..) => true,
-            Self::Shared(p) => !p.is_null(),
-        }
-    }
-}
-
-impl Deref for PageTableRef {
-    type Target = PageTable;
-
-    fn deref(&self) -> &Self::Target {
-        match self {
-            Self::Owned(p) => p,
-            // SAFETY: nobody else has access to `ptr` so it cannot be aliased.
-            Self::Shared(p) => unsafe { p.as_ref().unwrap() },
-        }
-    }
-}
-
-impl DerefMut for PageTableRef {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        match self {
-            Self::Owned(p) => p,
-            // SAFETY: nobody else has access to `ptr` so it cannot be aliased.
-            Self::Shared(p) => unsafe { p.as_mut().unwrap() },
-        }
-    }
-}
-
-/// SAFETY: `PageTableRef` is more or less equivalent to a mutable reference to
-///         a PageTable and so if `&mut PageTable` implements `Send` so does
-///         `PageTableRef`.
-unsafe impl Send for PageTableRef where &'static mut PageTable: Send {}
-
 /// Represents a sub-tree of a page-table which can be mapped at a top-level index
 #[derive(Default, Debug)]
 struct RawPageTablePart {

kernel/src/mm/vm/range.rs

@@ -8,7 +8,7 @@ use crate::address::{Address, VirtAddr};
 use crate::cpu::{flush_tlb_global_percpu, flush_tlb_global_sync};
 use crate::error::SvsmError;
 use crate::locking::RWLock;
-use crate::mm::pagetable::{PTEntryFlags, PageTable, PageTablePart, PageTableRef};
+use crate::mm::pagetable::{PTEntryFlags, PageTable, PageTablePart};
 use crate::types::{PageSize, PAGE_SHIFT, PAGE_SIZE};
 use crate::utils::{align_down, align_up};
 
@@ -122,16 +122,16 @@ impl VMR {
     ///
     /// # Arguments
     ///
-    /// * `pgtbl` - A [`PageTableRef`] pointing to the target page-table
-    pub fn populate(&self, pgtbl: &mut PageTableRef) {
+    /// * `pgtbl` - A [`PageTable`] pointing to the target page-table
+    pub fn populate(&self, pgtbl: &mut PageTable) {
         let parts = self.pgtbl_parts.lock_read();
 
         for part in parts.iter() {
             pgtbl.populate_pgtbl_part(part);
         }
     }
 
-    pub fn populate_addr(&self, pgtbl: &mut PageTableRef, vaddr: VirtAddr) {
+    pub fn populate_addr(&self, pgtbl: &mut PageTable, vaddr: VirtAddr) {
         let start = VirtAddr::from(self.start_pfn << PAGE_SHIFT);
         let end = VirtAddr::from(self.end_pfn << PAGE_SHIFT);
         assert!(vaddr >= start && vaddr < end);

kernel/src/platform/snp.rs

@@ -220,7 +220,8 @@ impl SvsmPlatform for SnpPlatform {
     }
 
     fn start_cpu(&self, cpu: &PerCpu, start_rip: u64) -> Result<(), SvsmError> {
-        cpu.setup(self)?;
+        let pgtable = this_cpu().get_pgtable().clone_shared()?;
+        cpu.setup(self, pgtable)?;
         let (vmsa_pa, sev_features) = cpu.alloc_svsm_vmsa(*VTOM as u64, start_rip)?;
 
         current_ghcb().ap_create(vmsa_pa, cpu.get_apic_id().into(), 0, sev_features)

kernel/src/sev/ghcb.rs

@@ -9,7 +9,6 @@ use crate::cpu::msr::{write_msr, SEV_GHCB};
 use crate::cpu::percpu::this_cpu;
 use crate::cpu::{flush_tlb_global_sync, X86GeneralRegs};
 use crate::error::SvsmError;
-use crate::mm::pagetable::get_init_pgtable_locked;
 use crate::mm::validate::{
     valid_bitmap_clear_valid_4k, valid_bitmap_set_valid_4k, valid_bitmap_valid_addr,
 };
@@ -151,7 +150,7 @@ impl GhcbPage {
         }
 
         // Map page unencrypted
-        get_init_pgtable_locked().set_shared_4k(vaddr)?;
+        this_cpu().get_pgtable().set_shared_4k(vaddr)?;
         flush_tlb_global_sync();
 
         // SAFETY: all zeros is a valid representation for the GHCB.
@@ -323,7 +322,7 @@ impl GHCB {
         let paddr = virt_to_phys(vaddr);
 
         // Re-encrypt page
-        get_init_pgtable_locked().set_encrypted_4k(vaddr)?;
+        this_cpu().get_pgtable().set_encrypted_4k(vaddr)?;
 
         // Unregister GHCB PA
         register_ghcb_gpa_msr(PhysAddr::null())?;

kernel/src/stage2.rs

@@ -28,10 +28,7 @@ use svsm::error::SvsmError;
 use svsm::fw_cfg::FwCfg;
 use svsm::igvm_params::IgvmParams;
 use svsm::mm::alloc::{memory_info, print_memory_info, root_mem_init};
-use svsm::mm::pagetable::{
-    get_init_pgtable_locked, paging_init_early, set_init_pgtable, PTEntryFlags, PageTable,
-    PageTableRef,
-};
+use svsm::mm::pagetable::{paging_init_early, PTEntryFlags, PageTable};
 use svsm::mm::validate::{
     init_valid_bitmap_alloc, valid_bitmap_addr, valid_bitmap_set_valid_range,
 };
@@ -41,7 +38,7 @@ use svsm::types::{PageSize, PAGE_SIZE, PAGE_SIZE_2M};
 use svsm::utils::{halt, is_aligned, MemoryRegion};
 
 extern "C" {
-    pub static mut pgtable: PageTable;
+    static mut pgtable: PageTable;
 }
 
 fn setup_stage2_allocator(heap_start: u64, heap_end: u64) {
@@ -55,9 +52,12 @@ fn setup_stage2_allocator(heap_start: u64, heap_end: u64) {
 
 fn init_percpu(platform: &mut dyn SvsmPlatform) -> Result<(), SvsmError> {
     let bsp_percpu = PerCpu::alloc(0)?;
-    unsafe {
-        bsp_percpu.set_pgtable(PageTableRef::shared(addr_of_mut!(pgtable)));
-    }
+    let init_pgtable = unsafe {
+        // SAFETY: pgtable is a static mut and this is the only place where we
+        // get a reference to it.
+        &mut *addr_of_mut!(pgtable)
+    };
+    bsp_percpu.set_pgtable(init_pgtable);
     bsp_percpu.map_self_stage2()?;
     platform.setup_guest_host_comm(bsp_percpu, true);
     Ok(())
@@ -106,8 +106,6 @@ fn setup_env(
     register_cpuid_table(cpuid_page);
     paging_init_early(platform).expect("Failed to initialize early paging");
 
-    set_init_pgtable(PageTableRef::shared(unsafe { addr_of_mut!(pgtable) }));
-
     // Configure the heap to exist from 64 KB to 640 KB.
     setup_stage2_allocator(0x10000, 0xA0000);
 
@@ -139,7 +137,7 @@ fn map_and_validate(
         | PTEntryFlags::ACCESSED
         | PTEntryFlags::DIRTY;
 
-    let mut pgtbl = get_init_pgtable_locked();
+    let mut pgtbl = this_cpu().get_pgtable();
     pgtbl.map_region(vregion, paddr, flags)?;
 
     if config.page_state_change_required() {

kernel/src/svsm.rs

@@ -11,7 +11,6 @@ use svsm::fw_meta::{print_fw_meta, validate_fw_memory, SevFWMetaData};
 
 use bootlib::kernel_launch::KernelLaunchInfo;
 use core::arch::global_asm;
-use core::mem::size_of;
 use core::panic::PanicInfo;
 use core::ptr;
 use core::slice;
@@ -336,16 +335,14 @@ pub extern "C" fn svsm_start(li: &KernelLaunchInfo, vb_addr: usize) {
     };
 
     paging_init(platform).expect("Failed to initialize paging");
-    init_page_table(&launch_info, &kernel_elf).expect("Could not initialize the page table");
+    let init_pgtable =
+        init_page_table(&launch_info, &kernel_elf).expect("Could not initialize the page table");
+    init_pgtable.load();
 
-    // SAFETY: this PerCpu has just been allocated and no other CPUs have been
-    // brought up, thus it cannot be aliased and we can get a mutable
-    // reference to it. We trust PerCpu::alloc() to return a valid and
-    // aligned pointer.
     let bsp_percpu = PerCpu::alloc(0).expect("Failed to allocate BSP per-cpu data");
 
     bsp_percpu
-        .setup(platform)
+        .setup(platform, init_pgtable)
         .expect("Failed to setup BSP per-cpu area");
     bsp_percpu
         .setup_on_cpu(platform)

kernel/src/svsm_paging.rs

@@ -8,8 +8,8 @@ use crate::address::{Address, PhysAddr, VirtAddr};
 use crate::config::SvsmConfig;
 use crate::error::SvsmError;
 use crate::igvm_params::IgvmParams;
-use crate::mm::pagetable::{set_init_pgtable, PTEntryFlags, PageTableRef};
-use crate::mm::PerCPUPageMappingGuard;
+use crate::mm::pagetable::{PTEntryFlags, PageTable};
+use crate::mm::{PageBox, PerCPUPageMappingGuard};
 use crate::platform::PageStateChangeOp;
 use crate::platform::SvsmPlatform;
 use crate::types::{PageSize, PAGE_SIZE};
@@ -24,8 +24,8 @@ struct IgvmParamInfo<'a> {
 pub fn init_page_table(
     launch_info: &KernelLaunchInfo,
     kernel_elf: &elf::Elf64File<'_>,
-) -> Result<(), SvsmError> {
-    let mut pgtable = PageTableRef::alloc()?;
+) -> Result<PageBox<PageTable>, SvsmError> {
+    let mut pgtable = PageBox::try_new(PageTable::default())?;
     let igvm_param_info = if launch_info.igvm_params_virt_addr != 0 {
         let addr = VirtAddr::from(launch_info.igvm_params_virt_addr);
         IgvmParamInfo {
@@ -91,8 +91,7 @@ pub fn init_page_table(
 
     pgtable.load();
 
-    set_init_pgtable(pgtable);
-    Ok(())
+    Ok(pgtable)
 }
 
 fn invalidate_boot_memory_region(