Merge pull request #80 from mrexodia/memory-partial-release

Memory partial release

.github/workflows/build.yml

@@ -70,7 +70,7 @@ jobs:
     - name: 'Test: DumpulatorTests'
       run: |
         cd tests
-        python harness-tests.py
+        python run-tests.py
 
     - name: 'Test: ExceptionTest_x64'
       run: |

src/dumpulator/memory.py

@@ -181,7 +181,7 @@ def find_free(self, size: int, allocation_align=True) -> Optional[int]:
                     return info.base
         return None
 
-    def reserve(self, start: int, size: int, protect: MemoryProtect, memory_type: MemoryType = MemoryType.MEM_PRIVATE, info: Any = None) -> None:
+    def reserve(self, start: int, size: int, protect: MemoryProtect, memory_type: MemoryType = MemoryType.MEM_PRIVATE, info: Any = None) -> MemoryRegion:
         assert isinstance(protect, MemoryProtect)
         assert isinstance(memory_type, MemoryType)
         assert size > 0 and self.align_page(size) == size
@@ -203,6 +203,7 @@ def check_overlaps(idx):
             check_overlaps(index - 1)
             check_overlaps(index)
         self._regions.insert(index, region)
+        return region
 
     def _decommit_region(self, parent_region: MemoryRegion, decommit_region: MemoryRegion):
         assert decommit_region in parent_region
@@ -227,19 +228,42 @@ def _decommit_region(self, parent_region: MemoryRegion, decommit_region: MemoryR
                 del self._committed[release_start + i * PAGE_SIZE]
                 parent_region.commit_count -= 1
 
-    def release(self, start: int) -> None:
+    def release(self, start: int, size: int = 0) -> None:
         assert self.align_allocation(start) == start
+        assert self.align_allocation(size) == size
 
-        parent_region = self.find_region(start)
-        if parent_region is None:
+        parent = self.find_region(start)
+        if parent is None:
             raise KeyError(f"Could not find parent for {hex(start)}")
-        if parent_region.start != start:
-            raise KeyError(f"You can only release the whole parent region")
 
-        self._decommit_region(parent_region, parent_region)
+        if size == 0:
+            size = parent.size
+        if start + size > parent.start + parent.size:
+            raise KeyError(f"You can only release part of the parent region")
+
+        decommit = MemoryRegion(start, size)
+        self._decommit_region(parent, decommit)
+        self._regions.remove(parent)
+
+        before = MemoryRegion(parent.start, decommit.start - parent.start)
+        after = MemoryRegion(decommit.end, parent.end - decommit.end)
+
+        if parent.size == decommit.size:
+            # Make sure the whole region was freed
+            assert parent.commit_count == 0
+
+        if before.size > 0:
+            before = self.reserve(before.start, before.size, parent.protect, parent.type, parent.info)
+            for page in before.pages():
+                if page in self._committed:
+                    before.commit_count += 1
+        if after.size > 0:
+            after = self.reserve(after.start, after.size, parent.protect, parent.type, parent.info)
+            for page in after.pages():
+                if page in self._committed:
+                    after.commit_count += 1
 
-        assert parent_region.commit_count == 0
-        self._regions.remove(parent_region)
+        assert before.commit_count + after.commit_count == parent.commit_count
 
     def commit(self, start: int, size: int, protect: MemoryProtect = MemoryProtect.UNDEFINED) -> None:
         assert isinstance(protect, MemoryProtect)

src/dumpulator/ntsyscalls.py

@@ -1820,11 +1820,10 @@ def ZwFreeVirtualMemory(dp: Dumpulator,
     size = RegionSize.read_ptr()
     if FreeType == MEM_RELEASE:
         print(f"release {hex(base)}[{hex(size)}]")
-        assert size == 0
         region = dp.memory.find_region(base)
         if region is None:
             return STATUS_MEMORY_NOT_ALLOCATED
-        dp.memory.release(base)
+        dp.memory.release(base, size)
         return STATUS_SUCCESS
     elif FreeType == MEM_DECOMMIT:
         print(f"decommit {hex(base)}[{hex(size)}]")

tests/DumpulatorTests/ExceptionTest/ExceptionTest.cpp

@@ -1,7 +1,5 @@
-#include <Windows.h>
-#include <cstdio>
-
 #include "../Tests/debug.h"
+#include <cstdlib>
 
 static LONG WINAPI VectoredHandler(struct _EXCEPTION_POINTERS* ExceptionInfo)
 {

tests/DumpulatorTests/Loader/Loader.cpp

@@ -1,5 +1,6 @@
-#include <Windows.h>
+#include "../Tests/debug.h"
 #include <cstdio>
+#include <cstdlib>
 
 int main(int argc, char** argv)
 {
@@ -9,6 +10,9 @@ int main(int argc, char** argv)
     auto dll = "Tests_x86.dll";
 #endif // _WIN64
     auto hLib = LoadLibraryA(dll);
+    auto p_DebugPrintf = (void**)GetProcAddress(hLib, "DebugPrintf");
+    if (p_DebugPrintf != nullptr)
+        *p_DebugPrintf = printf;
     if (argc < 2)
     {
         // TODO: implement enumerating all exports and running them

tests/DumpulatorTests/README.md

@@ -14,7 +14,7 @@ The harness dumps were created as follows:
 
 ## Adding a new test
 
-Add a new `mytest.cpp` file to the `Tests` project. The tests are exported as `bool <Prefix>_<description>Test();` and the result indicates whether the test was successful or not. If you need a custom environment add the following in `tests/harness-tests.py`:
+Add a new `mytest.cpp` file to the `Tests` project. The tests are exported as `bool <Prefix>_<description>Test();` and the result indicates whether the test was successful or not. If you need a custom environment add the following in `tests/run-tests.py`:
 
 ```python
 class <Prefix>Environment(TestEnvironment):

tests/DumpulatorTests/Tests/DllMain.cpp

@@ -1,4 +1,22 @@
-#include <Windows.h>
+#include "debug.h"
+
+inline ULONG
+STDAPIVCALLTYPE
+DbgPrintNop(
+	_In_z_ _Printf_format_string_ PCSTR Format,
+	...
+)
+{
+	return 0;
+}
+
+#ifdef _WIN64
+#pragma comment(linker, "/EXPORT:DebugPrintf=DebugPrintf")
+#else
+#pragma comment(linker, "/EXPORT:DebugPrintf=_DebugPrintf")
+#endif // _WIN64
+
+extern "C" decltype(&DbgPrint) DebugPrintf = DbgPrintNop;
 
 BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved)
 {

tests/DumpulatorTests/Tests/HandleTest.cpp

@@ -59,7 +59,7 @@ extern "C" __declspec(dllexport) bool Handle_ReadFileTest()
 	if (file_handle == INVALID_HANDLE_VALUE)
 	{
 		DebugPrint(L"Failed to open file");
-		return ret_value;
+		return false;
 	}
 
 	ret_value = ReadFile(
@@ -70,7 +70,8 @@ extern "C" __declspec(dllexport) bool Handle_ReadFileTest()
 		FALSE
 	);
 
+
 	CloseHandle(file_handle);
 
-	return ret_value;
+	return ret_value && memcmp(read_buffer, "Lorem ipsum", 11) == 0;
 }

tests/DumpulatorTests/Tests/MemoryTest.cpp

@@ -0,0 +1,43 @@
+#include "debug.h"
+
+extern "C" __declspec(dllexport) bool Memory_PartialReleaseTest()
+{
+	PVOID Base = 0;
+	SIZE_T RegionSize = 0x1c0000;
+	auto status = NtAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &RegionSize, MEM_RESERVE, PAGE_READWRITE);
+	DebugPrintf("status: 0x%08X, base: %p, size: %p\n", status, Base, RegionSize);
+	if (!NT_SUCCESS(status))
+		return false;
+
+	SIZE_T PartialSize = RegionSize - 0x10000;
+	status = NtFreeVirtualMemory(NtCurrentProcess(), &Base, &PartialSize, MEM_RELEASE);
+	DebugPrintf("status: 0x%08X, base: %p, size: %p\n", status, Base, PartialSize);
+	if (!NT_SUCCESS(status))
+		return false;
+
+	SIZE_T PageSize = 0x1000;
+	PVOID PartialBase = (char*)Base + PartialSize;
+	status = NtAllocateVirtualMemory(NtCurrentProcess(), &PartialBase, 0, &PageSize, MEM_COMMIT, PAGE_READWRITE);
+	DebugPrintf("status: 0x%08X, base: %p, size: %p\n", status, PartialBase, PageSize);
+	if (!NT_SUCCESS(status))
+		return false;
+	return true;
+}
+
+extern "C" __declspec(dllexport) bool Memory_MiddleReleaseTest()
+{
+	PVOID Base = 0;
+	SIZE_T RegionSize = 0x30000;
+	auto status = NtAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &RegionSize, MEM_RESERVE, PAGE_READWRITE);
+	DebugPrintf("status: 0x%08X, base: %p, size: %p\n", status, Base, RegionSize);
+	if (!NT_SUCCESS(status))
+		return false;
+
+	SIZE_T MiddleSize = 0x10000;
+	PVOID MiddleBase = (char*)Base + 0x10000;
+	status = NtFreeVirtualMemory(NtCurrentProcess(), &MiddleBase, &MiddleSize, MEM_RELEASE);
+	DebugPrintf("status: 0x%08X, base: %p, size: %p\n", status, MiddleBase, MiddleSize);
+	if (!NT_SUCCESS(status))
+		return false;
+	return true;
+}

tests/DumpulatorTests/Tests/Tests.vcxproj

@@ -13,9 +13,13 @@
   <ItemGroup>
     <ClCompile Include="DllMain.cpp" />
     <ClCompile Include="HandleTest.cpp" />
+    <ClCompile Include="MemoryTest.cpp" />
+    <ClCompile Include="ntstatusdb.cpp" />
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="debug.h" />
+    <ClInclude Include="ntstatusdb.h" />
+    <ClInclude Include="phnt.h" />
   </ItemGroup>
   <PropertyGroup Label="Globals">
     <VCProjectVersion>16.0</VCProjectVersion>

tests/DumpulatorTests/Tests/Tests.vcxproj.filters

@@ -21,10 +21,22 @@
     <ClCompile Include="DllMain.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="MemoryTest.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
+    <ClCompile Include="ntstatusdb.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="debug.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="phnt.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
+    <ClInclude Include="ntstatusdb.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
   </ItemGroup>
 </Project>

tests/DumpulatorTests/Tests/debug.h

@@ -1,22 +1,11 @@
 #pragma once
 
-#include <Windows.h>
-#include <winternl.h>
-
-#ifdef __cplusplus
-extern "C"
-#endif // __cplusplus
-NTSYSCALLAPI
-NTSTATUS
-NTAPI
-NtDisplayString(
-	PUNICODE_STRING String
-);
+#include "phnt.h"
 
 #define WIDEN_EXPAND(str) L ## str
 #define WIDEN(str) WIDEN_EXPAND(str)
 
-#ifdef __c1plusplus
+#ifdef __cplusplus
 // Helper function to directly call NtDisplayString with a string
 // This simplifies the trace output of Dumpulator
 template<size_t Count>
@@ -38,3 +27,5 @@ static void DebugPrint(const wchar_t* str)
 	NtDisplayString(&ustr);
 }
 #endif // __cplusplus
+
+extern "C" decltype(&DbgPrint) DebugPrintf;