From 00e126f7d2de4278db78dac173ee9c200286453e Mon Sep 17 00:00:00 2001
From: Pawel Wieczorkiewicz <wipawel@grsecurity.net>
Date: Thu, 2 Nov 2023 11:57:56 +0100
Subject: [PATCH] arch,pt: add frame finder for virtual address functionality

Signed-off-by: Pawel Wieczorkiewicz <wipawel@grsecurity.net>
---
 arch/x86/pagetables.c        | 20 ++++++++++++++++++++
 include/arch/x86/pagetable.h |  3 +++
 tests/unittests.c            |  8 ++++++--
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/arch/x86/pagetables.c b/arch/x86/pagetables.c
index b82face7..06258a12 100644
--- a/arch/x86/pagetables.c
+++ b/arch/x86/pagetables.c
@@ -523,6 +523,26 @@ int get_user_va_mfn_order(void *va, mfn_t *mfn, unsigned int *order) {
     return err;
 }
 
+static frame_t *find_va_frame(const cr3_t *cr3_ptr, const void *va) {
+    unsigned int order;
+    mfn_t mfn;
+    int err;
+
+    spin_lock(&vmap_lock);
+    err = get_va_mfn_order(cr3_ptr, va, &mfn, &order);
+    spin_unlock(&vmap_lock);
+
+    return err ? NULL : find_mfn_frame(mfn, order);
+}
+
+frame_t *find_kern_va_frame(const void *va) {
+    return find_va_frame(&cr3, va);
+}
+
+frame_t *find_user_va_frame(const void *va) {
+    return find_va_frame(&user_cr3, va);
+}
+
 static inline void init_cr3(cr3_t *cr3_ptr) {
     memset(cr3_ptr, 0, sizeof(*cr3_ptr));
     cr3_ptr->mfn = MFN_INVALID;
diff --git a/include/arch/x86/pagetable.h b/include/arch/x86/pagetable.h
index fb95e444..405d7973 100644
--- a/include/arch/x86/pagetable.h
+++ b/include/arch/x86/pagetable.h
@@ -304,6 +304,9 @@ extern void dump_user_pagetable_va(void *va);
 extern int get_kern_va_mfn_order(void *va, mfn_t *mfn, unsigned int *order);
 extern int get_user_va_mfn_order(void *va, mfn_t *mfn, unsigned int *order);
 
+extern frame_t *find_kern_va_frame(const void *va);
+extern frame_t *find_user_va_frame(const void *va);
+
 #endif /* __ASSEMBLY__ */
 
 #endif /* KTF_PAGETABLE_H */
diff --git a/tests/unittests.c b/tests/unittests.c
index 82d114b8..28e9b3cf 100644
--- a/tests/unittests.c
+++ b/tests/unittests.c
@@ -26,6 +26,7 @@
 #include <console.h>
 #include <cpuid.h>
 #include <ktf.h>
+#include <pagetable.h>
 #include <real_mode.h>
 #include <sched.h>
 #include <string.h>
@@ -203,13 +204,16 @@ int unit_tests(void *_unused) {
     task_user3 = new_user_task("test3 user", test_user_task_func3, NULL);
     task_user4 = new_user_task("test4 user", test_user_task_func4, NULL);
 
-    vmap_4k(HIGH_USER_PTR + 0x1000, get_free_frame()->mfn, L1_PROT);
+    frame_t *frame = get_free_frame();
+    vmap_4k(HIGH_USER_PTR + 0x1000, frame->mfn, L1_PROT);
     memset(HIGH_USER_PTR + 0x1000, 0, 0x1000);
-    vmap_user_4k(HIGH_USER_PTR, get_free_frame()->mfn, L1_PROT_USER);
+    vmap_user_4k(HIGH_USER_PTR, frame->mfn, L1_PROT_USER);
 
     /* Be sure that we can still touch this vmap despite the user vmap. */
     BUG_ON(*(unsigned long *) (HIGH_USER_PTR + 0x1000) != 0);
 
+    BUG_ON(frame != find_kern_va_frame(HIGH_USER_PTR + 0x1000));
+
     set_task_repeat(task1, 10);
     schedule_task(task1, get_bsp_cpu());
     schedule_task(task2, get_cpu(1));