support more than 64 VPs

Author: Sangho Lee

Cargo.lock

@@ -17,6 +17,7 @@ spin = { version = "0.10.0", default-features = false, features = [
   "once",
   "rwlock",
 ] }
+bitvec = { version = "1.0.1", default-features = false, features = ["alloc"] }
 arrayvec = { version = "0.7.6", default-features = false }
 rangemap = { version = "1.5.1", features = ["const_fn"] }
 thiserror = { version = "2.0.6", default-features = false }

litebox_platform_lvbs/Cargo.toml

@@ -484,6 +484,7 @@ impl<M: MemoryProvider, const ALIGN: usize> X64PageTable<'_, M, ALIGN> {
     /// # Behavior
     /// - Any existing mapping is treated as an error
     /// - On error, all pages mapped by this call are unmapped (atomic)
+    #[cfg(feature = "optee_syscall")]
     pub(crate) fn map_non_contiguous_phys_frames(
         &self,
         frames: &[PhysFrame<Size4KiB>],
@@ -551,6 +552,7 @@ impl<M: MemoryProvider, const ALIGN: usize> X64PageTable<'_, M, ALIGN> {
     ///
     /// Note: The caller must already hold the page table lock (`self.inner`).
     /// This function accepts the locked `MappedPageTable` directly.
+    #[cfg(feature = "optee_syscall")]
     fn rollback_mapped_pages(
         inner: &mut MappedPageTable<'_, FrameMapping<M>>,
         pages: x86_64::structures::paging::page::PageRangeInclusive<Size4KiB>,

litebox_platform_lvbs/src/arch/x86/mm/paging.rs

@@ -24,14 +24,12 @@ use litebox::{
     shim::ContinueOperation,
     utils::TruncateExt,
 };
-use litebox_common_linux::{
-    PunchthroughSyscall,
-    errno::Errno,
-    vmap::{
-        PhysPageAddr, PhysPageAddrArray, PhysPageMapInfo, PhysPageMapPermissions, PhysPointerError,
-        VmapManager,
-    },
+#[cfg(feature = "optee_syscall")]
+use litebox_common_linux::vmap::{
+    PhysPageAddr, PhysPageAddrArray, PhysPageMapInfo, PhysPageMapPermissions, PhysPointerError,
+    VmapManager,
 };
+use litebox_common_linux::{PunchthroughSyscall, errno::Errno};
 use x86_64::{
     VirtAddr,
     structures::paging::{

litebox_platform_lvbs/src/lib.rs

@@ -50,6 +50,7 @@ pub trait MemoryProvider {
     fn pa_to_va(pa: PhysAddr) -> VirtAddr {
         let pa = pa.as_u64() & !Self::PRIVATE_PTE_MASK;
         let va = VirtAddr::new_truncate(pa + Self::GVA_OFFSET.as_u64());
+        #[cfg(feature = "optee_syscall")]
         assert!(
             va.as_u64() < vmap::VMAP_START as u64,
             "VA {va:#x} is out of range for direct mapping"

litebox_platform_lvbs/src/mm/mod.rs

@@ -5,9 +5,10 @@
 
 #[cfg(not(test))]
 use crate::mshv::{
-    HV_FLUSH_ALL_VIRTUAL_ADDRESS_SPACES, HVCALL_FLUSH_VIRTUAL_ADDRESS_LIST,
-    HVCALL_FLUSH_VIRTUAL_ADDRESS_SPACE, HvInputFlushVirtualAddressList,
-    HvInputFlushVirtualAddressSpace, hvcall::hv_do_hypercall, vtl_switch::vtl1_vp_mask,
+    HV_FLUSH_ALL_VIRTUAL_ADDRESS_SPACES, HV_FLUSH_EX_VP_SET_BANKS, HV_GENERIC_SET_SPARSE_4K,
+    HVCALL_FLUSH_VIRTUAL_ADDRESS_LIST_EX, HVCALL_FLUSH_VIRTUAL_ADDRESS_SPACE_EX,
+    HvInputFlushVirtualAddressListEx, HvInputFlushVirtualAddressSpaceEx, hvcall::hv_do_hypercall,
+    vtl_switch::vtl1_vp_mask,
 };
 use crate::{
     host::per_cpu_variables::with_per_cpu_variables_mut,
@@ -19,6 +20,21 @@ use crate::{
     },
 };
 
+#[cfg(not(test))]
+#[inline]
+fn vp_set_valid_bank_mask(vp_set_bank_contents: [u64; HV_FLUSH_EX_VP_SET_BANKS]) -> u64 {
+    vp_set_bank_contents
+        .iter()
+        .enumerate()
+        .fold(0u64, |mask, (bank, contents)| {
+            if *contents != 0 {
+                mask | (1u64 << bank)
+            } else {
+                mask
+            }
+        })
+}
+
 /// Hyper-V Hypercall to prevent lower VTLs (i.e., VTL0) from accessing a specified range of
 /// guest physical memory pages with a given protection flag.
 pub fn hv_modify_vtl_protection_mask(
@@ -74,20 +90,27 @@ pub fn hv_modify_vtl_protection_mask(
 /// This is the cross-core equivalent of a local CR3 reload.
 #[cfg(not(test))]
 pub(crate) fn hv_flush_virtual_address_space() -> Result<(), HypervCallError> {
+    let vp_mask = vtl1_vp_mask();
+    let valid_bank_mask = vp_set_valid_bank_mask(vp_mask);
+    if valid_bank_mask == 0 {
+        return Ok(());
+    }
     let input = with_per_cpu_variables_mut(|pcv| unsafe {
         &mut *pcv
             .hv_hypercall_input_page_as_mut_ptr()
-            .cast::<HvInputFlushVirtualAddressSpace>()
+            .cast::<HvInputFlushVirtualAddressSpaceEx>()
     });
 
-    *input = HvInputFlushVirtualAddressSpace {
+    *input = HvInputFlushVirtualAddressSpaceEx {
         address_space: 0,
         flags: HV_FLUSH_ALL_VIRTUAL_ADDRESS_SPACES,
-        processor_mask: vtl1_vp_mask(),
+        vp_set_format: HV_GENERIC_SET_SPARSE_4K,
+        vp_set_valid_bank_mask: valid_bank_mask,
+        vp_set_bank_contents: vp_mask,
     };
 
     hv_do_hypercall(
-        u64::from(HVCALL_FLUSH_VIRTUAL_ADDRESS_SPACE),
+        u64::from(HVCALL_FLUSH_VIRTUAL_ADDRESS_SPACE_EX),
         (&raw const *input).cast::<core::ffi::c_void>(),
         core::ptr::null_mut(),
     )?;
@@ -115,15 +138,22 @@ pub(crate) fn hv_flush_virtual_address_list(
     );
     debug_assert!(page_count > 0, "page_count must not be 0");
 
+    let vp_mask = vtl1_vp_mask();
+    let valid_bank_mask = vp_set_valid_bank_mask(vp_mask);
+    if valid_bank_mask == 0 {
+        return Ok(());
+    }
     let input = with_per_cpu_variables_mut(|pcv| unsafe {
         &mut *pcv
             .hv_hypercall_input_page_as_mut_ptr()
-            .cast::<HvInputFlushVirtualAddressList>()
+            .cast::<HvInputFlushVirtualAddressListEx>()
     });
 
     input.address_space = 0;
     input.flags = HV_FLUSH_ALL_VIRTUAL_ADDRESS_SPACES;
-    input.processor_mask = vtl1_vp_mask();
+    input.vp_set_format = HV_GENERIC_SET_SPARSE_4K;
+    input.vp_set_valid_bank_mask = valid_bank_mask;
+    input.vp_set_bank_contents = vp_mask;
 
     let mut remaining = page_count;
     let mut current_va = start_va;
@@ -132,7 +162,7 @@ pub(crate) fn hv_flush_virtual_address_list(
         let mut gva_count: u16 = 0;
 
         while remaining > 0
-            && (gva_count as usize) < HvInputFlushVirtualAddressList::MAX_GVAS_PER_REQUEST
+            && (gva_count as usize) < HvInputFlushVirtualAddressListEx::MAX_GVAS_PER_REQUEST
         {
             // Each entry can cover up to `MAX_ADDITIONAL_PAGES + 1` pages.
             let additional = remaining.saturating_sub(1).min(MAX_ADDITIONAL_PAGES);
@@ -148,9 +178,9 @@ pub(crate) fn hv_flush_virtual_address_list(
         }
 
         hv_do_rep_hypercall(
-            HVCALL_FLUSH_VIRTUAL_ADDRESS_LIST,
+            HVCALL_FLUSH_VIRTUAL_ADDRESS_LIST_EX,
             gva_count,
-            0, // simple processor mask, not HV_VP_SET
+            HvInputFlushVirtualAddressListEx::VP_SET_QWORD_COUNT,
             (&raw const *input).cast::<core::ffi::c_void>(),
             core::ptr::null_mut(),
         )?;

litebox_platform_lvbs/src/mshv/hvcall_mm.rs

@@ -15,6 +15,7 @@ pub mod vsm_intercept;
 pub mod vtl1_mem_layout;
 pub mod vtl_switch;
 
+use crate::arch::MAX_CORES;
 use crate::mshv::vtl1_mem_layout::PAGE_SIZE;
 use modular_bitfield::prelude::*;
 use modular_bitfield::specifiers::{B3, B4, B7, B8, B16, B31, B32, B45, B51, B62};
@@ -73,8 +74,8 @@ pub const VTL_ENTRY_REASON_RESERVED: u32 = 0x0;
 pub const VTL_ENTRY_REASON_LOWER_VTL_CALL: u32 = 0x1;
 pub const VTL_ENTRY_REASON_INTERRUPT: u32 = 0x2;
 
-pub const HVCALL_FLUSH_VIRTUAL_ADDRESS_SPACE: u16 = 0x_0002;
-pub const HVCALL_FLUSH_VIRTUAL_ADDRESS_LIST: u16 = 0x_0003;
+pub const HVCALL_FLUSH_VIRTUAL_ADDRESS_SPACE_EX: u16 = 0x_0013;
+pub const HVCALL_FLUSH_VIRTUAL_ADDRESS_LIST_EX: u16 = 0x_0014;
 pub const HVCALL_MODIFY_VTL_PROTECTION_MASK: u16 = 0x_000c;
 pub const HVCALL_ENABLE_VP_VTL: u16 = 0x_000f;
 pub const HVCALL_GET_VP_REGISTERS: u16 = 0x_0050;
@@ -525,62 +526,65 @@ impl Default for HvInputModifyVtlProtectionMask {
     }
 }
 
-/// Input structure for `HvCallFlushVirtualAddressSpace` (0x0002).
+/// VP-set format for sparse 4K virtual processor numbering.
+pub const HV_GENERIC_SET_SPARSE_4K: u64 = 0;
+
+/// Number of VP banks encoded in EX flush requests.
 ///
-/// Layout (TLFS §6.5.2):
-/// - Fixed header: `address_space` (u64) + `flags` (u64) + `processor_mask` (u64)
+/// Each bank contains 64 VPs.
+pub const HV_FLUSH_EX_VP_SET_BANKS: usize = MAX_CORES.div_ceil(64);
+
+/// Input structure for `HvCallFlushVirtualAddressSpaceEx` (0x0013).
 ///
-/// This is the **non-extended** version; it uses a simple 64-bit processor
-/// mask (one bit per VP, max 64 VPs).  When `HV_FLUSH_ALL_PROCESSORS` is
-/// set in `flags` the mask is ignored, but must still be present.
+/// Layout (TLFS §6.5.2 EX):
+/// - `address_space` (u64)
+/// - `flags` (u64)
+/// - VP set header: `vp_set_format` (u64), `vp_set_valid_bank_mask` (u64)
+/// - VP set banks: `vp_set_bank_contents` (u64 per bank)
 #[derive(Clone, Copy)]
 #[repr(C, packed)]
-pub struct HvInputFlushVirtualAddressSpace {
-    /// CR3 / EPTP value identifying the address space. Ignored when
-    /// `HV_FLUSH_ALL_VIRTUAL_ADDRESS_SPACES` is set.
+pub struct HvInputFlushVirtualAddressSpaceEx {
     pub address_space: u64,
-    /// Combination of `HV_FLUSH_*` flags.
     pub flags: u64,
-    /// 64-bit processor mask (bit *i* → VP *i*). Ignored when
-    /// `HV_FLUSH_ALL_PROCESSORS` is set in `flags`.
-    pub processor_mask: u64,
+    pub vp_set_format: u64,
+    pub vp_set_valid_bank_mask: u64,
+    pub vp_set_bank_contents: [u64; HV_FLUSH_EX_VP_SET_BANKS],
 }
 
-/// Input structure for `HvCallFlushVirtualAddressList` (0x0003).
+/// Input structure for `HvCallFlushVirtualAddressListEx` (0x0014).
 ///
-/// Layout (TLFS §6.5.3):
-/// - Fixed header: `address_space` (u64) + `flags` (u64) + `processor_mask` (u64)
+/// Layout (TLFS §6.5.3 EX):
+/// - Fixed header: `address_space` (u64), `flags` (u64)
+/// - Variable header: VP set (`vp_set_*` and bank contents)
 /// - Rep elements: array of `HV_GVA_RANGE` entries (u64 each)
-///
-/// This is the **non-extended** version; it uses a simple 64-bit processor
-/// mask.  The GVA range list is placed immediately after the fixed header.
 #[derive(Clone, Copy)]
 #[repr(C, packed)]
-pub struct HvInputFlushVirtualAddressList {
-    /// CR3 / EPTP value identifying the address space.
+pub struct HvInputFlushVirtualAddressListEx {
     pub address_space: u64,
-    /// Combination of `HV_FLUSH_*` flags.
     pub flags: u64,
-    /// 64-bit processor mask. Ignored when `HV_FLUSH_ALL_PROCESSORS` is set.
-    pub processor_mask: u64,
-    /// GVA range entries. Each entry specifies a GVA range to flush.
-    /// Bits 11:0 encode `additional_pages` (number of *extra* pages beyond the
-    /// first; 0 = flush exactly 1 page).
-    /// Bits 63:12 encode the GVA page number (GVA >> 12).
-    pub gva_range_list: [u64; HV_FLUSH_MAX_GVAS],
+    pub vp_set_format: u64,
+    pub vp_set_valid_bank_mask: u64,
+    pub vp_set_bank_contents: [u64; HV_FLUSH_EX_VP_SET_BANKS],
+    pub gva_range_list: [u64; HV_FLUSH_EX_MAX_GVAS],
 }
 
-/// Maximum number of GVA range entries that fit in a single hypercall input page
-/// after the fixed header of `HvInputFlushVirtualAddressList`.
+/// Maximum number of GVA entries that fit in one input page for
+/// `HvInputFlushVirtualAddressListEx` with a VP set sized for `MAX_CORES`.
 ///
-/// Input page = 4096 bytes. Header = 3 × 8 = 24 bytes. Remaining = 4072 / 8 = 509.
+/// Input page = 4096 bytes.
+/// Header = (4 + `HV_FLUSH_EX_VP_SET_BANKS`) * 8 bytes.
 #[expect(clippy::cast_possible_truncation)]
-const HV_FLUSH_MAX_GVAS: usize =
-    ((PAGE_SIZE as u32 - 3 * (u64::BITS / 8)) / (u64::BITS / 8)) as usize;
+const HV_FLUSH_EX_MAX_GVAS: usize = ((PAGE_SIZE as u32
+    - (4 + HV_FLUSH_EX_VP_SET_BANKS as u32) * (u64::BITS / 8))
+    / (u64::BITS / 8)) as usize;
+
+impl HvInputFlushVirtualAddressListEx {
+    /// Number of 64-bit words occupied by the VP-set variable header.
+    #[allow(clippy::cast_possible_truncation)]
+    pub const VP_SET_QWORD_COUNT: u16 = (2 + HV_FLUSH_EX_VP_SET_BANKS) as u16;
 
-impl HvInputFlushVirtualAddressList {
-    /// Maximum number of GVA range entries per hypercall invocation.
-    pub const MAX_GVAS_PER_REQUEST: usize = HV_FLUSH_MAX_GVAS;
+    /// Maximum number of GVA range entries per EX hypercall invocation.
+    pub const MAX_GVAS_PER_REQUEST: usize = HV_FLUSH_EX_MAX_GVAS;
 }
 
 #[bitfield]

litebox_platform_lvbs/src/mshv/mod.rs

@@ -3,6 +3,7 @@
 
 //! VTL switch related functions
 
+use crate::arch::MAX_CORES;
 use crate::arch::instrs::rdmsr;
 use crate::host::{
     hv_hypercall_page_address,
@@ -12,25 +13,26 @@ use crate::host::{
     },
 };
 use crate::mshv::{
-    HV_REGISTER_VP_INDEX, HV_REGISTER_VSM_CODEPAGE_OFFSETS, HvRegisterVsmCodePageOffsets,
-    NUM_VTLCALL_PARAMS, VTL_ENTRY_REASON_INTERRUPT, VTL_ENTRY_REASON_LOWER_VTL_CALL,
-    VTL_ENTRY_REASON_RESERVED, error::VsmError, hvcall_vp::hvcall_get_vp_registers,
-    vsm_intercept::vsm_handle_intercept,
+    HV_FLUSH_EX_VP_SET_BANKS, HV_REGISTER_VP_INDEX, HV_REGISTER_VSM_CODEPAGE_OFFSETS,
+    HvRegisterVsmCodePageOffsets, NUM_VTLCALL_PARAMS, VTL_ENTRY_REASON_INTERRUPT,
+    VTL_ENTRY_REASON_LOWER_VTL_CALL, VTL_ENTRY_REASON_RESERVED, error::VsmError,
+    hvcall_vp::hvcall_get_vp_registers, vsm_intercept::vsm_handle_intercept,
 };
-use core::sync::atomic::{AtomicU64, Ordering};
+use bitvec::prelude::{BitArray, Lsb0};
 use litebox::utils::{ReinterpretUnsignedExt, TruncateExt};
 use num_enum::TryFromPrimitive;
+use spin::mutex::SpinMutex;
 
-/// Bitmask of VPs that are currently executing VTL1 code.  Bit *N* is set
+/// Bitmask of VPs that are currently executing VTL1 code. Bit *N* is set
 /// when VP index *N* is inside VTL1 (between VTL entry and VTL return).
 ///
 /// The TLB flush hypercalls read this mask so they only target VPs that
 /// are in VTL1. VPs running in VTL0 use a separate address space and
 /// will receive a full TLB flush on their next VTL1 entry.
 ///
-/// Supports up to 64 VPs (the non-extended hypercall processor mask is
-/// 64 bits wide).  [`vtl1_vp_enter`] asserts that the VP index fits.
-static VTL1_VP_MASK: AtomicU64 = AtomicU64::new(0);
+/// Supports up to `MAX_CORES` VPs. [`vtl1_vp_enter`] asserts that the VP index fits.
+type Vtl1VpMaskBits = BitArray<[u64; HV_FLUSH_EX_VP_SET_BANKS], Lsb0>;
+static VTL1_VP_MASK: SpinMutex<Vtl1VpMaskBits> = SpinMutex::new(Vtl1VpMaskBits::ZERO);
 
 /// Read the current VP index from the hypervisor MSR.
 #[inline]
@@ -41,35 +43,38 @@ fn current_vp_index() -> u32 {
 /// Mark the current VP as executing in VTL1.
 ///
 /// # Panics
-/// Panics (debug-only) if the VP index ≥ 64.
+/// Panics if the VP index ≥ `MAX_CORES`.
 #[inline]
 fn vtl1_vp_enter() {
     let vp_index = current_vp_index();
-    debug_assert!(
-        vp_index < 64,
-        "VP index {vp_index} exceeds the 64-bit processor mask"
+    assert!(
+        vp_index < u32::try_from(MAX_CORES).unwrap(),
+        "VP index {vp_index} exceeds the configured processor mask"
     );
-    VTL1_VP_MASK.fetch_or(1u64 << vp_index, Ordering::Release);
+    VTL1_VP_MASK.lock().set(vp_index as usize, true);
 }
 
 /// Remove the current VP from the VTL1 mask (it is returning to VTL0).
 #[inline]
 fn vtl1_vp_exit() {
     let vp_index = current_vp_index();
-    VTL1_VP_MASK.fetch_and(!(1u64 << vp_index), Ordering::Release);
+    VTL1_VP_MASK.lock().set(vp_index as usize, false);
 }
 
 /// Return the current VTL1 VP mask for use in TLB flush hypercalls.
 ///
-/// The returned value is a snapshot — VPs may enter or leave VTL1
+/// The returned value is a snapshot. VPs may enter or leave VTL1
 /// between the load and the hypercall, but that is benign:
 /// - A VP that left VTL1 after the snapshot merely receives a redundant flush.
-/// - A VP that entered VTL1 after the snapshot will perform a full TLB flush
-///   as part of the VTL transition.
+/// - A VP that entered VTL1 after the snapshot performs a full TLB flush
+///   because CR3 is reloaded at the entry and PCID is not enabled in VTL1.
 #[cfg(not(test))]
 #[inline]
-pub(crate) fn vtl1_vp_mask() -> u64 {
-    VTL1_VP_MASK.load(Ordering::Acquire)
+pub(crate) fn vtl1_vp_mask() -> [u64; HV_FLUSH_EX_VP_SET_BANKS] {
+    let mask = VTL1_VP_MASK.lock();
+    let mut banks = [0u64; HV_FLUSH_EX_VP_SET_BANKS];
+    banks.copy_from_slice(mask.as_raw_slice());
+    banks
 }
 
 // ============================================================================