virt_mshv_vtl/mshv: Cleanup TLB lock tracking

openhcl/virt_mshv_vtl/src/processor/mshv/arm64.rs

@@ -157,6 +157,7 @@ impl BackingPrivate for HypervisorBackedArm64 {
                 this.backing.next_deliverability_notifications;
         }
 
+        this.unlock_tlb_lock(Vtl::Vtl2);
         let intercepted = this
             .runner
             .run()
@@ -593,7 +594,7 @@ impl<T: CpuIo> EmulatorSupport for UhEmulationState<'_, '_, T, HypervisorBackedA
         }
     }
 
-    fn initial_gva_translation(&self) -> Option<emulate::InitialTranslation> {
+    fn initial_gva_translation(&mut self) -> Option<emulate::InitialTranslation> {
         if (self.vp.runner.exit_message().header.typ != HvMessageType::HvMessageTypeGpaIntercept)
             && (self.vp.runner.exit_message().header.typ != HvMessageType::HvMessageTypeUnmappedGpa)
             && (self.vp.runner.exit_message().header.typ
@@ -616,13 +617,20 @@ impl<T: CpuIo> EmulatorSupport for UhEmulationState<'_, '_, T, HypervisorBackedA
         let translate_mode = emulate::TranslateMode::try_from(message.header.intercept_access_type)
             .expect("unexpected intercept access type");
 
-        tracing::trace!(?message.guest_virtual_address, ?message.guest_physical_address, ?translate_mode, "initial translation");
-
-        Some(emulate::InitialTranslation {
+        let translation = emulate::InitialTranslation {
             gva: message.guest_virtual_address,
             gpa: message.guest_physical_address,
             translate_mode,
-        })
+        };
+
+        tracing::trace!(?translation, "initial translation");
+
+        // If we have a valid translation, the hypervisor must have set the TLB lock
+        // so the translation remains valid for the duration of this exit.
+        // Update our local cache appropriately.
+        self.vp.mark_tlb_locked(Vtl::Vtl2, self.vtl);
+
+        Some(translation)
     }
 
     fn interruption_pending(&self) -> bool {
@@ -653,6 +661,7 @@ impl<T: CpuIo> EmulatorSupport for UhEmulationState<'_, '_, T, HypervisorBackedA
             )
         {
             // Should always be called after translate gva with the tlb lock flag
+            // or with an initial translation.
             debug_assert!(self.vp.is_tlb_locked(Vtl::Vtl2, self.vtl));
 
             let cpsr: Cpsr64 = self
@@ -729,10 +738,13 @@ impl<T: CpuIo> EmulatorSupport for UhEmulationState<'_, '_, T, HypervisorBackedA
             Ok(ioctl::TranslateResult {
                 gpa_page,
                 overlay_page,
-            }) => Ok(Ok(EmuTranslateResult {
-                gpa: (gpa_page << hvdef::HV_PAGE_SHIFT) + (gva & (hvdef::HV_PAGE_SIZE - 1)),
-                overlay_page: Some(overlay_page),
-            })),
+            }) => {
+                self.vp.mark_tlb_locked(Vtl::Vtl2, self.vtl);
+                Ok(Ok(EmuTranslateResult {
+                    gpa: (gpa_page << hvdef::HV_PAGE_SHIFT) + (gva & (hvdef::HV_PAGE_SIZE - 1)),
+                    overlay_page: Some(overlay_page),
+                }))
+            }
             Err(ioctl::aarch64::TranslateErrorAarch64 { code }) => Ok(Err(EmuTranslateError {
                 code: hypercall::TranslateGvaResultCode(code),
                 event_info: None,

openhcl/virt_mshv_vtl/src/processor/mshv/tlb_lock.rs

@@ -8,7 +8,6 @@ use crate::UhProcessor;
 use hcl::GuestVtl;
 use hvdef::HvAllArchRegisterName;
 use hvdef::Vtl;
-use hvdef::hypercall::HvInputVtl;
 
 impl UhProcessor<'_, HypervisorBacked> {
     /// Causes the specified VTL on the current VP to wait on all TLB locks.
@@ -26,11 +25,11 @@ impl UhProcessor<'_, HypervisorBacked> {
     }
 
     /// Lock the TLB of the target VTL on the current VP.
-    #[cfg_attr(guest_arch = "aarch64", expect(dead_code))]
+    #[expect(dead_code)]
     pub(crate) fn set_tlb_lock(&mut self, requesting_vtl: Vtl, target_vtl: GuestVtl) {
         debug_assert_eq!(requesting_vtl, Vtl::Vtl2);
 
-        if self.is_tlb_locked(requesting_vtl, target_vtl) {
+        if self.vtls_tlb_locked.get(requesting_vtl, target_vtl) {
             return;
         }
 
@@ -47,33 +46,45 @@ impl UhProcessor<'_, HypervisorBacked> {
         self.vtls_tlb_locked.set(requesting_vtl, target_vtl, true);
     }
 
+    /// Mark the TLB of the target VTL on the current VP as locked without
+    /// informing the hypervisor. Only should be used when the hypervisor
+    /// is expected to have already locked the TLB.
+    pub(crate) fn mark_tlb_locked(&mut self, requesting_vtl: Vtl, target_vtl: GuestVtl) {
+        debug_assert_eq!(requesting_vtl, Vtl::Vtl2);
+        debug_assert!(self.is_tlb_locked_in_hypervisor(target_vtl));
+        self.vtls_tlb_locked.set(requesting_vtl, target_vtl, true);
+    }
+
     /// Check the status of the TLB lock of the target VTL on the current VP.
-    pub(crate) fn is_tlb_locked(&mut self, requesting_vtl: Vtl, target_vtl: GuestVtl) -> bool {
+    pub(crate) fn is_tlb_locked(&self, requesting_vtl: Vtl, target_vtl: GuestVtl) -> bool {
+        // This function should only be called in debug assertions.
+        assert!(cfg!(debug_assertions));
         debug_assert_eq!(requesting_vtl, Vtl::Vtl2);
         let local_status = self.vtls_tlb_locked.get(requesting_vtl, target_vtl);
         // The hypervisor may lock the TLB without us knowing, but the inverse should never happen.
         if local_status {
-            debug_assert!(self.is_tlb_locked_in_hypervisor(target_vtl))
-        };
+            debug_assert!(self.is_tlb_locked_in_hypervisor(target_vtl));
+        }
         local_status
     }
 
     fn is_tlb_locked_in_hypervisor(&self, target_vtl: GuestVtl) -> bool {
+        // This function should only be called in debug assertions.
+        assert!(cfg!(debug_assertions));
         let name = HvAllArchRegisterName(
             HvAllArchRegisterName::VsmVpSecureConfigVtl0.0 + target_vtl as u32,
         );
         let result = self
             .partition
             .hcl
-            .get_vp_register(name, HvInputVtl::CURRENT_VTL)
+            .get_vp_register(name, hvdef::hypercall::HvInputVtl::CURRENT_VTL)
             .expect("failure is a misconfiguration");
         let config = hvdef::HvRegisterVsmVpSecureVtlConfig::from(result.as_u64());
         config.tlb_locked()
     }
 
     /// Marks the TLBs of all lower VTLs as unlocked.
     /// The hypervisor does the actual unlocking required upon VTL exit.
-    #[cfg_attr(guest_arch = "aarch64", expect(dead_code))]
     pub(crate) fn unlock_tlb_lock(&mut self, unlocking_vtl: Vtl) {
         debug_assert_eq!(unlocking_vtl, Vtl::Vtl2);
         self.vtls_tlb_locked.fill(unlocking_vtl, false);

openhcl/virt_mshv_vtl/src/processor/mshv/x64.rs

@@ -1163,7 +1163,9 @@ impl<T: CpuIo> EmulatorSupport for UhEmulationState<'_, '_, T, HypervisorBackedX
         }
     }
 
-    fn initial_gva_translation(&self) -> Option<virt_support_x86emu::emulate::InitialTranslation> {
+    fn initial_gva_translation(
+        &mut self,
+    ) -> Option<virt_support_x86emu::emulate::InitialTranslation> {
         if (self.vp.runner.exit_message().header.typ != HvMessageType::HvMessageTypeGpaIntercept)
             && (self.vp.runner.exit_message().header.typ != HvMessageType::HvMessageTypeUnmappedGpa)
             && (self.vp.runner.exit_message().header.typ
@@ -1188,13 +1190,20 @@ impl<T: CpuIo> EmulatorSupport for UhEmulationState<'_, '_, T, HypervisorBackedX
         )
         .expect("unexpected intercept access type");
 
-        tracing::trace!(?message.guest_virtual_address, ?message.guest_physical_address, ?translate_mode, "initial translation");
-
-        Some(virt_support_x86emu::emulate::InitialTranslation {
+        let translation = virt_support_x86emu::emulate::InitialTranslation {
             gva: message.guest_virtual_address,
             gpa: message.guest_physical_address,
             translate_mode,
-        })
+        };
+
+        tracing::trace!(?translation, "initial translation");
+
+        // If we have a valid translation, the hypervisor must have set the TLB lock
+        // so the translation remains valid for the duration of this exit.
+        // Update our local cache appropriately.
+        self.vp.mark_tlb_locked(Vtl::Vtl2, self.vtl);
+
+        Some(translation)
     }
 
     fn interruption_pending(&self) -> bool {
@@ -1225,6 +1234,7 @@ impl<T: CpuIo> EmulatorSupport for UhEmulationState<'_, '_, T, HypervisorBackedX
             )
         {
             // Should always be called after translate gva with the tlb lock flag
+            // or with an initial translation.
             debug_assert!(self.vp.is_tlb_locked(Vtl::Vtl2, self.vtl));
 
             let mbec_user_execute = self
@@ -1285,7 +1295,6 @@ impl<T: CpuIo> EmulatorSupport for UhEmulationState<'_, '_, T, HypervisorBackedX
         // remains usable until the VP is resumed back to direct execution.
         control_flags.set_set_page_table_bits(true);
         control_flags.set_tlb_flush_inhibit(true);
-        self.vp.set_tlb_lock(Vtl::Vtl2, target_vtl);
 
         // In case we're not running ring 0, check privileges against VP state
         // as of when the original intercept came in - since the emulator
@@ -1307,10 +1316,13 @@ impl<T: CpuIo> EmulatorSupport for UhEmulationState<'_, '_, T, HypervisorBackedX
             Ok(ioctl::TranslateResult {
                 gpa_page,
                 overlay_page,
-            }) => Ok(Ok(EmuTranslateResult {
-                gpa: (gpa_page << hvdef::HV_PAGE_SHIFT) + (gva & (HV_PAGE_SIZE - 1)),
-                overlay_page: Some(overlay_page),
-            })),
+            }) => {
+                self.vp.mark_tlb_locked(Vtl::Vtl2, GuestVtl::Vtl0);
+                Ok(Ok(EmuTranslateResult {
+                    gpa: (gpa_page << hvdef::HV_PAGE_SHIFT) + (gva & (HV_PAGE_SIZE - 1)),
+                    overlay_page: Some(overlay_page),
+                }))
+            }
             Err(ioctl::x64::TranslateErrorX64 { code, event_info }) => Ok(Err(EmuTranslateError {
                 code: hypercall::TranslateGvaResultCode(code),
                 event_info: Some(event_info),

openhcl/virt_mshv_vtl/src/processor/snp/mod.rs

@@ -1735,7 +1735,9 @@ impl<T: CpuIo> X86EmulatorSupport for UhEmulationState<'_, '_, T, SnpBacked> {
         Some(self.vp.runner.vmsa(self.vtl).exit_info2())
     }
 
-    fn initial_gva_translation(&self) -> Option<virt_support_x86emu::emulate::InitialTranslation> {
+    fn initial_gva_translation(
+        &mut self,
+    ) -> Option<virt_support_x86emu::emulate::InitialTranslation> {
         None
     }
 

openhcl/virt_mshv_vtl/src/processor/tdx/mod.rs

@@ -2792,7 +2792,9 @@ impl<T: CpuIo> X86EmulatorSupport for UhEmulationState<'_, '_, T, TdxBacked> {
         Some(TdxExit(self.vp.runner.tdx_vp_enter_exit_info()).gpa())
     }
 
-    fn initial_gva_translation(&self) -> Option<virt_support_x86emu::emulate::InitialTranslation> {
+    fn initial_gva_translation(
+        &mut self,
+    ) -> Option<virt_support_x86emu::emulate::InitialTranslation> {
         let exit_info = TdxExit(self.vp.runner.tdx_vp_enter_exit_info());
         let ept_info = VmxEptExitQualification::from(exit_info.qualification());
 

vmm_core/virt_mshv/src/lib.rs

@@ -935,7 +935,9 @@ impl EmulatorSupport for MshvEmulationState<'_> {
         }
     }
 
-    fn initial_gva_translation(&self) -> Option<virt_support_x86emu::emulate::InitialTranslation> {
+    fn initial_gva_translation(
+        &mut self,
+    ) -> Option<virt_support_x86emu::emulate::InitialTranslation> {
         if (self.message.header.message_type != HvMessageType::HvMessageTypeGpaIntercept.0)
             && (self.message.header.message_type != HvMessageType::HvMessageTypeUnmappedGpa.0)
             && (self.message.header.message_type != HvMessageType::HvMessageTypeUnacceptedGpa.0)

vmm_core/virt_support_aarch64emu/src/emulate.rs

@@ -35,7 +35,7 @@ pub trait EmulatorSupport: AccessCpuState {
     fn physical_address(&self) -> Option<u64>;
 
     /// The gva translation included in the intercept message header, if valid.
-    fn initial_gva_translation(&self) -> Option<InitialTranslation>;
+    fn initial_gva_translation(&mut self) -> Option<InitialTranslation>;
 
     /// If interrupt pending is marked in the intercept message
     fn interruption_pending(&self) -> bool;
@@ -98,6 +98,7 @@ pub struct EmuTranslateResult {
 }
 
 /// The translation, if any, provided in the intercept message and provided by [`EmulatorSupport`].
+#[derive(Debug)]
 pub struct InitialTranslation {
     /// GVA for the translation
     pub gva: u64,

vmm_core/virt_support_x86emu/src/emulate.rs

@@ -79,7 +79,7 @@ pub trait EmulatorSupport {
     fn physical_address(&self) -> Option<u64>;
 
     /// The gva translation included in the intercept message header, if valid.
-    fn initial_gva_translation(&self) -> Option<InitialTranslation>;
+    fn initial_gva_translation(&mut self) -> Option<InitialTranslation>;
 
     /// If interrupt pending is marked in the intercept message
     fn interruption_pending(&self) -> bool;
@@ -194,6 +194,7 @@ pub struct EmuTranslateResult {
 }
 
 /// The translation, if any, provided in the intercept message and provided by [`EmulatorSupport`].
+#[derive(Debug)]
 pub struct InitialTranslation {
     /// GVA for the translation
     pub gva: u64,

vmm_core/virt_support_x86emu/tests/tests/checkvtlaccess.rs

@@ -118,7 +118,7 @@ impl EmulatorSupport for MockSupport {
     }
 
     /// The gva translation included in the intercept message header, if valid.
-    fn initial_gva_translation(&self) -> Option<InitialTranslation> {
+    fn initial_gva_translation(&mut self) -> Option<InitialTranslation> {
         None
     }
 

vmm_core/virt_support_x86emu/tests/tests/emulate.rs

@@ -99,7 +99,9 @@ impl EmulatorSupport for MockSupport {
     }
 
     /// The gva translation included in the intercept message header, if valid.
-    fn initial_gva_translation(&self) -> Option<virt_support_x86emu::emulate::InitialTranslation> {
+    fn initial_gva_translation(
+        &mut self,
+    ) -> Option<virt_support_x86emu::emulate::InitialTranslation> {
         None
     }
 

vmm_core/virt_support_x86emu/tests/tests/translate.rs

@@ -283,7 +283,7 @@ impl EmulatorSupport for MockSupport {
     }
 
     /// The gva translation included in the intercept message header, if valid.
-    fn initial_gva_translation(&self) -> Option<InitialTranslation> {
+    fn initial_gva_translation(&mut self) -> Option<InitialTranslation> {
         Some(InitialTranslation {
             gva: INITIAL_GVA,
             gpa: INITIAL_GPA,

vmm_core/virt_whp/src/emu.rs

@@ -197,7 +197,9 @@ impl<T: CpuIo> virt_support_x86emu::emulate::EmulatorSupport for WhpEmulationSta
         }
     }
 
-    fn initial_gva_translation(&self) -> Option<virt_support_x86emu::emulate::InitialTranslation> {
+    fn initial_gva_translation(
+        &mut self,
+    ) -> Option<virt_support_x86emu::emulate::InitialTranslation> {
         if let WhpVpRefEmulation::MemoryAccessContext(access) = self.access {
             if !(access.AccessInfo.GvaValid()) {
                 return None;

vmm_core/virt_whp/src/hypercalls.rs

@@ -1531,8 +1531,8 @@ mod x86 {
                         return Err(HvError::AccessDenied);
                     }
 
-                    u64::from(HvRegisterVsmVpSecureVtlConfig::new().with_tlb_locked(self.tlb_lock))
-                        .into()
+                    // WHP locks the TLB on every exit, so this is already locked.
+                    u64::from(HvRegisterVsmVpSecureVtlConfig::new().with_tlb_locked(true)).into()
                 }
                 reg => {
                     if let Ok(name) = regs::hv_register_to_whp(reg) {
@@ -1629,8 +1629,11 @@ mod x86 {
                         return Err(HvError::AccessDenied);
                     }
 
-                    self.tlb_lock =
-                        HvRegisterVsmVpSecureVtlConfig::from(value.as_u64()).tlb_locked();
+                    // WHP locks the TLB on every exit, so this is already locked.
+                    // Make sure the guest isn't trying to unlock.
+                    if !HvRegisterVsmVpSecureVtlConfig::from(value.as_u64()).tlb_locked() {
+                        return Err(HvError::InvalidParameter);
+                    }
                 }
 
                 reg => {

vmm_core/virt_whp/src/lib.rs

@@ -624,8 +624,6 @@ impl virt::BindProcessor for WhpProcessorBinder {
                 .vtl2
                 .as_ref()
                 .map(|vtl2| &vtl2.vplcs[self.index.index() as usize]),
-
-            tlb_lock: false,
         };
 
         let vp_info = &vp.inner.vp_info;
@@ -677,11 +675,6 @@ pub struct WhpProcessor<'a> {
     state: RunState,
     vplc0: &'a Vplc,
     vplc2: Option<&'a Vplc>,
-
-    /// Whether the VTL 0 TLB is locked by VTL 2 or not.
-    // TODO: This doesn't actually control anything, we just
-    // track it so we can report it back correctly when asked.
-    tlb_lock: bool,
 }
 
 #[derive(Copy, Clone)]

vmm_core/virt_whp/src/vp.rs

@@ -134,11 +134,6 @@ impl<'a> WhpProcessor<'a> {
 
             self.state.halted = false;
 
-            // Pretend to unlock the TLB when we return to VTL 0
-            if new_vtl == Vtl::Vtl0 {
-                self.tlb_lock = false;
-            }
-
             if new_vtl == Vtl::Vtl2 {
                 // No need to schedule any wakeups until the next VTL0 entry.
                 self.state