Skip to content

Commit 9fa2c3f

Browse files
mrexodiamrexodia
committed
Allow releasing partial memory ranges
1 parent ef27fb8 commit 9fa2c3f

File tree

13 files changed

+41432
-30
lines changed

13 files changed

+41432
-30
lines changed

‎src/dumpulator/memory.py‎

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,7 @@ def find_free(self, size: int, allocation_align=True) -> Optional[int]:
181181
return info.base
182182
return None
183183

184-
def reserve(self, start: int, size: int, protect: MemoryProtect, memory_type: MemoryType = MemoryType.MEM_PRIVATE, info: Any = None) -> None:
184+
def reserve(self, start: int, size: int, protect: MemoryProtect, memory_type: MemoryType = MemoryType.MEM_PRIVATE, info: Any = None) -> MemoryRegion:
185185
assert isinstance(protect, MemoryProtect)
186186
assert isinstance(memory_type, MemoryType)
187187
assert size > 0 and self.align_page(size) == size
@@ -203,6 +203,7 @@ def check_overlaps(idx):
203203
check_overlaps(index - 1)
204204
check_overlaps(index)
205205
self._regions.insert(index, region)
206+
return region
206207

207208
def _decommit_region(self, parent_region: MemoryRegion, decommit_region: MemoryRegion):
208209
assert decommit_region in parent_region
@@ -227,19 +228,42 @@ def _decommit_region(self, parent_region: MemoryRegion, decommit_region: MemoryR
227228
del self._committed[release_start + i * PAGE_SIZE]
228229
parent_region.commit_count -= 1
229230

230-
def release(self, start: int) -> None:
231+
def release(self, start: int, size: int = 0) -> None:
231232
assert self.align_allocation(start) == start
233+
assert self.align_allocation(size) == size
232234

233-
parent_region = self.find_region(start)
234-
if parent_region is None:
235+
parent = self.find_region(start)
236+
if parent is None:
235237
raise KeyError(f"Could not find parent for {hex(start)}")
236-
if parent_region.start != start:
237-
raise KeyError(f"You can only release the whole parent region")
238238

239-
self._decommit_region(parent_region, parent_region)
239+
if size == 0:
240+
size = parent.size
241+
if start + size > parent.start + parent.size:
242+
raise KeyError(f"You can only release part of the parent region")
243+
244+
decommit = MemoryRegion(start, size)
245+
self._decommit_region(parent, decommit)
246+
self._regions.remove(parent)
247+
248+
before = MemoryRegion(parent.start, decommit.start - parent.start)
249+
after = MemoryRegion(decommit.end, parent.end - decommit.end)
250+
251+
if parent.size == decommit.size:
252+
# Make sure the whole region was freed
253+
assert parent.commit_count == 0
254+
255+
if before.size > 0:
256+
before = self.reserve(before.start, before.size, parent.protect, parent.type, parent.info)
257+
for page in before.pages():
258+
if page in self._committed:
259+
before.commit_count += 1
260+
if after.size > 0:
261+
after = self.reserve(after.start, after.size, parent.protect, parent.type, parent.info)
262+
for page in after.pages():
263+
if page in self._committed:
264+
after.commit_count += 1
240265

241-
assert parent_region.commit_count == 0
242-
self._regions.remove(parent_region)
266+
assert before.commit_count + after.commit_count == parent.commit_count
243267

244268
def commit(self, start: int, size: int, protect: MemoryProtect = MemoryProtect.UNDEFINED) -> None:
245269
assert isinstance(protect, MemoryProtect)

‎src/dumpulator/ntsyscalls.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1820,11 +1820,10 @@ def ZwFreeVirtualMemory(dp: Dumpulator,
18201820
size = RegionSize.read_ptr()
18211821
if FreeType == MEM_RELEASE:
18221822
print(f"release {hex(base)}[{hex(size)}]")
1823-
assert size == 0
18241823
region = dp.memory.find_region(base)
18251824
if region is None:
18261825
return STATUS_MEMORY_NOT_ALLOCATED
1827-
dp.memory.release(base)
1826+
dp.memory.release(base, size)
18281827
return STATUS_SUCCESS
18291828
elif FreeType == MEM_DECOMMIT:
18301829
print(f"decommit {hex(base)}[{hex(size)}]")

‎tests/DumpulatorTests/ExceptionTest/ExceptionTest.cpp‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
1-
#include <Windows.h>
2-
#include <cstdio>
3-
41
#include "../Tests/debug.h"
2+
#include <cstdlib>
53

64
static LONG WINAPI VectoredHandler(struct _EXCEPTION_POINTERS* ExceptionInfo)
75
{

‎tests/DumpulatorTests/Loader/Loader.cpp‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
1-
#include <Windows.h>
1+
#include "../Tests/debug.h"
22
#include <cstdio>
3+
#include <cstdlib>
34

45
int main(int argc, char** argv)
56
{
@@ -9,6 +10,9 @@ int main(int argc, char** argv)
910
auto dll = "Tests_x86.dll";
1011
#endif // _WIN64
1112
auto hLib = LoadLibraryA(dll);
13+
auto p_DebugPrintf = (void**)GetProcAddress(hLib, "DebugPrintf");
14+
if (p_DebugPrintf != nullptr)
15+
*p_DebugPrintf = printf;
1216
if (argc < 2)
1317
{
1418
// TODO: implement enumerating all exports and running them

‎tests/DumpulatorTests/Tests/DllMain.cpp‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,22 @@
1-
#include <Windows.h>
1+
#include "debug.h"
2+
3+
inline ULONG
4+
STDAPIVCALLTYPE
5+
DbgPrintNop(
6+
_In_z_ _Printf_format_string_ PCSTR Format,
7+
...
8+
)
9+
{
10+
return 0;
11+
}
12+
13+
#ifdef _WIN64
14+
#pragma comment(linker, "/EXPORT:DebugPrintf=DebugPrintf")
15+
#else
16+
#pragma comment(linker, "/EXPORT:DebugPrintf=_DebugPrintf")
17+
#endif // _WIN64
18+
19+
extern "C" decltype(&DbgPrint) DebugPrintf = DbgPrintNop;
220

321
BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved)
422
{

‎tests/DumpulatorTests/Tests/HandleTest.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ extern "C" __declspec(dllexport) bool Handle_WriteAndCreateFileTest()
1818
FILE_ATTRIBUTE_NORMAL,
1919
NULL
2020
);
21-
21+
2222
if (file_handle == INVALID_HANDLE_VALUE)
2323
{
2424
DebugPrint(L"Failed to create file");

‎tests/DumpulatorTests/Tests/MemoryTest.cpp‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
#include "debug.h"
2+
3+
extern "C" __declspec(dllexport) bool Memory_PartialReleaseTest()
4+
{
5+
PVOID Base = 0;
6+
SIZE_T RegionSize = 0x1c0000;
7+
auto status = NtAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &RegionSize, MEM_RESERVE, PAGE_READWRITE);
8+
DebugPrintf("status: 0x%08X, base: %p, size: %p\n", status, Base, RegionSize);
9+
if (!NT_SUCCESS(status))
10+
return false;
11+
12+
SIZE_T PartialSize = RegionSize - 0x10000;
13+
status = NtFreeVirtualMemory(NtCurrentProcess(), &Base, &PartialSize, MEM_RELEASE);
14+
DebugPrintf("status: 0x%08X, base: %p, size: %p\n", status, Base, PartialSize);
15+
if (!NT_SUCCESS(status))
16+
return false;
17+
18+
SIZE_T PageSize = 0x1000;
19+
PVOID PartialBase = (char*)Base + PartialSize;
20+
status = NtAllocateVirtualMemory(NtCurrentProcess(), &PartialBase, 0, &PageSize, MEM_COMMIT, PAGE_READWRITE);
21+
DebugPrintf("status: 0x%08X, base: %p, size: %p\n", status, PartialBase, PageSize);
22+
if (!NT_SUCCESS(status))
23+
return false;
24+
return true;
25+
}
26+
27+
extern "C" __declspec(dllexport) bool Memory_MiddleReleaseTest()
28+
{
29+
PVOID Base = 0;
30+
SIZE_T RegionSize = 0x30000;
31+
auto status = NtAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &RegionSize, MEM_RESERVE, PAGE_READWRITE);
32+
DebugPrintf("status: 0x%08X, base: %p, size: %p\n", status, Base, RegionSize);
33+
if (!NT_SUCCESS(status))
34+
return false;
35+
36+
SIZE_T MiddleSize = 0x10000;
37+
PVOID MiddleBase = (char*)Base + 0x10000;
38+
status = NtFreeVirtualMemory(NtCurrentProcess(), &MiddleBase, &MiddleSize, MEM_RELEASE);
39+
DebugPrintf("status: 0x%08X, base: %p, size: %p\n", status, MiddleBase, MiddleSize);
40+
if (!NT_SUCCESS(status))
41+
return false;
42+
return true;
43+
}

‎tests/DumpulatorTests/Tests/Tests.vcxproj‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,13 @@
1313
<ItemGroup>
1414
<ClCompile Include="DllMain.cpp" />
1515
<ClCompile Include="HandleTest.cpp" />
16+
<ClCompile Include="MemoryTest.cpp" />
17+
<ClCompile Include="ntstatusdb.cpp" />
1618
</ItemGroup>
1719
<ItemGroup>
1820
<ClInclude Include="debug.h" />
21+
<ClInclude Include="ntstatusdb.h" />
22+
<ClInclude Include="phnt.h" />
1923
</ItemGroup>
2024
<PropertyGroup Label="Globals">
2125
<VCProjectVersion>16.0</VCProjectVersion>

‎tests/DumpulatorTests/Tests/Tests.vcxproj.filters‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,22 @@
2121
<ClCompile Include="DllMain.cpp">
2222
<Filter>Source Files</Filter>
2323
</ClCompile>
24+
<ClCompile Include="MemoryTest.cpp">
25+
<Filter>Source Files</Filter>
26+
</ClCompile>
27+
<ClCompile Include="ntstatusdb.cpp">
28+
<Filter>Source Files</Filter>
29+
</ClCompile>
2430
</ItemGroup>
2531
<ItemGroup>
2632
<ClInclude Include="debug.h">
2733
<Filter>Header Files</Filter>
2834
</ClInclude>
35+
<ClInclude Include="phnt.h">
36+
<Filter>Header Files</Filter>
37+
</ClInclude>
38+
<ClInclude Include="ntstatusdb.h">
39+
<Filter>Header Files</Filter>
40+
</ClInclude>
2941
</ItemGroup>
3042
</Project>

‎tests/DumpulatorTests/Tests/debug.h‎

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,11 @@
11
#pragma once
22

3-
#include <Windows.h>
4-
#include <winternl.h>
5-
6-
#ifdef __cplusplus
7-
extern "C"
8-
#endif // __cplusplus
9-
NTSYSCALLAPI
10-
NTSTATUS
11-
NTAPI
12-
NtDisplayString(
13-
PUNICODE_STRING String
14-
);
3+
#include "phnt.h"
154

165
#define WIDEN_EXPAND(str) L ## str
176
#define WIDEN(str) WIDEN_EXPAND(str)
187

19-
#ifdef __c1plusplus
8+
#ifdef __cplusplus
209
// Helper function to directly call NtDisplayString with a string
2110
// This simplifies the trace output of Dumpulator
2211
template<size_t Count>
@@ -38,3 +27,5 @@ static void DebugPrint(const wchar_t* str)
3827
NtDisplayString(&ustr);
3928
}
4029
#endif // __cplusplus
30+
31+
extern "C" decltype(&DbgPrint) DebugPrintf;

0 commit comments

Comments
 (0)