From 8cd8cf39382444195dbe2b0c6e1c7d6eebbc6452 Mon Sep 17 00:00:00 2001
From: Zhang Junyang <junyang@stu.pku.edu.cn>
Date: Thu, 27 Jun 2024 08:39:37 +0000
Subject: [PATCH] Add missing TLB flushes for the kernel stack guard pages

---
 ostd/src/arch/x86/mm/mod.rs |  7 +++++++
 ostd/src/task/task.rs       | 13 +++++++------
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/ostd/src/arch/x86/mm/mod.rs b/ostd/src/arch/x86/mm/mod.rs
index f2ea6d12e..93d38baf7 100644
--- a/ostd/src/arch/x86/mm/mod.rs
+++ b/ostd/src/arch/x86/mm/mod.rs
@@ -63,20 +63,27 @@ bitflags::bitflags! {
     }
 }
 
+/// Flush any TLB entry that contains the map of the given virtual address.
+///
+/// This flush performs regardless of the global-page bit. So it can flush both global
+/// and non-global entries.
 pub(crate) fn tlb_flush_addr(vaddr: Vaddr) {
     tlb::flush(VirtAddr::new(vaddr as u64));
 }
 
+/// Flush any TLB entry that intersects with the given address range.
 pub(crate) fn tlb_flush_addr_range(range: &Range<Vaddr>) {
     for vaddr in range.clone().step_by(PAGE_SIZE) {
         tlb_flush_addr(vaddr);
     }
 }
 
+/// Flush all TLB entries except for the global-page entries.
 pub(crate) fn tlb_flush_all_excluding_global() {
     tlb::flush_all();
 }
 
+/// Flush all TLB entries, including global-page entries.
 pub(crate) fn tlb_flush_all_including_global() {
     // SAFETY: updates to CR4 here only change the global-page bit, the side effect
     // is only to invalidate the TLB, which doesn't affect the memory safety.
diff --git a/ostd/src/task/task.rs b/ostd/src/task/task.rs
index 036359be7..e18b07995 100644
--- a/ostd/src/task/task.rs
+++ b/ostd/src/task/task.rs
@@ -16,6 +16,7 @@ use super::{
 };
 pub(crate) use crate::arch::task::{context_switch, TaskContext};
 use crate::{
+    arch::mm::tlb_flush_addr_range,
     cpu::CpuSet,
     mm::{kspace::KERNEL_PAGE_TABLE, FrameAllocOptions, PageFlags, Segment, PAGE_SIZE},
     prelude::*,
@@ -66,11 +67,11 @@ impl KernelStack {
         };
         // SAFETY: the segment allocated is not used by others so we can protect it.
         unsafe {
+            let vaddr_range = guard_page_vaddr..guard_page_vaddr + PAGE_SIZE;
             page_table
-                .protect(&(guard_page_vaddr..guard_page_vaddr + PAGE_SIZE), |p| {
-                    p.flags -= PageFlags::RW
-                })
+                .protect(&vaddr_range, |p| p.flags -= PageFlags::RW)
                 .unwrap();
+            tlb_flush_addr_range(&vaddr_range);
         }
         Ok(Self {
             segment: stack_segment,
@@ -94,11 +95,11 @@ impl Drop for KernelStack {
             };
             // SAFETY: the segment allocated is not used by others so we can protect it.
             unsafe {
+                let vaddr_range = guard_page_vaddr..guard_page_vaddr + PAGE_SIZE;
                 page_table
-                    .protect(&(guard_page_vaddr..guard_page_vaddr + PAGE_SIZE), |p| {
-                        p.flags |= PageFlags::RW
-                    })
+                    .protect(&vaddr_range, |p| p.flags |= PageFlags::RW)
                     .unwrap();
+                tlb_flush_addr_range(&vaddr_range);
             }
         }
     }