From a2c05de60482fc5ec0bc3af878cca11bb7c8ad67 Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Mon, 17 Oct 2022 19:32:24 +0200
Subject: [PATCH] mmap(), mprotect(), munmap(): Check more stuff

---
 kernel/src/sys/mem.cpp | 44 ++++++++++++++++++++++++++++++++----------
 1 file changed, 34 insertions(+), 10 deletions(-)

diff --git a/kernel/src/sys/mem.cpp b/kernel/src/sys/mem.cpp
index 6712d118..50018740 100644
--- a/kernel/src/sys/mem.cpp
+++ b/kernel/src/sys/mem.cpp
@@ -3,6 +3,7 @@
 #include "errno.h"
 #include "interrupts/Context.h"
 #include "log/Log.h"
+#include "memory/Memory.h"
 #include "memory/MemoryManager.h"
 #include "memory/VMM.h"
 #include "misc/utils.h"
@@ -50,6 +51,12 @@ void sys_mmap(Context* context, void* address, size_t size, int prot)
     if (address)
     {
         kdbgln("mmap(): %ld pages at address %p, %s", size / PAGE_SIZE, address, format_prot(prot));
+        if (Memory::is_kernel_address((uintptr_t)address))
+        {
+            kwarnln("munmap() failed: attempted to unmap a kernel page");
+            context->rax = MAP_FAIL(ENOMEM);
+            return;
+        }
         if (VMM::get_physical((uint64_t)address) != (uint64_t)-1) // Address is already used.
         {
             kwarnln("attempt to map an already mapped address");
@@ -67,7 +74,7 @@ void sys_mmap(Context* context, void* address, size_t size, int prot)
         }
         else
         {
-            kwarnln("mmap() failed");
+            kwarnln("mmap() failed: failed to allocate physical memory");
             context->rax = MAP_FAIL(ENOMEM);
             return;
         }
@@ -75,6 +82,12 @@ void sys_mmap(Context* context, void* address, size_t size, int prot)
     kdbgln("mmap(): %ld pages at any address, %s", Utilities::get_blocks_from_size(PAGE_SIZE, size), format_prot(prot));
     uint64_t ptr =
         Scheduler::current_task()->allocator.request_virtual_pages(Utilities::get_blocks_from_size(PAGE_SIZE, size));
+    if (!ptr)
+    {
+        kwarnln("mmap() failed: failed to allocate virtual address");
+        context->rax = MAP_FAIL(ENOMEM);
+        return;
+    }
     void* result = MemoryManager::get_pages_at(ptr, Utilities::get_blocks_from_size(PAGE_SIZE, size), real_flags);
     if (result)
     {
@@ -84,7 +97,7 @@ void sys_mmap(Context* context, void* address, size_t size, int prot)
     }
     else
     {
-        kwarnln("mmap() failed");
+        kwarnln("mmap() failed: failed to allocate physical memory");
         context->rax = MAP_FAIL(ENOMEM);
         return;
     }
@@ -111,10 +124,16 @@ void sys_munmap(Context* context, void* address, size_t size)
         context->rax = -EINVAL;
         return;
     }
-    uint64_t flags = VMM::get_flags((uint64_t)address);
-    if (!(flags & MAP_USER))
+    if (Memory::is_kernel_address((uintptr_t)address))
     {
-        kwarnln("munmap() failed: attempted to unmap a non-existent or kernel page");
+        kwarnln("munmap() failed: attempted to unmap a kernel page");
+        context->rax = -EINVAL;
+        return;
+    }
+    uint64_t phys = VMM::get_physical((uint64_t)address);
+    if (phys == (uint64_t)-1)
+    {
+        kwarnln("munmap() failed: attempted to unmap a non-existent page");
         context->rax = -EINVAL;
         return;
     }
@@ -145,15 +164,20 @@ void sys_mprotect(Context* context, void* address, size_t size, int prot)
     }
     if (!address)
     {
-        kwarnln("mprotect() failed: attempted to unmap page 0");
+        kwarnln("mprotect() failed: attempted to protect page 0");
         context->rax = -EINVAL;
         return;
     }
-
-    uint64_t flags = VMM::get_flags((uint64_t)address);
-    if (!(flags & MAP_USER))
+    if (Memory::is_kernel_address((uintptr_t)address))
     {
-        kwarnln("mprotect() failed: attempted to protect a non-existent or kernel page");
+        kwarnln("mprotect() failed: attempted to protect a kernel page");
+        context->rax = -EINVAL;
+        return;
+    }
+    uint64_t phys = VMM::get_physical((uint64_t)address);
+    if (phys == (uint64_t)-1)
+    {
+        kwarnln("mprotect() failed: attempted to protect a non-existent page");
         context->rax = -EINVAL;
         return;
     }