From 1f01ddfc6dfe2951110d12133b0481bc5edf4cc3 Mon Sep 17 00:00:00 2001
From: xvanc <xvancm@gmail.com>
Date: Sun, 4 Jun 2023 13:16:25 -0500
Subject: [PATCH] limine: scale the maximum KASLR offset based on the address
 space size

---
 common/mm/vmm.c        | 13 +++++++++++++
 common/mm/vmm.h        |  4 ++++
 common/protos/limine.c |  3 ++-
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/common/mm/vmm.c b/common/mm/vmm.c
index eca817b9..4ee192fa 100644
--- a/common/mm/vmm.c
+++ b/common/mm/vmm.c
@@ -273,6 +273,19 @@ uint64_t paging_mode_higher_half(int paging_mode) {
     }
 }
 
+int paging_mode_va_bits(int paging_mode) {
+    switch (paging_mode) {
+        case PAGING_MODE_RISCV_SV39:
+            return 39;
+        case PAGING_MODE_RISCV_SV48:
+            return 48;
+        case PAGING_MODE_RISCV_SV57:
+            return 57;
+        default:
+            panic(false, "paging_mode_va_bits: invalid mode");
+    }
+}
+
 int vmm_max_paging_mode(void)
 {
     static int max_level;
diff --git a/common/mm/vmm.h b/common/mm/vmm.h
index cddbd1eb..a435818d 100644
--- a/common/mm/vmm.h
+++ b/common/mm/vmm.h
@@ -15,6 +15,8 @@
 #define PAGING_MODE_X86_64_4LVL 0
 #define PAGING_MODE_X86_64_5LVL 1
 
+#define paging_mode_va_bits(mode) ((mode) ? 57 : 48)
+
 static inline uint64_t paging_mode_higher_half(int paging_mode) {
     if (paging_mode == PAGING_MODE_X86_64_5LVL) {
         return 0xff00000000000000;
@@ -89,6 +91,8 @@ void map_page(pagemap_t pagemap, uint64_t virt_addr, uint64_t phys_addr, uint64_
 #define PAGING_MODE_RISCV_SV48 9
 #define PAGING_MODE_RISCV_SV57 10
 
+int paging_mode_va_bits(int paging_mode);
+
 enum page_size {
     Size4KiB,
     Size2MiB,
diff --git a/common/protos/limine.c b/common/protos/limine.c
index 3f2862be..46c08850 100644
--- a/common/protos/limine.c
+++ b/common/protos/limine.c
@@ -194,7 +194,8 @@ static void **requests;
 static void set_paging_mode(int paging_mode, bool kaslr) {
     direct_map_offset = paging_mode_higher_half(paging_mode);
     if (kaslr) {
-        direct_map_offset += (rand64() & ~((uint64_t)0x40000000 - 1)) & 0xfffffffffff;
+        uint64_t mask = ((uint64_t)1 << (paging_mode_va_bits(paging_mode) - 4)) - 1;
+        direct_map_offset += (rand64() & ~((uint64_t)0x40000000 - 1)) & mask;
     }
 }