From 5fcfe6c516c3c12269e8ac17f8f698d098692beb Mon Sep 17 00:00:00 2001
From: OEOTYAN <OEOTYAN@outlook.com>
Date: Wed, 20 Mar 2024 22:39:55 +0800
Subject: [PATCH] refactor: make NativeClosure safer

---
 src/ll/api/memory/Closure.cpp | 23 +++++++++++++----------
 src/ll/api/memory/Closure.h   |  6 +++---
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/src/ll/api/memory/Closure.cpp b/src/ll/api/memory/Closure.cpp
index cf0e8a5867..7bbfdfef91 100644
--- a/src/ll/api/memory/Closure.cpp
+++ b/src/ll/api/memory/Closure.cpp
@@ -26,7 +26,12 @@ using T = NativeClosure<void*>;
 void initNativeClosure(void* t, void* impl, size_t offset, size_t size) {
     impl                       = unwrapFuncPtrJmp(impl);
     auto                  self = (T*)t;
-    NativeClosurePrologue asmc = {
+
+    self->closure = VirtualAlloc(nullptr, size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
+
+    memcpy(self->closure, impl, offset);
+
+    new ((char*)(self->closure) + offset) NativeClosurePrologue{
         .data     = (uintptr_t)&self->stored,
         .push_rax = 0x50,
         .mov_rax  = {0x48, 0xB8},
@@ -34,16 +39,14 @@ void initNativeClosure(void* t, void* impl, size_t offset, size_t size) {
         .jmp_rax  = {0xFF, 0xE0}
     };
 
-    self->closure = std::make_unique_for_overwrite<uchar[]>(size);
-
-    VirtualProtect(self->closure.get(), size, PAGE_EXECUTE_READWRITE, &self->oldProtectFlags);
-
-    memcpy(self->closure.get(), impl, offset);
-    memcpy(self->closure.get() + offset, &asmc, sizeof(asmc));
+    ulong _;
+    VirtualProtect(self->closure, size, PAGE_EXECUTE_READ, &_);
 }
 void releaseNativeClosure(void* t, size_t size) {
-    auto  self = (T*)t;
-    ulong _;
-    VirtualProtect(self->closure.get(), size, self->oldProtectFlags, &_);
+    auto self = (T*)t;
+    if (self->closure != nullptr) {
+        VirtualFree(self->closure, 0, MEM_RELEASE);
+        self->closure = nullptr;
+    }
 }
 } // namespace ll::memory::detail
diff --git a/src/ll/api/memory/Closure.h b/src/ll/api/memory/Closure.h
index 56b048805f..599a2a4123 100644
--- a/src/ll/api/memory/Closure.h
+++ b/src/ll/api/memory/Closure.h
@@ -72,14 +72,14 @@ class NativeClosure {
         origin_fn* func;
         uintptr_t  data;
     } stored;
-    ulong                    oldProtectFlags{};
-    std::unique_ptr<uchar[]> closure;
+
+    void* closure;
 
     NativeClosure(origin_fn* func, uintptr_t data) : stored({func, data}) {
         detail::initNativeClosure(this, closureImpl, implOffset, closureSize);
     }
 
-    closure_fn* get() const { return (closure_fn*)closure.get(); }
+    closure_fn* get() const { return (closure_fn*)closure; }
 
     ~NativeClosure() { detail::releaseNativeClosure(this, closureSize); }