Merge pull request #186 from Freax13/fix/tlb-flush

 fix race condition in TLB flushing code

Author: Freax13

common/constants/src/lib.rs

@@ -5,7 +5,7 @@
 use core::{
     fmt::{self, Debug, Display},
     marker::PhantomData,
-    ops::{BitAnd, BitAndAssign, BitOrAssign, Index, Not, RangeInclusive},
+    ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign, Index, Not, RangeInclusive},
     sync::atomic::{AtomicU32, Ordering},
 };
 
@@ -165,10 +165,18 @@ impl ApBitmap {
     }
 }
 
+impl BitOr for ApBitmap {
+    type Output = Self;
+
+    fn bitor(self, rhs: Self) -> Self::Output {
+        Self(self.0 | rhs.0)
+    }
+}
+
 impl BitOrAssign for ApBitmap {
     #[inline]
     fn bitor_assign(&mut self, rhs: Self) {
-        self.0 |= rhs.0;
+        *self = *self | rhs;
     }
 }
 

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

@@ -205,6 +205,8 @@ impl ActivePageTableGuard {
             }
             if num_pages < usize::from(invlpgb.invlpgb_count_max()) {
                 builder = builder.pages(Page::range(*pages.start(), *pages.end() + 1));
+            } else {
+                state.needs_flush = ApBitmap::empty();
             }
             builder.flush();
             invlpgb.tlbsync();
@@ -298,18 +300,30 @@ impl FlushGuard for GlobalFlushGuard {
         // IPI, so the other thread thinks that the bit was never cleared and
         // it won't get cleared again because we didn't set another IPI.
         let active_aps = ACTIVE_APS.get_all();
-        let other_active_aps = active_aps & all_other_aps;
         let inactive_aps = !active_aps & all_other_aps;
-        PENDING_GLOBAL_TLB_SHOOTDOWN.set_all(other_active_aps);
         LAZY_PENDING_GLOBAL_TLB_SHOOTDOWN.set_all(inactive_aps);
+        // Re-fetch `ACTIVE_APS` after setting `LAZY_PENDING_GLOBAL_TLB_SHOOTDOWN`.
+        // If we don't do this, there's a race condition:
+        // +----------------------------------------------------------------------------------+
+        // | this AP                                 | other AP                               |
+        // | ...                                     | inactive                               |
+        // | read ACTIVE_APS                         |                                        |
+        // |                                         | wake up                                |
+        // |                                         | set ACTIVE_APS                         |
+        // |                                         | read LAZY_PENDING_GLOBAL_TLB_SHOOTDOWN |
+        // | write LAZY_PENDING_GLOBAL_TLB_SHOOTDOWN |                                        |
+        // +----------------------------------------------------------------------------------+
+        let active_aps = active_aps | ACTIVE_APS.get_all();
+        let other_active_aps = active_aps & all_other_aps;
+        PENDING_GLOBAL_TLB_SHOOTDOWN.set_all(other_active_aps);
 
         // Send IPIs to all other currently active APs.
         send_tlb_ipis(other_active_aps);
 
         // Flush the local TLB entry.
         tlb::flush(page.start_address());
 
-        // Wait for the APS to acknowledge the IPI.
+        // Wait for the APs to acknowledge the IPI.
         let mut remaining_aps = other_active_aps;
         while !remaining_aps.is_empty() {
             remaining_aps &= PENDING_GLOBAL_TLB_SHOOTDOWN.get_all();