From 001550b387cd0fa4804a20cac7b842beafa69700 Mon Sep 17 00:00:00 2001
From: Paul Pan <i@0x7f.app>
Date: Mon, 20 May 2024 13:50:58 +0800
Subject: [PATCH] feat: vspace: addr: add is_kernel

---
 kernel/src/arch/riscv/vspace/addr/mod.rs  | 48 +++++++++++++++++++----
 kernel/src/arch/riscv/vspace/addr/sv39.rs | 10 ++++-
 kernel/src/vspace/addr.rs                 |  6 +++
 3 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/kernel/src/arch/riscv/vspace/addr/mod.rs b/kernel/src/arch/riscv/vspace/addr/mod.rs
index 4555a04..9e9e74e 100644
--- a/kernel/src/arch/riscv/vspace/addr/mod.rs
+++ b/kernel/src/arch/riscv/vspace/addr/mod.rs
@@ -1,9 +1,5 @@
-use utils::{
-    addr::{AddressOps, VirtAddr},
-    MASK,
-};
-
-use crate::vspace::TableLevel;
+use crate::{arch::layout::*, vspace::TableLevel};
+use utils::{addr::*, MASK};
 
 #[cfg(feature = "riscv.pagetable.sv39")]
 mod sv39;
@@ -27,6 +23,13 @@ trait GenericPhysAddrPage: AddressOps {
     const PA_PPN_MASK: usize = MASK!(Self::PPN_BITS) << Self::PG_OFFSET;
     const PTE_PPN_MASK: usize = MASK!(Self::PPN_BITS) << Self::PPN_OFFSET;
 
+    fn is_kernel(&self) -> bool {
+        let kernel_start = unsafe { kernel_virt_to_phys(KERNEL_START.as_virt_addr()) };
+        let kernel_end = unsafe { kernel_virt_to_phys(KERNEL_END.as_virt_addr()).align_up(PAGE_SIZE) };
+
+        self.as_usize() >= kernel_start.as_usize() && self.as_usize() < kernel_end.as_usize()
+    }
+
     fn extract_ppn(&self) -> usize {
         (self.as_usize() & Self::PA_PPN_MASK) >> Self::PG_OFFSET
     }
@@ -51,16 +54,45 @@ trait GenericVirtAddrPage: AddressOps {
     const PG_OFFSET: usize;
     const MAX_LEVEL: usize;
 
+    fn is_kernel(&self, bits: usize) -> bool {
+        let mask = !MASK!(bits + 1);
+        (self.as_usize() & mask) != 0
+    }
+
     fn extract_vpn<T: TableLevel>(&self) -> usize {
         let mask = MASK!(T::LEVEL_BITS);
         (self.as_usize() >> (Self::PG_OFFSET + T::LEVEL_BITS * (Self::MAX_LEVEL - T::LEVEL))) & mask
     }
 
-    fn merge_vpn<T: TableLevel>(&self, vpn: usize) -> VirtAddr {
+    fn merge_vpn<T: TableLevel>(&self, vpn: usize, bits: usize) -> VirtAddr {
         let shift = Self::PG_OFFSET + T::LEVEL_BITS * (Self::MAX_LEVEL - T::LEVEL);
         let mask = MASK!(T::LEVEL_BITS);
         let addr = (self.as_usize() & !(mask << shift)) | ((vpn & mask) << shift);
 
-        VirtAddr(sign_extend(addr, 39 - 1))
+        VirtAddr(sign_extend(addr, bits))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test_case]
+    fn test_sign_extend() {
+        assert_eq!(sign_extend(0b01010, 4), 0b01010);
+        assert_eq!(sign_extend(0b01010, 3), !0b101);
+    }
+
+    #[test_case]
+    fn test_is_kernel() {
+        let addr = PhysAddr(0x80200000);
+        assert!(addr.is_kernel());
+        let addr = PhysAddr(0x80000000);
+        assert!(!addr.is_kernel());
+
+        let addr = VirtAddr(0x0000_0080_0000_0000);
+        assert!(addr.is_kernel(39 - 1));
+        let addr = VirtAddr(0x0000_007f_ffff_ffff);
+        assert!(!addr.is_kernel(39 - 1));
     }
 }
diff --git a/kernel/src/arch/riscv/vspace/addr/sv39.rs b/kernel/src/arch/riscv/vspace/addr/sv39.rs
index c926b03..a16f1db 100644
--- a/kernel/src/arch/riscv/vspace/addr/sv39.rs
+++ b/kernel/src/arch/riscv/vspace/addr/sv39.rs
@@ -9,6 +9,10 @@ impl GenericPhysAddrPage for PhysAddr {
 }
 
 impl PhysAddrPage for PhysAddr {
+    fn is_kerenel(&self) -> bool {
+        GenericPhysAddrPage::is_kernel(self)
+    }
+
     fn extract_ppn(&self) -> usize {
         GenericPhysAddrPage::extract_ppn(self)
     }
@@ -32,11 +36,15 @@ impl GenericVirtAddrPage for VirtAddr {
 }
 
 impl VirtAddrPage for VirtAddr {
+    fn is_kerenel(&self) -> bool {
+        GenericVirtAddrPage::is_kernel(self, 39 - 1)
+    }
+
     fn extract_vpn<T: TableLevel>(&self) -> usize {
         GenericVirtAddrPage::extract_vpn::<T>(self)
     }
 
     fn merge_vpn<T: TableLevel>(&self, vpn: usize) -> usize {
-        GenericVirtAddrPage::merge_vpn::<T>(self, vpn).into()
+        GenericVirtAddrPage::merge_vpn::<T>(self, vpn, 39 - 1).into()
     }
 }
diff --git a/kernel/src/vspace/addr.rs b/kernel/src/vspace/addr.rs
index b7fbda3..1cd6382 100644
--- a/kernel/src/vspace/addr.rs
+++ b/kernel/src/vspace/addr.rs
@@ -13,6 +13,9 @@ use super::TableLevel;
 use utils::addr::{AddressOps, PhysAddr, VirtAddr};
 
 pub trait PhysAddrPage: AddressOps {
+    /// Checks if the address is a kernel address.
+    fn is_kerenel(&self) -> bool;
+
     /// Extracts the Physical Page Number (PPN).
     fn extract_ppn(&self) -> usize;
 
@@ -27,6 +30,9 @@ pub trait PhysAddrPage: AddressOps {
 }
 
 pub trait VirtAddrPage: AddressOps {
+    /// Checks if the address is a kernel address.
+    fn is_kerenel(&self) -> bool;
+
     /// Extracts the Virtual Page Number (VPN).
     fn extract_vpn<T: TableLevel>(&self) -> usize;