refactor!: memory related apis

Author: std-microblock

include/blook/function.h

@@ -1,5 +1,6 @@
 #pragma once
 #include "blook/memo.h"
+#include "blook/process.h"
 #include "blook/utils.h"
 
 #include <cstddef>
@@ -122,7 +123,7 @@ class Function {
     }
 
     Serializer serializer;
-    void *pCodePage = Pointer::malloc_rwx(utils::estimateCodeSize(program));
+    void *pCodePage = Process::self()->malloc(utils::estimateCodeSize(program), Process::MemoryProtection::rwx);
     if (auto err =
             serializer.serialize(program, reinterpret_cast<int64_t>(pCodePage));
         err != zasm::ErrorCode::None)

include/blook/memo.h

@@ -17,6 +17,8 @@
 #include <vector>
 
 
+#include "process.h"
+
 namespace blook {
 class Process;
 
@@ -40,34 +42,23 @@ class Pointer {
   friend MemoryPatch;
 
 public:
+  using MemoryProtection = Process::MemoryProtection;
+
   std::shared_ptr<Process> proc = nullptr;
   bool is_self() const;
-  static void *malloc_rwx(size_t size);
-
-  static void protect_rwx(void *p, size_t size);
-
-  static void *malloc_near_rwx(void *near, size_t size);
 
   bool operator==(const Pointer &other) const = default;
 
-  enum class MemoryProtection {
-    Read = 0x0001,
-    Write = 0x0010,
-    Execute = 0x0100,
-    ReadWrite = Read | Write,
-    ReadWriteExecute = Read | Write | Execute,
-    ReadExecute = Read | Execute,
-    rw = ReadWrite,
-    rwx = ReadWriteExecute,
-    rx = ReadExecute
-  };
-
   Pointer malloc(size_t size, void *near,
                  MemoryProtection protection = MemoryProtection::rw);
 
   Pointer malloc(size_t size,
                  MemoryProtection protection = MemoryProtection::rw);
 
+  Pointer malloc_rx_near_this(size_t size);
+
+  void free(size_t size = 0);
+
   std::optional<Thread> create_thread(bool suspended = false);
 
   // Read operations
@@ -287,13 +278,13 @@ class Pointer {
 };
 
 struct ScopedSetMemoryRWX {
-  void *ptr;
+  Pointer ptr;
   size_t size;
-  void *old_protect;
+  Process::MemoryProtection old_protect;
 
-  ScopedSetMemoryRWX(void *ptr, size_t size);
-  ScopedSetMemoryRWX(const Pointer &p, size_t size)
-      : ScopedSetMemoryRWX(p.data(), size) {}
+  ScopedSetMemoryRWX(Pointer ptr, size_t size);
+  ScopedSetMemoryRWX(void *ptr, size_t size)
+      : ScopedSetMemoryRWX(Pointer(ptr), size) {}
   ScopedSetMemoryRWX(const ScopedSetMemoryRWX &) = delete;
   ScopedSetMemoryRWX &operator=(const ScopedSetMemoryRWX &) = delete;
   ScopedSetMemoryRWX(const MemoryRange &r);

include/blook/module.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "process.h"
 #include <filesystem>
 #include <map>
 #include <optional>

include/blook/process.h

@@ -1,50 +1,62 @@
 #pragma once
 
-#include "module.h"
-
 #include "platform_types.h"
-
 #include "utils.h"
 #include <cstdint>
 #include <expected>
 #include <filesystem>
 #include <iostream>
+#include <map>
+#include <memory>
 #include <optional>
 #include <span>
 #include <string>
 #include <vector>
 
 namespace blook {
 class Pointer;
+class Module;
 struct Thread;
-class Process : public std::enable_shared_from_this<Process> {
-#ifdef _WIN32
+class Process;
 
+using Allocator = struct ProcessAllocator {
+  explicit ProcessAllocator(std::shared_ptr<Process> proc);
 
-  explicit Process(HANDLE h);
+  std::optional<Pointer> allocate(size_t size, void *nearAddr = nullptr);
 
-  explicit Process(DWORD pid);
+  void deallocate(Pointer addr);
 
-#endif
+private:
+  struct AllocatedPageData {
+    size_t size;
+    std::map<void *, size_t> allocated;
+  };
+  std::map<void *, AllocatedPageData> allocatedPages;
+  std::shared_ptr<Process> proc;
+};
 
-  struct Allocator {
-    explicit Allocator(std::shared_ptr<Process> proc);
+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
+  };
 
-    std::optional<Pointer> allocate(size_t size, void *nearAddr = nullptr);
+#ifdef _WIN32
 
-    void deallocate(Pointer addr);
+  explicit Process(HANDLE h);
 
-  private:
-    struct AllocatedPageData {
-      size_t size;
-      std::map<void *, size_t> allocated;
-    };
-    std::map<void *, AllocatedPageData> allocatedPages;
-    std::shared_ptr<Process> proc;
-  };
+  explicit Process(DWORD pid);
 
-  std::optional<Pointer> _memo{};
-  std::optional<Allocator> _allocator{};
+#endif
 
   explicit Process(std::string name, size_t skip = 0);
   bool is_self_cached;
@@ -53,17 +65,49 @@ class Process : public std::enable_shared_from_this<Process> {
   friend Module;
   friend Pointer;
 
+  std::shared_ptr<Pointer> _memo_ptr;
+  std::shared_ptr<struct ProcessAllocator> _allocator_ptr;
+
 public:
   CLASS_MOVE_ONLY(Process)
 #ifdef _WIN32
   DWORD pid;
 #endif
-  [[nodiscard]] std::optional<std::vector<std::uint8_t>>
-  read(void *addr, size_t size) const;
 
-  void *read(void *dest, void *addr, size_t size) const;
+  // Memory APIs
+  std::expected<void, std::string> try_read(void *dst, void *src,
+                                            size_t size) const;
+  void read(void *dst, void *src, size_t size) const;
+
+  std::expected<void, std::string> try_write(void *dst, const void *src,
+                                             size_t size) const;
+  void write(void *dst, const void *src, size_t size) const;
+
+  std::expected<bool, std::string> try_check_readable(void *addr,
+                                                      size_t size) const;
+  bool check_readable(void *addr, size_t size) const;
+
+  std::expected<bool, std::string> try_check_writable(void *addr,
+                                                      size_t size) const;
+  bool check_writable(void *addr, size_t size) const;
+
+  std::expected<MemoryProtection, 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;
+
+  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<void, std::string> write(void *addr, std::span<uint8_t>) const;
+  std::expected<void, std::string> try_free(void *addr, size_t size = 0) const;
+  void free(void *addr, size_t size = 0) const;
 
   [[nodiscard]] std::optional<std::shared_ptr<Module>>
   module(const std::string &name);
@@ -82,7 +126,7 @@ class Process : public std::enable_shared_from_this<Process> {
 
   Pointer memo();
 
-  Allocator allocator();
+  struct ProcessAllocator& allocator();
 
   std::vector<Thread> threads();
 

include/blook/utils.h

@@ -1,15 +1,17 @@
 #pragma once
+
 #include "ctype.h"
 #include "zasm/base/mode.hpp"
 #include "zasm/decoder/decoder.hpp"
 #include "zasm/zasm.hpp"
 
+#ifndef CLASS_MOVE_ONLY
 #define CLASS_MOVE_ONLY(classname)                                             \
-  classname() = delete;                                                        \
-  classname(classname &) = delete;                                             \
-  classname &operator=(classname &) = delete;                                  \
+  classname(const classname &) = delete;                                       \
+  classname &operator=(const classname &) = delete;                            \
   classname(classname &&) noexcept = default;                                  \
   classname &operator=(classname &&) noexcept = default;
+#endif
 
 
 #if defined(__x86_64__) || defined(_M_X64)

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)Pointer::malloc_near_rwx(func, 300));
+  auto ptr = Pointer(Process::self(), Process::self()->malloc(300, Process::MemoryProtection::rwx, func));
   ptr.reassembly([=](auto a) {
 #ifdef BLOOK_ARCHITECTURE_X86_64
        auto sp = x86::rsp;

src/memory.cpp

@@ -46,16 +46,16 @@ MemoryPatch::MemoryPatch(Pointer ptr, std::vector<uint8_t> buffer)
 
 void MemoryPatch::swap() {
   const auto target = reinterpret_cast<void *>(ptr._offset);
+  ScopedSetMemoryRWX protector(ptr, buffer.size());
   if (ptr.proc->is_self()) {
-    ScopedSetMemoryRWX protector(ptr, buffer.size());
     std::vector<uint8_t> tmp(buffer.size());
     std::memcpy(tmp.data(), buffer.data(), buffer.size());
     std::memcpy(buffer.data(), target, buffer.size());
     std::memcpy(target, tmp.data(), tmp.size());
   } else {
     std::vector<uint8_t> tmp(buffer.size());
     ptr.proc->read(tmp.data(), target, buffer.size());
-    if (!ptr.proc->write(target, buffer))
+    if (!ptr.proc->try_write(target, buffer.data(), buffer.size()))
       throw std::runtime_error("Failed to write memory");
     buffer = std::move(tmp);
   }