Skip to content

Commit 06968d8

Browse files
bb107bb107
committed
#49 #52 add initialize/cleanup to achieve graceful shutdown
1 parent 61a80f4 commit 06968d8

File tree

15 files changed

+176
-87
lines changed

15 files changed

+176
-87
lines changed

‎3rdparty/phnt/include/ntioapi.h‎

Lines changed: 0 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -774,22 +774,6 @@ typedef struct _FILE_ID_EXTD_BOTH_DIR_INFORMATION
774774
WCHAR FileName[1];
775775
} FILE_ID_EXTD_BOTH_DIR_INFORMATION, *PFILE_ID_EXTD_BOTH_DIR_INFORMATION;
776776

777-
// private
778-
typedef struct _FILE_STAT_INFORMATION
779-
{
780-
LARGE_INTEGER FileId;
781-
LARGE_INTEGER CreationTime;
782-
LARGE_INTEGER LastAccessTime;
783-
LARGE_INTEGER LastWriteTime;
784-
LARGE_INTEGER ChangeTime;
785-
LARGE_INTEGER AllocationSize;
786-
LARGE_INTEGER EndOfFile;
787-
ULONG FileAttributes;
788-
ULONG ReparseTag;
789-
ULONG NumberOfLinks;
790-
ACCESS_MASK EffectiveAccess;
791-
} FILE_STAT_INFORMATION, *PFILE_STAT_INFORMATION;
792-
793777
// private
794778
typedef struct _FILE_MEMORY_PARTITION_INFORMATION
795779
{
@@ -805,43 +789,6 @@ typedef struct _FILE_MEMORY_PARTITION_INFORMATION
805789
} Flags;
806790
} FILE_MEMORY_PARTITION_INFORMATION, *PFILE_MEMORY_PARTITION_INFORMATION;
807791

808-
// LxFlags
809-
#define LX_FILE_METADATA_HAS_UID 0x1
810-
#define LX_FILE_METADATA_HAS_GID 0x2
811-
#define LX_FILE_METADATA_HAS_MODE 0x4
812-
#define LX_FILE_METADATA_HAS_DEVICE_ID 0x8
813-
#define LX_FILE_CASE_SENSITIVE_DIR 0x10
814-
815-
// private
816-
typedef struct _FILE_STAT_LX_INFORMATION
817-
{
818-
LARGE_INTEGER FileId;
819-
LARGE_INTEGER CreationTime;
820-
LARGE_INTEGER LastAccessTime;
821-
LARGE_INTEGER LastWriteTime;
822-
LARGE_INTEGER ChangeTime;
823-
LARGE_INTEGER AllocationSize;
824-
LARGE_INTEGER EndOfFile;
825-
ULONG FileAttributes;
826-
ULONG ReparseTag;
827-
ULONG NumberOfLinks;
828-
ACCESS_MASK EffectiveAccess;
829-
ULONG LxFlags;
830-
ULONG LxUid;
831-
ULONG LxGid;
832-
ULONG LxMode;
833-
ULONG LxDeviceIdMajor;
834-
ULONG LxDeviceIdMinor;
835-
} FILE_STAT_LX_INFORMATION, *PFILE_STAT_LX_INFORMATION;
836-
837-
#define FILE_CS_FLAG_CASE_SENSITIVE_DIR 0x00000001
838-
839-
// private
840-
typedef struct _FILE_CASE_SENSITIVE_INFORMATION
841-
{
842-
ULONG Flags;
843-
} FILE_CASE_SENSITIVE_INFORMATION, *PFILE_CASE_SENSITIVE_INFORMATION;
844-
845792
// private
846793
typedef enum _FILE_KNOWN_FOLDER_TYPE
847794
{

‎MemoryModule/Initialize.cpp‎

Lines changed: 70 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -438,6 +438,7 @@ NTSTATUS InitializeLockHeld() {
438438
status = STATUS_NOT_SUPPORTED;
439439
}
440440
else {
441+
++MmpGlobalDataPtr->ReferenceCount;
441442
status = STATUS_SUCCESS;
442443
}
443444
}
@@ -448,6 +449,7 @@ NTSTATUS InitializeLockHeld() {
448449
MmpGlobalDataPtr->MajorVersion = MEMORY_MODULE_MAJOR_VERSION;
449450
MmpGlobalDataPtr->MinorVersion = MEMORY_MODULE_MINOR_VERSION;
450451
MmpGlobalDataPtr->BaseAddress = MmpGlobalDataPtr;
452+
MmpGlobalDataPtr->ReferenceCount = 1;
451453

452454
GetSystemInfo(&MmpGlobalDataPtr->SystemInfo);
453455

@@ -504,16 +506,74 @@ NTSTATUS InitializeLockHeld() {
504506
return status;
505507
}
506508

507-
NTSTATUS NTAPI Initialize() {
509+
NTSTATUS NTAPI MmInitialize() {
508510
NTSTATUS status;
509511

510-
RtlAcquirePebLock();
511-
status = InitializeLockHeld();
512-
RtlReleasePebLock();
512+
PVOID cookie;
513+
LdrLockLoaderLock(LDR_LOCK_LOADER_LOCK_FLAG_RAISE_ON_ERRORS, nullptr, &cookie);
514+
515+
__try {
516+
status = InitializeLockHeld();
517+
}
518+
__finally {
519+
LdrUnlockLoaderLock(LDR_UNLOCK_LOADER_LOCK_FLAG_RAISE_ON_ERRORS, cookie);
520+
}
513521

514522
return status;
515523
}
516524

525+
NTSTATUS CleanupLockHeld() {
526+
527+
PLIST_ENTRY ListHead = &NtCurrentPeb()->Ldr->InLoadOrderModuleList, ListEntry = ListHead->Flink;
528+
PLDR_DATA_TABLE_ENTRY CurEntry;
529+
530+
while (ListEntry != ListHead) {
531+
CurEntry = CONTAINING_RECORD(ListEntry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
532+
ListEntry = ListEntry->Flink;
533+
534+
if (IsValidMemoryModuleHandle((HMEMORYMODULE)CurEntry->DllBase)) {
535+
536+
//
537+
// Make sure all memory module is unloaded.
538+
//
539+
540+
return STATUS_NOT_SUPPORTED;
541+
}
542+
}
543+
544+
if (--MmpGlobalDataPtr->ReferenceCount > 0) {
545+
return STATUS_SUCCESS;
546+
}
547+
548+
MmpTlsCleanup();
549+
MmpCleanupDotNetHooks();
550+
551+
NtUnmapViewOfSection(NtCurrentProcess(), MmpGlobalDataPtr->BaseAddress);
552+
MmpGlobalDataPtr = nullptr;
553+
return STATUS_SUCCESS;
554+
}
555+
556+
NTSTATUS NTAPI MmCleanup() {
557+
NTSTATUS status;
558+
PVOID cookie;
559+
LdrLockLoaderLock(LDR_LOCK_LOADER_LOCK_FLAG_RAISE_ON_ERRORS, nullptr, &cookie);
560+
561+
__try {
562+
563+
if (MmpGlobalDataPtr == nullptr) {
564+
status = STATUS_ACCESS_VIOLATION;
565+
__leave;
566+
}
567+
568+
status = CleanupLockHeld();
569+
}
570+
__finally {
571+
LdrUnlockLoaderLock(LDR_UNLOCK_LOADER_LOCK_FLAG_RAISE_ON_ERRORS, cookie);
572+
}
573+
574+
return status;
575+
}
576+
517577
#ifdef _USRDLL
518578
extern "C" __declspec(dllexport) BOOL WINAPI ReflectiveMapDll(HMODULE hModule) {
519579
PIMAGE_NT_HEADERS headers = RtlImageNtHeader(hModule);
@@ -542,7 +602,8 @@ extern "C" __declspec(dllexport) BOOL WINAPI ReflectiveMapDll(HMODULE hModule) {
542602

543603
BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved) {
544604
if (ul_reason_for_call == DLL_PROCESS_ATTACH) {
545-
if (NT_SUCCESS(Initialize())) {
605+
#ifdef _HAS_AUTO_INITIALIZE
606+
if (NT_SUCCESS(MmInitialize())) {
546607
if (lpReserved == (PVOID)-1) {
547608
if (!ReflectiveMapDll(hModule)) {
548609
RtlRaiseStatus(STATUS_NOT_SUPPORTED);
@@ -553,10 +614,13 @@ BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReser
553614
}
554615

555616
return FALSE;
617+
#endif
556618
}
557619

558620
return TRUE;
559621
}
560622
#else
561-
const NTSTATUS Initializer = Initialize();
623+
#ifdef _HAS_AUTO_INITIALIZE
624+
const NTSTATUS Initializer = MmInitialize();
625+
#endif
562626
#endif

‎MemoryModule/Initialize.h‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
#pragma once
2+
3+
NTSTATUS NTAPI MmInitialize();
4+
NTSTATUS NTAPI MmCleanup();
5+
6+
//
7+
// This function is available only if the MMPP is compiled as a DLL.
8+
//
9+
BOOL WINAPI ReflectiveMapDll(HMODULE hModule);

‎MemoryModule/LoadDllMemoryApi.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
typedef HMODULE HMEMORYMODULE;
55
#include "Loader.h"
6+
#include "Initialize.h"
67

78
#define MemoryModuleToModule(_hMemoryModule_) (_hMemoryModule_)
89

‎MemoryModule/Loader.cpp‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,14 @@ NTSTATUS NTAPI LdrLoadDllMemoryExW(
6161
__try {
6262
*BaseAddress = nullptr;
6363
if (LdrEntry)*LdrEntry = nullptr;
64-
if (!RtlIsValidImageBuffer(BufferAddress, &BufferSize) && !(dwFlags & LOAD_FLAGS_PASS_IMAGE_CHECK))status = STATUS_INVALID_IMAGE_FORMAT;
64+
65+
if (!RtlIsValidImageBuffer(BufferAddress, &BufferSize) && !(dwFlags & LOAD_FLAGS_PASS_IMAGE_CHECK)) {
66+
status = STATUS_INVALID_IMAGE_FORMAT;
67+
}
68+
69+
if (MmpGlobalDataPtr == nullptr) {
70+
status = STATUS_INVALID_PARAMETER;
71+
}
6572
}
6673
__except (EXCEPTION_EXECUTE_HANDLER) {
6774
status = GetExceptionCode();
@@ -229,6 +236,11 @@ NTSTATUS NTAPI LdrUnloadDllMemory(_In_ HMEMORYMODULE BaseAddress) {
229236
break;
230237
}
231238

239+
if (MmpGlobalDataPtr == nullptr) {
240+
status = STATUS_INVALID_PARAMETER;
241+
break;
242+
}
243+
232244
//Mapping dll failed
233245
if (!module->MappedDll) {
234246
module->underUnload = true;

‎MemoryModule/MemoryModule.vcxproj‎

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,7 @@
152152
<ClInclude Include="..\3rdparty\phnt\include\subprocesstag.h" />
153153
<ClInclude Include="..\3rdparty\phnt\include\winsta.h" />
154154
<ClInclude Include="ImportTable.h" />
155+
<ClInclude Include="Initialize.h" />
155156
<ClInclude Include="LoadDllMemoryApi.h" />
156157
<ClInclude Include="LoaderPrivate.h" />
157158
<ClInclude Include="MemoryModule.h" />
@@ -419,7 +420,7 @@
419420
<PrecompiledHeader>NotUsing</PrecompiledHeader>
420421
<WarningLevel>Level3</WarningLevel>
421422
<SDLCheck>true</SDLCheck>
422-
<PreprocessorDefinitions>_MEMORY_MODULE;WIN32;_DEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
423+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;WIN32;_DEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
423424
<ConformanceMode>true</ConformanceMode>
424425
<PrecompiledHeaderFile>
425426
</PrecompiledHeaderFile>
@@ -437,7 +438,7 @@
437438
<PrecompiledHeader>NotUsing</PrecompiledHeader>
438439
<WarningLevel>Level3</WarningLevel>
439440
<SDLCheck>true</SDLCheck>
440-
<PreprocessorDefinitions>_MEMORY_MODULE;WIN32;_DEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
441+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;WIN32;_DEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
441442
<ConformanceMode>true</ConformanceMode>
442443
<PrecompiledHeaderFile>
443444
</PrecompiledHeaderFile>
@@ -456,7 +457,7 @@
456457
<PrecompiledHeader>NotUsing</PrecompiledHeader>
457458
<WarningLevel>Level3</WarningLevel>
458459
<SDLCheck>true</SDLCheck>
459-
<PreprocessorDefinitions>_MEMORY_MODULE;WIN32;_DEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
460+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;WIN32;_DEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
460461
<ConformanceMode>true</ConformanceMode>
461462
<PrecompiledHeaderFile>
462463
</PrecompiledHeaderFile>
@@ -475,7 +476,7 @@
475476
<PrecompiledHeader>NotUsing</PrecompiledHeader>
476477
<WarningLevel>Level3</WarningLevel>
477478
<SDLCheck>true</SDLCheck>
478-
<PreprocessorDefinitions>_MEMORY_MODULE;WIN32;_DEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
479+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;WIN32;_DEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
479480
<ConformanceMode>true</ConformanceMode>
480481
<PrecompiledHeaderFile>
481482
</PrecompiledHeaderFile>
@@ -494,7 +495,7 @@
494495
<PrecompiledHeader>NotUsing</PrecompiledHeader>
495496
<WarningLevel>Level3</WarningLevel>
496497
<SDLCheck>true</SDLCheck>
497-
<PreprocessorDefinitions>_MEMORY_MODULE;_DEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
498+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;_DEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
498499
<ConformanceMode>true</ConformanceMode>
499500
<PrecompiledHeaderFile>
500501
</PrecompiledHeaderFile>
@@ -512,7 +513,7 @@
512513
<PrecompiledHeader>NotUsing</PrecompiledHeader>
513514
<WarningLevel>Level3</WarningLevel>
514515
<SDLCheck>true</SDLCheck>
515-
<PreprocessorDefinitions>_MEMORY_MODULE;_DEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
516+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;_DEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
516517
<ConformanceMode>true</ConformanceMode>
517518
<PrecompiledHeaderFile>
518519
</PrecompiledHeaderFile>
@@ -531,7 +532,7 @@
531532
<PrecompiledHeader>NotUsing</PrecompiledHeader>
532533
<WarningLevel>Level3</WarningLevel>
533534
<SDLCheck>true</SDLCheck>
534-
<PreprocessorDefinitions>_MEMORY_MODULE;_DEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
535+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;_DEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
535536
<ConformanceMode>true</ConformanceMode>
536537
<PrecompiledHeaderFile>
537538
</PrecompiledHeaderFile>
@@ -550,7 +551,7 @@
550551
<PrecompiledHeader>NotUsing</PrecompiledHeader>
551552
<WarningLevel>Level3</WarningLevel>
552553
<SDLCheck>true</SDLCheck>
553-
<PreprocessorDefinitions>_MEMORY_MODULE;_DEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
554+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;_DEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
554555
<ConformanceMode>true</ConformanceMode>
555556
<PrecompiledHeaderFile>
556557
</PrecompiledHeaderFile>
@@ -571,7 +572,7 @@
571572
<FunctionLevelLinking>true</FunctionLevelLinking>
572573
<IntrinsicFunctions>true</IntrinsicFunctions>
573574
<SDLCheck>true</SDLCheck>
574-
<PreprocessorDefinitions>_MEMORY_MODULE;WIN32;NDEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
575+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;WIN32;NDEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
575576
<ConformanceMode>true</ConformanceMode>
576577
<PrecompiledHeaderFile>
577578
</PrecompiledHeaderFile>
@@ -593,7 +594,7 @@
593594
<FunctionLevelLinking>true</FunctionLevelLinking>
594595
<IntrinsicFunctions>true</IntrinsicFunctions>
595596
<SDLCheck>true</SDLCheck>
596-
<PreprocessorDefinitions>_MEMORY_MODULE;WIN32;NDEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
597+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;WIN32;NDEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
597598
<ConformanceMode>true</ConformanceMode>
598599
<PrecompiledHeaderFile>
599600
</PrecompiledHeaderFile>
@@ -616,7 +617,7 @@
616617
<FunctionLevelLinking>true</FunctionLevelLinking>
617618
<IntrinsicFunctions>true</IntrinsicFunctions>
618619
<SDLCheck>true</SDLCheck>
619-
<PreprocessorDefinitions>_MEMORY_MODULE;WIN32;NDEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
620+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;WIN32;NDEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
620621
<ConformanceMode>true</ConformanceMode>
621622
<PrecompiledHeaderFile>
622623
</PrecompiledHeaderFile>
@@ -639,7 +640,7 @@
639640
<FunctionLevelLinking>true</FunctionLevelLinking>
640641
<IntrinsicFunctions>true</IntrinsicFunctions>
641642
<SDLCheck>true</SDLCheck>
642-
<PreprocessorDefinitions>_MEMORY_MODULE;WIN32;NDEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
643+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;WIN32;NDEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
643644
<ConformanceMode>true</ConformanceMode>
644645
<PrecompiledHeaderFile>
645646
</PrecompiledHeaderFile>
@@ -662,7 +663,7 @@
662663
<FunctionLevelLinking>true</FunctionLevelLinking>
663664
<IntrinsicFunctions>true</IntrinsicFunctions>
664665
<SDLCheck>true</SDLCheck>
665-
<PreprocessorDefinitions>_MEMORY_MODULE;NDEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
666+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;NDEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
666667
<ConformanceMode>true</ConformanceMode>
667668
<PrecompiledHeaderFile>
668669
</PrecompiledHeaderFile>
@@ -684,7 +685,7 @@
684685
<FunctionLevelLinking>true</FunctionLevelLinking>
685686
<IntrinsicFunctions>true</IntrinsicFunctions>
686687
<SDLCheck>true</SDLCheck>
687-
<PreprocessorDefinitions>_MEMORY_MODULE;NDEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
688+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;NDEBUG;_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
688689
<ConformanceMode>true</ConformanceMode>
689690
<PrecompiledHeaderFile>
690691
</PrecompiledHeaderFile>
@@ -707,7 +708,7 @@
707708
<FunctionLevelLinking>true</FunctionLevelLinking>
708709
<IntrinsicFunctions>true</IntrinsicFunctions>
709710
<SDLCheck>true</SDLCheck>
710-
<PreprocessorDefinitions>_MEMORY_MODULE;NDEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
711+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;NDEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
711712
<ConformanceMode>true</ConformanceMode>
712713
<PrecompiledHeaderFile>
713714
</PrecompiledHeaderFile>
@@ -730,7 +731,7 @@
730731
<FunctionLevelLinking>true</FunctionLevelLinking>
731732
<IntrinsicFunctions>true</IntrinsicFunctions>
732733
<SDLCheck>true</SDLCheck>
733-
<PreprocessorDefinitions>_MEMORY_MODULE;NDEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
734+
<PreprocessorDefinitions>_MEMORY_MODULE;_HAS_AUTO_INITIALIZE;NDEBUG;_USRDLL;%(PreprocessorDefinitions)</PreprocessorDefinitions>
734735
<ConformanceMode>true</ConformanceMode>
735736
<PrecompiledHeaderFile>
736737
</PrecompiledHeaderFile>

‎MemoryModule/MemoryModule.vcxproj.filters‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -260,6 +260,9 @@
260260
<ClInclude Include="ImportTable.h">
261261
<Filter>Header Files</Filter>
262262
</ClInclude>
263+
<ClInclude Include="Initialize.h">
264+
<Filter>Header Files</Filter>
265+
</ClInclude>
263266
</ItemGroup>
264267
<ItemGroup>
265268
<None Include="..\README.md">

0 commit comments

Comments
 (0)