refactor!: impl a allocator for allocations and unify api

Author: std-microblock

include/blook/allocator.h

@@ -0,0 +1,76 @@
+#pragma once
+
+#include "process.h"
+#include "protect.h"
+#include <cstdint>
+#include <expected>
+#include <map>
+#include <memory>
+#include <optional>
+#include <string>
+#include <vector>
+
+namespace blook {
+class Pointer;
+
+class ProcessAllocator {
+public:
+  explicit ProcessAllocator(std::shared_ptr<Process> proc);
+  ~ProcessAllocator();
+
+  CLASS_MOVE_ONLY(ProcessAllocator)
+
+  // Allocate memory with optional near address constraint
+  std::expected<Pointer, std::string>
+  try_allocate(size_t size, void *nearAddr = nullptr,
+               Protect protection = Protect::rw);
+
+  Pointer allocate(size_t size, void *nearAddr = nullptr,
+                   Protect protection = Protect::rw);
+
+  // Deallocate memory
+  std::expected<void, std::string> try_deallocate(Pointer ptr);
+  void deallocate(Pointer ptr);
+
+  // Get allocation statistics
+  size_t allocated_count() const;
+  size_t total_allocated_bytes() const;
+  size_t total_reserved_bytes() const;
+
+private:
+  struct AllocationInfo {
+    size_t size;
+    Protect protection;
+  };
+
+  struct PageInfo {
+    void *base_address;
+    size_t total_size;
+    size_t used_size;
+    Protect protection;
+    std::map<void *, AllocationInfo> allocations;
+  };
+
+  // Direct allocation using VirtualAllocEx
+  std::expected<void *, std::string>
+  allocate_direct(size_t size, Protect protection,
+                  void *nearAddr = nullptr);
+
+  // Find or create a page suitable for allocation
+  std::expected<void *, std::string>
+  find_or_allocate_page(size_t size, void *nearAddr,
+                        Protect protection);
+
+  // Allocate a new page
+  std::expected<void *, std::string>
+  allocate_new_page(size_t size, void *nearAddr,
+                    Protect protection);
+
+  std::shared_ptr<Process> proc;
+  std::map<void *, PageInfo> pages;
+
+  static constexpr size_t DEFAULT_PAGE_SIZE = 64 * 1024; // 64KB
+  static constexpr size_t LARGE_ALLOCATION_THRESHOLD = 32 * 1024; // 32KB
+};
+
+} // namespace blook

include/blook/function.h

@@ -123,9 +123,9 @@ class Function {
     }
 
     Serializer serializer;
-    void *pCodePage = Process::self()->malloc(utils::estimateCodeSize(program), Process::MemoryProtection::rwx);
+    auto pCodePage = Process::self()->malloc(utils::estimateCodeSize(program), Protect::rwx);
     if (auto err =
-            serializer.serialize(program, reinterpret_cast<int64_t>(pCodePage));
+            serializer.serialize(program, reinterpret_cast<int64_t>((void*)pCodePage));
         err != zasm::ErrorCode::None)
       throw std::runtime_error(
           std::format("JIT Serialization failure: {}", err.getErrorName()));

include/blook/memo.h

@@ -41,19 +41,17 @@ class Pointer {
   friend MemoryPatch;
 
 public:
-  using MemoryProtection = Process::MemoryProtection;
-
   std::shared_ptr<Process> proc = nullptr;
   bool is_self() const;
   bool is_valid() const;
 
   bool operator==(const Pointer &other) const = default;
 
-  Pointer malloc(size_t size, void *near,
-                 MemoryProtection protection = MemoryProtection::rw);
+  bool operator==(std::nullptr_t) const { return _offset == 0; }
+
+  Pointer malloc(size_t size, void *near, Protect protection = Protect::rw);
 
-  Pointer malloc(size_t size,
-                 MemoryProtection protection = MemoryProtection::rw);
+  Pointer malloc(size_t size, Protect protection = Protect::rw);
 
   Pointer malloc_rx_near_this(size_t size);
 
@@ -227,14 +225,16 @@ class Pointer {
 
   operator size_t() const { return (size_t)this->_offset; }
 
-  inline Pointer absolute(const auto &t) const { return {proc, (void *)t}; }
-
-  Function as_function();
+  operator void *() const { return (void *)this->_offset; }
 
-  [[nodiscard]] void *data() const;
+  [[nodiscard]] inline void *data() const { return (void *)_offset; }
 
   [[nodiscard]] inline size_t offset() const { return _offset; }
 
+  inline Pointer absolute(const auto &t) const { return {proc, (void *)t}; }
+
+  Function as_function();
+
   // Pointer operations
   inline Pointer add(const auto &t) const {
     return {proc, (void *)(_offset + (size_t)t)};
@@ -257,7 +257,7 @@ class Pointer {
   }
 
   [[nodiscard]] MemoryPatch
-      reassembly(std::function<void(zasm::x86::Assembler&)>);
+      reassembly(std::function<void(zasm::x86::Assembler &)>);
 
   [[nodiscard]] MemoryPatch reassembly_thread_pause();
 
@@ -266,10 +266,10 @@ class Pointer {
   // Original location gets a jump to the new memory
   [[nodiscard]] std::expected<MemoryPatch, std::string>
   try_reassembly_with_trampoline(
-      std::function<void(zasm::x86::Assembler&)> func);
+      std::function<void(zasm::x86::Assembler &)> func);
 
   [[nodiscard]] MemoryPatch
-  reassembly_with_trampoline(std::function<void(zasm::x86::Assembler&)> func);
+  reassembly_with_trampoline(std::function<void(zasm::x86::Assembler &)> func);
 
   std::optional<Function> guess_function(size_t max_scan_size = 50000);
 
@@ -296,7 +296,7 @@ class Pointer {
 struct ScopedSetMemoryRWX {
   Pointer ptr;
   size_t size;
-  Process::MemoryProtection old_protect;
+  Protect old_protect;
 
   ScopedSetMemoryRWX(Pointer ptr, size_t size);
   ScopedSetMemoryRWX(const ScopedSetMemoryRWX &) = delete;
@@ -528,10 +528,10 @@ class MemoryRange : public Pointer {
   // If the new code is smaller, pad with NOPs
   // If the new code is larger, return unexpected
   [[nodiscard]] std::expected<MemoryPatch, std::string>
-  try_reassembly_with_padding(std::function<void(zasm::x86::Assembler&)> func);
+  try_reassembly_with_padding(std::function<void(zasm::x86::Assembler &)> func);
 
   [[nodiscard]] MemoryPatch
-  reassembly_with_padding(std::function<void(zasm::x86::Assembler&)> func);
+  reassembly_with_padding(std::function<void(zasm::x86::Assembler &)> func);
 };
 
 static_assert(std::sentinel_for<decltype(std::declval<MemoryRange>().begin()),

include/blook/process.h

@@ -2,6 +2,7 @@
 
 #include "platform_types.h"
 #include "utils.h"
+#include "protect.h"
 #include <cstdint>
 #include <expected>
 #include <filesystem>
@@ -18,38 +19,10 @@ class Pointer;
 class Module;
 struct Thread;
 class Process;
-
-using Allocator = struct ProcessAllocator {
-  explicit ProcessAllocator(std::shared_ptr<Process> proc);
-
-  std::optional<Pointer> allocate(size_t size, void *nearAddr = nullptr);
-
-  void deallocate(Pointer addr);
-
-private:
-  struct AllocatedPageData {
-    size_t size;
-    std::map<void *, size_t> allocated;
-  };
-  std::map<void *, AllocatedPageData> allocatedPages;
-  std::shared_ptr<Process> proc;
-};
+class ProcessAllocator;
 
 class Process : public std::enable_shared_from_this<Process> {
 public:
-  enum class MemoryProtection {
-    None = 0,
-    Read = 0x0001,
-    Write = 0x0010,
-    Execute = 0x0100,
-    ReadWrite = Read | Write,
-    ReadWriteExecute = Read | Write | Execute,
-    ReadExecute = Read | Execute,
-    rw = ReadWrite,
-    rwx = ReadWriteExecute,
-    rx = ReadExecute
-  };
-
 #ifdef _WIN32
 
   explicit Process(HANDLE h);
@@ -64,9 +37,10 @@ class Process : public std::enable_shared_from_this<Process> {
   WIN_ONLY(HANDLE h = nullptr);
   friend Module;
   friend Pointer;
+  friend ProcessAllocator;
 
   std::shared_ptr<Pointer> _memo_ptr;
-  std::shared_ptr<struct ProcessAllocator> _allocator_ptr;
+  std::shared_ptr<ProcessAllocator> _allocator_ptr;
 
 public:
   CLASS_MOVE_ONLY(Process)
@@ -91,23 +65,23 @@ class Process : public std::enable_shared_from_this<Process> {
                                                       size_t size) const;
   bool check_writable(void *addr, size_t size) const;
 
-  std::expected<MemoryProtection, std::string>
+  std::expected<Protect, std::string>
   try_set_memory_protect(void *addr, size_t size,
-                         MemoryProtection protect) const;
-  MemoryProtection set_memory_protect(void *addr, size_t size,
-                                      MemoryProtection protect) const;
+                         Protect protect) const;
+  Protect set_memory_protect(void *addr, size_t size,
+                                      Protect protect) const;
 
   std::expected<bool, std::string> try_check_valid(void *addr) const;
   bool check_valid(void *addr) const;
 
-  std::expected<void *, std::string>
-  try_malloc(size_t size, MemoryProtection protect = MemoryProtection::rw,
-             void *nearAddr = nullptr) const;
-  void *malloc(size_t size, MemoryProtection protect = MemoryProtection::rw,
-               void *nearAddr = nullptr) const;
+  std::expected<Pointer, std::string>
+  try_malloc(size_t size, Protect protect = Protect::rw,
+             void *nearAddr = nullptr);
+  Pointer malloc(size_t size, Protect protect = Protect::rw,
+                 void *nearAddr = nullptr);
 
-  std::expected<void, std::string> try_free(void *addr, size_t size = 0) const;
-  void free(void *addr, size_t size = 0) const;
+  std::expected<void, std::string> try_free(void *addr, size_t size = 0);
+  void free(void *addr, size_t size = 0);
 
   [[nodiscard]] std::optional<std::shared_ptr<Module>>
   module(const std::string &name);
@@ -126,7 +100,7 @@ class Process : public std::enable_shared_from_this<Process> {
 
   Pointer memo();
 
-  struct ProcessAllocator& allocator();
+  ProcessAllocator& allocator();
 
   std::vector<Thread> threads();
 

include/blook/protect.h

@@ -0,0 +1,17 @@
+#pragma once
+
+namespace blook {
+enum class Protect {
+  None = 0,
+  Read = 0x0001,
+  Write = 0x0010,
+  Execute = 0x0100,
+  ReadWrite = Read | Write,
+  ReadWriteExecute = Read | Write | Execute,
+  ReadExecute = Read | Execute,
+  rw = ReadWrite,
+  rwx = ReadWriteExecute,
+  rx = ReadExecute
+};
+
+}

src/function.cpp

@@ -54,7 +54,7 @@ void *Function::into_safe_function_pointer(void *func, bool thread_safety) {
   auto program = Program(utils::compileMachineMode());
   auto a = x86::Assembler(program);
 
-  auto ptr = Pointer(Process::self(), Process::self()->malloc(300, Process::MemoryProtection::rwx, func));
+  auto ptr = Process::self()->malloc(300, Protect::rwx, func);
   ptr.reassembly([=](auto a) {
 #ifdef BLOOK_ARCHITECTURE_X86_64
        auto sp = x86::rsp;

src/hook.cpp

@@ -102,7 +102,7 @@ Trampoline Trampoline::make(Pointer pCode, size_t minByteSize,
   zasm::Serializer serializer;
   auto pCodePage = pCode.malloc(utils::estimateCodeSize(program),
                                 near_alloc ? pCode.data() : nullptr,
-                                Pointer::MemoryProtection::rwx);
+                                Protect::rwx);
   if (pCodePage == nullptr) {
     throw std::runtime_error("Failed to allocate memory for trampoline.");
   }

src/memory.cpp

@@ -83,8 +83,6 @@ bool MemoryPatch::restore() {
   return true;
 }
 
-void *Pointer::data() const { return (void *)_offset; }
-
 bool Pointer::is_self() const { return proc && proc->is_self(); }
 
 std::optional<Function> Pointer::guess_function(size_t max_scan_size) {
@@ -126,7 +124,7 @@ Pointer::find_upwards(std::initializer_list<uint8_t> pattern,
   Pointer p = *this;
 
   if (p.proc->is_self())
-    for (; p > this - max_scan_size; p -= 1) {
+    for (; p > (*this - max_scan_size); p -= 1) {
       if (memcmp(pattern.begin(), (unsigned char *)p.data(), pattern.size()) ==
           0) {
         return p;

src/platform/linux/memory.cpp

@@ -32,22 +32,22 @@ void Pointer::protect_rwx(void *p, size_t size) {
 
 void *Pointer::malloc_near_rwx(void *targetAddr, size_t size) {
     return Process::self()->memo().malloc(size, targetAddr,
-                                          MemoryProtection::rwx);
+                                          Protect::rwx);
 }
 
-void *Pointer::malloc(size_t size, Pointer::MemoryProtection protection) {
+void *Pointer::malloc(size_t size, Pointer::Protect protection) {
     int prot_flags = 0;
     switch (protection) {
-        case MemoryProtection::Read:
+        case Protect::Read:
             prot_flags = PROT_READ;
             break;
-        case MemoryProtection::ReadWrite:
+        case Protect::ReadWrite:
             prot_flags = PROT_READ | PROT_WRITE;
             break;
-        case MemoryProtection::ReadWriteExecute:
+        case Protect::ReadWriteExecute:
             prot_flags = PROT_READ | PROT_WRITE | PROT_EXEC;
             break;
-        case MemoryProtection::ReadExecute:
+        case Protect::ReadExecute:
             prot_flags = PROT_READ | PROT_EXEC;
             break;
         default:
@@ -62,20 +62,20 @@ void *Pointer::malloc(size_t size, Pointer::MemoryProtection protection) {
 Pointer::Pointer(std::shared_ptr<Process> proc) : proc(std::move(proc)) {}
 
 void *Pointer::malloc(size_t size, void *nearby,
-                      Pointer::MemoryProtection protection) {
+                      Pointer::Protect protection) {
     long pagesize = sysconf(_SC_PAGESIZE);
     int prot_flags = 0;
     switch (protection) {
-        case MemoryProtection::Read:
+        case Protect::Read:
             prot_flags = PROT_READ;
             break;
-        case MemoryProtection::ReadWrite:
+        case Protect::ReadWrite:
             prot_flags = PROT_READ | PROT_WRITE;
             break;
-        case MemoryProtection::ReadWriteExecute:
+        case Protect::ReadWriteExecute:
             prot_flags = PROT_READ | PROT_WRITE | PROT_EXEC;
             break;
-        case MemoryProtection::ReadExecute:
+        case Protect::ReadExecute:
             prot_flags = PROT_READ | PROT_EXEC;
             break;
         default: