From 00b9f016e3c29ea01df7f472eeeee2c1ed8e7300 Mon Sep 17 00:00:00 2001
From: DH <dh.rpcs3@gmail.com>
Date: Mon, 14 Oct 2024 19:15:05 +0300
Subject: [PATCH] orbis-kernel: implement debug kalloc

---
 orbis-kernel/src/KernelContext.cpp | 77 ++++++++++++++++++++++++------
 1 file changed, 62 insertions(+), 15 deletions(-)

diff --git a/orbis-kernel/src/KernelContext.cpp b/orbis-kernel/src/KernelContext.cpp
index 20ce6e5e4..5598bbfb6 100644
--- a/orbis-kernel/src/KernelContext.cpp
+++ b/orbis-kernel/src/KernelContext.cpp
@@ -2,15 +2,18 @@
 #include "orbis/thread/Process.hpp"
 #include "orbis/thread/ProcessOps.hpp"
 #include "orbis/utils/Logs.hpp"
+#include <bit>
 #include <chrono>
 #include <csignal>
+#include <mutex>
 #include <sys/mman.h>
 #include <thread>
 #include <unistd.h>
 
 static const std::uint64_t g_allocProtWord = 0xDEADBEAFBADCAFE1;
-static constexpr auto kHeapBaseAddress = 0x600'0000'0000;
-static constexpr auto kHeapSize = 0x2'0000'0000;
+static constexpr auto kHeapBaseAddress = 0x00000800'0000'0000;
+static constexpr auto kHeapSize = 0x10'0000'0000;
+static constexpr int kDebugHeap = 2;
 
 namespace orbis {
 thread_local Thread *g_currentThread;
@@ -53,7 +56,6 @@ Process *KernelContext::createProcess(pid_t pid) {
 
 void KernelContext::deleteProcess(Process *proc) {
   auto procNode = reinterpret_cast<utils::LinkedNode<Process> *>(proc);
-  auto pid = proc->pid;
 
   {
     std::lock_guard lock(m_proc_mtx);
@@ -154,11 +156,19 @@ void *KernelContext::kalloc(std::size_t size, std::size_t align) {
     // Try to reuse previously freed block
     for (auto [it, end] = m_free_heap.equal_range(size); it != end; ++it) {
       auto result = it->second;
-      if (!(reinterpret_cast<std::uintptr_t>(result) & (align - 1))) {
+      if (!(std::bit_cast<std::uintptr_t>(result) & (align - 1))) {
         auto node = m_free_heap.extract(it);
         node.key() = 0;
         node.mapped() = nullptr;
         m_used_node.insert(m_used_node.begin(), std::move(node));
+
+        // std::fprintf(stderr, "kalloc: reuse %p-%p, size = %lx\n", result,
+        //              (char *)result + size, size);
+
+        if (kDebugHeap > 0) {
+          std::memcpy(std::bit_cast<std::byte *>(result) + size,
+                      &g_allocProtWord, sizeof(g_allocProtWord));
+        }
         return result;
       }
     }
@@ -169,6 +179,13 @@ void *KernelContext::kalloc(std::size_t size, std::size_t align) {
   auto heap = reinterpret_cast<std::uintptr_t>(m_heap_next);
   heap = (heap + (align - 1)) & ~(align - 1);
 
+  if (kDebugHeap > 1) {
+    if (auto diff = (heap + size + sizeof(g_allocProtWord)) % 4096; diff != 0) {
+      heap += 4096 - diff;
+      heap &= ~(align - 1);
+    }
+  }
+
   if (heap + size > kHeapBaseAddress + kHeapSize) {
     std::fprintf(stderr, "out of kernel memory");
     std::abort();
@@ -179,18 +196,36 @@ void *KernelContext::kalloc(std::size_t size, std::size_t align) {
     std::abort();
   }
 
-  auto result = reinterpret_cast<void *>(heap);
-  std::memcpy(std::bit_cast<std::byte *>(result) + size, &g_allocProtWord,
-              sizeof(g_allocProtWord));
-  m_heap_next = reinterpret_cast<void *>(heap + size + sizeof(g_allocProtWord));
+  std::fprintf(stderr, "kalloc: allocate %lx-%lx, size = %lx, align=%lx\n",
+               heap, heap + size, size, align);
 
-  if (true) {
+  auto result = reinterpret_cast<void *>(heap);
+  if (kDebugHeap > 0) {
+    std::memcpy(std::bit_cast<std::byte *>(result) + size, &g_allocProtWord,
+                sizeof(g_allocProtWord));
+  }
+
+  if (kDebugHeap > 0) {
+    m_heap_next =
+        reinterpret_cast<void *>(heap + size + sizeof(g_allocProtWord));
+  } else {
+    m_heap_next = reinterpret_cast<void *>(heap + size);
+  }
+
+  if (kDebugHeap > 1) {
     heap = reinterpret_cast<std::uintptr_t>(m_heap_next);
     align = std::min<std::size_t>(align, 4096);
     heap = (heap + (align - 1)) & ~(align - 1);
     size = 4096;
-    ::mmap(reinterpret_cast<void *>(heap), size, PROT_NONE, MAP_FIXED, -1, 0);
+    std::fprintf(stderr, "kalloc: protect %lx-%lx, size = %lx, align=%lx\n",
+                 heap, heap + size, size, align);
 
+    auto result = ::mmap(reinterpret_cast<void *>(heap), size, PROT_NONE,
+                         MAP_FIXED | MAP_ANONYMOUS | MAP_SHARED, -1, 0);
+    if (result == MAP_FAILED) {
+      std::fprintf(stderr, "failed to protect memory");
+      std::abort();
+    }
     m_heap_next = reinterpret_cast<void *>(heap + size);
   }
 
@@ -203,14 +238,26 @@ void KernelContext::kfree(void *ptr, std::size_t size) {
   if (!size)
     std::abort();
 
-  std::memset(ptr, 0xcc, size);
-
-  if (std::memcmp(std::bit_cast<std::byte *>(ptr) + size, &g_allocProtWord,
-                  sizeof(g_allocProtWord)) != 0) {
-    std::fprintf(stderr, "kernel heap corruption\n");
+  if (std::bit_cast<std::uintptr_t>(ptr) < kHeapBaseAddress ||
+      std::bit_cast<std::uintptr_t>(ptr) + size >
+          kHeapBaseAddress + kHeapSize) {
+    std::fprintf(stderr, "kfree: invalid address");
     std::abort();
   }
 
+  if (kDebugHeap > 0) {
+    if (std::memcmp(std::bit_cast<std::byte *>(ptr) + size, &g_allocProtWord,
+                    sizeof(g_allocProtWord)) != 0) {
+      std::fprintf(stderr, "kernel heap corruption\n");
+      std::abort();
+    }
+
+    std::memset(ptr, 0xcc, size + sizeof(g_allocProtWord));
+  }
+
+  // std::fprintf(stderr, "kfree: release %p-%p, size = %lx\n", ptr,
+  //              (char *)ptr + size, size);
+
   std::lock_guard lock(m_heap_map_mtx);
   if (!m_used_node.empty()) {
     auto node = m_used_node.extract(m_used_node.begin());