From 1d128f43ab18a790474711e5cbfd72fd9597c84c Mon Sep 17 00:00:00 2001
From: clement rouault <clement.rouault@exatrack.com>
Date: Tue, 15 Jun 2021 15:35:46 +0200
Subject: [PATCH] native_function.CustomAllocator now has the ability to be
 closed, freeing its pages.

---
 windows/native_exec/native_function.py | 68 +++++++-------------------
 1 file changed, 19 insertions(+), 49 deletions(-)

diff --git a/windows/native_exec/native_function.py b/windows/native_exec/native_function.py
index b0cf654..38892b0 100644
--- a/windows/native_exec/native_function.py
+++ b/windows/native_exec/native_function.py
@@ -5,59 +5,12 @@ import sys
 
 import windows
 import windows.winproxy
+import windows.generated_def as gdef
 
 from . import simple_x86 as x86
 from . import simple_x64 as x64
 
 
-class PyObj(ctypes.Structure):
-    _fields_ = [("ob_refcnt", ctypes.c_size_t),
-                ("ob_type", ctypes.c_void_p)]  # must be cast
-
-
-class PyMmap(PyObj):
-    _fields_ = [("ob_addr", ctypes.c_size_t), ("ob_size", ctypes.c_size_t)]
-
-
-# Specific mmap class for code injection
-class MyMap(mmap.mmap):
-    """ A mmap that is never unmapped and that contains the page address """
-    def __init__(self, *args, **kwarg):
-        # Get the page address by 'introspection' of the C struct
-        m = PyMmap.from_address(id(self))
-        self.addr = m.ob_addr
-        # Prevent garbage collection (so unmaping) of the page
-        m.ob_refcnt += 1
-
-    @classmethod
-    def get_map(cls, size):
-        """ Dispatch to the good mmap implem depending on the current system """
-        systems = {'windows': Win32MyMap,
-                   'linux': UnixMyMap}
-        x = platform.system().lower()
-        if x not in systems:
-            raise ValueError("Unknow system {0}".format(x))
-        return systems[x].get_map(size)
-
-
-class Win32MyMap(MyMap):
-    @classmethod
-    def get_map(cls, size):
-        addr = windows.winproxy.VirtualAlloc(0, size, 0x1000, 0x40)
-        new_map = (ctypes.c_char * size).from_address(addr)
-        new_map.addr = addr
-        if new_map.addr == 0:
-            raise ctypes.WinError()
-        return new_map
-
-
-class UnixMyMap(MyMap):
-    @classmethod
-    def get_map(cls, size):
-        prot = mmap.PROT_EXEC | mmap.PROT_WRITE | mmap.PROT_READ
-        return cls(-1, size, prot=prot)
-
-
 class CustomAllocator(object):
     int_size = {'32bit': 4, '64bit': 8}
 
@@ -75,7 +28,10 @@ class CustomAllocator(object):
         return cls.int_size[bits]
 
     def get_new_page(self, size):
-        self.maps.append(MyMap.get_map(size))
+        addr = windows.winproxy.VirtualAlloc(0, size, 0x1000, gdef.PAGE_EXECUTE_READWRITE)
+        mymap = (ctypes.c_char * size).from_address(addr)
+        mymap.addr = addr
+        self.maps.append(mymap)
         self.cur_offset = 0
         self.cur_page_size = size
 
@@ -99,6 +55,20 @@ class CustomAllocator(object):
         self.cur_offset += size
         return addr
 
+    def close(self):
+        maps = self.maps
+        self.maps = []
+        self.cur_offset = 0
+        self.cur_page_size = 0
+        if getattr(sys, "path", None) is None:
+            # Path is None -> Python shutdown
+            return
+        for mymap in maps:
+            windows.winproxy.VirtualFree(mymap.addr, dwFreeType=gdef.MEM_RELEASE)
+
+    def __del__(self):
+        self.close()
+
 allocator = CustomAllocator()