TLB flushing fixes

host/mushroom/src/insecure.rs

@@ -16,7 +16,7 @@ use anyhow::{ensure, Context, Result};
 use bit_field::BitField;
 use constants::{
     physical_address::{kernel, supervisor, DYNAMIC_2MIB, SUPERVISOR_SERVICES},
-    FIRST_AP, MAX_APS_COUNT,
+    MAX_APS_COUNT,
 };
 use loader::Input;
 use nix::sys::pthread::pthread_kill;
@@ -154,10 +154,7 @@ pub fn main(
     let cpuid_entries = Arc::<[KvmCpuidEntry2]>::from(cpuid_entries);
     let done = Arc::new(AtomicBool::new(false));
     let ap_threads = (0..MAX_APS_COUNT)
-        .map(|i| {
-            let id = FIRST_AP + i;
-            run_kernel_vcpu(id, vm.clone(), cpuid_entries.clone(), done.clone())
-        })
+        .map(|i| run_kernel_vcpu(i, vm.clone(), cpuid_entries.clone(), done.clone()))
         .collect::<Vec<_>>();
     ap_threads[0].thread().unpark();
     start_log_collection(&memory_slots, kernel::LOG_BUFFER)?;
@@ -234,8 +231,8 @@ fn run_kernel_vcpu(
 ) -> JoinHandle<()> {
     let supervisor_thread = std::thread::current();
 
+    let ap = vm.create_vcpu(i32::from(id)).unwrap();
     std::thread::spawn(move || {
-        let ap = vm.create_vcpu(i32::from(id)).unwrap();
         ap.set_cpuid(&cpuid_entries).unwrap();
 
         let kvm_run = ap.get_kvm_run_block().unwrap();

tee/kernel/src/memory/pagetable.rs

@@ -454,7 +454,6 @@ impl Pagetables {
         if needs_flush {
             flush_state_guard.needs_flush.set(ap_index, false);
         }
-        drop(flush_state_guard);
 
         if update_required || needs_flush {
             if let Some(pcid_allocation) = allocations.pcid_allocation.as_ref() {
@@ -472,10 +471,11 @@ impl Pagetables {
                     Cr3::write(allocations.pml4, Cr3Flags::empty());
                 }
             }
+        }
+        drop(flush_state_guard);
 
-            if update_required {
-                *guard = Some(allocations.clone());
-            }
+        if update_required {
+            *guard = Some(allocations.clone());
         }
 
         let guard = RefMut::map(guard, |a| a.as_mut().unwrap());

tee/kernel/src/memory/pagetable/flush.rs

@@ -189,6 +189,9 @@ impl ActivePageTableGuard {
         let mut guard = self.guard.flush_state.lock();
         let state = &mut *guard;
 
+        let mut other_used = state.used;
+        other_used.set(idx, false);
+
         // Check if the current AP is the only active AP.
         let mut other_active_aps = state.active;
         other_active_aps.set(idx, false);
@@ -197,6 +200,7 @@ impl ActivePageTableGuard {
         // APs to flush immediately. We only need to flush the TLB on the
         // current AP.
         if other_active_aps.is_empty() {
+            state.needs_flush |= other_used;
             drop(guard);
             return self.flush_pages_local(pages);
         }
@@ -233,7 +237,7 @@ impl ActivePageTableGuard {
         // We've run out of optimizations :(
         // Flush on the current processor and send IPIs to the other relevant
         // APs.
-        state.needs_flush |= state.used;
+        state.needs_flush |= other_used;
         drop(guard);
 
         PENDING_TLB_SHOOTDOWN.set_all(other_active_aps);