From acf9a439806552cf4e765005489f1081753c2d06 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Thomas=20Lov=C3=A9n?= <thomasloven@gmail.com>
Date: Wed, 27 Dec 2017 23:30:14 +0100
Subject: [PATCH] VMM - touch_page fails if out of pages

---
 src/kernel/memory/vmm.c  | 22 +++++++++++++++-------
 src/kernel/memory/vmm.tt |  7 +++++++
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/kernel/memory/vmm.c b/src/kernel/memory/vmm.c
index 707fe41..e0bd1c3 100644
--- a/src/kernel/memory/vmm.c
+++ b/src/kernel/memory/vmm.c
@@ -52,14 +52,22 @@ int vmm_set_page(void *P4, uintptr_t addr, uintptr_t page, uint16_t flags)
 
 int touch_page(void *P4, uintptr_t addr, uint16_t flags)
 {
-  (void)flags;
   if(!P4) return -1;
-  if(!P4e(P4, addr).present)
-    P4e(P4, addr).value = pmm_alloc() | flags | PAGE_PRESENT;
-  if(!P3e(P4, addr).present)
-    P3e(P4, addr).value = pmm_alloc() | flags | PAGE_PRESENT;
-  if(!P2e(P4, addr).present)
-    P2e(P4, addr).value = pmm_alloc() | flags | PAGE_PRESENT;
+
+  if((!P4e(P4, addr).present)
+      && (!(P4e(P4, addr).value = pmm_alloc())))
+    return -1;
+  P4e(P4, addr).value |= flags | PAGE_PRESENT;
+
+  if((!P3e(P4, addr).present)
+      && (!(P3e(P4, addr).value = pmm_alloc())))
+    return -1;
+  P3e(P4, addr).value |= flags | PAGE_PRESENT;
+
+  if((!P2e(P4, addr).present)
+      && (!(P2e(P4, addr).value = pmm_alloc())))
+    return -1;
+  P2e(P4, addr).value |= flags | PAGE_PRESENT;
 
   return 0;
 }
diff --git a/src/kernel/memory/vmm.tt b/src/kernel/memory/vmm.tt
index 79178c5..b9b7d57 100644
--- a/src/kernel/memory/vmm.tt
+++ b/src/kernel/memory/vmm.tt
@@ -119,6 +119,7 @@ uintptr_t pmm_alloc()
 {
   uintptr_t *pages[] = {p3, p2, p1};
   static int counter=0;
+  if(counter >= 3) return 0;
   return (uintptr_t)pages[counter++];
 }
 
@@ -146,3 +147,9 @@ TEST(touch_page_sets_flags)
 
   ASSERT_EQ_PTR(p2[0], (uintptr_t)p1 | 0x123 | PAGE_PRESENT);
 }
+TEST(touch_page_fails_if_out_of_pages)
+{
+  pmm_alloc();
+  int retval = touch_page(p4, 0, 0);
+  ASSERT_NEQ_INT(retval, 0);
+}