Mmp tls fixup

3rdparty/phnt/include/ntioapi.h

@@ -774,22 +774,6 @@ typedef struct _FILE_ID_EXTD_BOTH_DIR_INFORMATION
     WCHAR FileName[1];
 } FILE_ID_EXTD_BOTH_DIR_INFORMATION, *PFILE_ID_EXTD_BOTH_DIR_INFORMATION;
 
-// private
-typedef struct _FILE_STAT_INFORMATION
-{
-    LARGE_INTEGER FileId;
-    LARGE_INTEGER CreationTime;
-    LARGE_INTEGER LastAccessTime;
-    LARGE_INTEGER LastWriteTime;
-    LARGE_INTEGER ChangeTime;
-    LARGE_INTEGER AllocationSize;
-    LARGE_INTEGER EndOfFile;
-    ULONG FileAttributes;
-    ULONG ReparseTag;
-    ULONG NumberOfLinks;
-    ACCESS_MASK EffectiveAccess;
-} FILE_STAT_INFORMATION, *PFILE_STAT_INFORMATION;
-
 // private
 typedef struct _FILE_MEMORY_PARTITION_INFORMATION
 {
@@ -805,43 +789,6 @@ typedef struct _FILE_MEMORY_PARTITION_INFORMATION
     } Flags;
 } FILE_MEMORY_PARTITION_INFORMATION, *PFILE_MEMORY_PARTITION_INFORMATION;
 
-// LxFlags
-#define LX_FILE_METADATA_HAS_UID 0x1
-#define LX_FILE_METADATA_HAS_GID 0x2
-#define LX_FILE_METADATA_HAS_MODE 0x4
-#define LX_FILE_METADATA_HAS_DEVICE_ID 0x8
-#define LX_FILE_CASE_SENSITIVE_DIR 0x10
-
-// private
-typedef struct _FILE_STAT_LX_INFORMATION
-{
-    LARGE_INTEGER FileId;
-    LARGE_INTEGER CreationTime;
-    LARGE_INTEGER LastAccessTime;
-    LARGE_INTEGER LastWriteTime;
-    LARGE_INTEGER ChangeTime;
-    LARGE_INTEGER AllocationSize;
-    LARGE_INTEGER EndOfFile;
-    ULONG FileAttributes;
-    ULONG ReparseTag;
-    ULONG NumberOfLinks;
-    ACCESS_MASK EffectiveAccess;
-    ULONG LxFlags;
-    ULONG LxUid;
-    ULONG LxGid;
-    ULONG LxMode;
-    ULONG LxDeviceIdMajor;
-    ULONG LxDeviceIdMinor;
-} FILE_STAT_LX_INFORMATION, *PFILE_STAT_LX_INFORMATION;
-
-#define FILE_CS_FLAG_CASE_SENSITIVE_DIR     0x00000001
-
-// private
-typedef struct _FILE_CASE_SENSITIVE_INFORMATION
-{
-    ULONG Flags;
-} FILE_CASE_SENSITIVE_INFORMATION, *PFILE_CASE_SENSITIVE_INFORMATION;
-
 // private
 typedef enum _FILE_KNOWN_FOLDER_TYPE
 {

MemoryModule/Initialize.cpp

@@ -438,6 +438,7 @@ NTSTATUS InitializeLockHeld() {
 					status = STATUS_NOT_SUPPORTED;
 				}
 				else {
+					++MmpGlobalDataPtr->ReferenceCount;
 					status = STATUS_SUCCESS;
 				}
 			}
@@ -448,6 +449,7 @@ NTSTATUS InitializeLockHeld() {
         MmpGlobalDataPtr->MajorVersion = MEMORY_MODULE_MAJOR_VERSION;
         MmpGlobalDataPtr->MinorVersion = MEMORY_MODULE_MINOR_VERSION;
 		MmpGlobalDataPtr->BaseAddress = MmpGlobalDataPtr;
+		MmpGlobalDataPtr->ReferenceCount = 1;
 
 		GetSystemInfo(&MmpGlobalDataPtr->SystemInfo);
 
@@ -504,16 +506,74 @@ NTSTATUS InitializeLockHeld() {
     return status;
 }
 
-NTSTATUS NTAPI Initialize() {
+NTSTATUS NTAPI MmInitialize() {
     NTSTATUS status;
 
-    RtlAcquirePebLock();
-    status = InitializeLockHeld();
-    RtlReleasePebLock();
+	PVOID cookie;
+	LdrLockLoaderLock(LDR_LOCK_LOADER_LOCK_FLAG_RAISE_ON_ERRORS, nullptr, &cookie);
+
+	__try {
+		status = InitializeLockHeld();
+	}
+	__finally {
+		LdrUnlockLoaderLock(LDR_UNLOCK_LOADER_LOCK_FLAG_RAISE_ON_ERRORS, cookie);
+	}
 
     return status;
 }
 
+NTSTATUS CleanupLockHeld() {
+
+	PLIST_ENTRY ListHead = &NtCurrentPeb()->Ldr->InLoadOrderModuleList, ListEntry = ListHead->Flink;
+	PLDR_DATA_TABLE_ENTRY CurEntry;
+
+	while (ListEntry != ListHead) {
+		CurEntry = CONTAINING_RECORD(ListEntry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
+		ListEntry = ListEntry->Flink;
+
+		if (IsValidMemoryModuleHandle((HMEMORYMODULE)CurEntry->DllBase)) {
+
+			//
+			// Make sure all memory module is unloaded.
+			//
+
+			return STATUS_NOT_SUPPORTED;
+		}
+	}
+
+	if (--MmpGlobalDataPtr->ReferenceCount > 0) {
+		return STATUS_SUCCESS;
+	}
+
+	MmpTlsCleanup();
+	MmpCleanupDotNetHooks();
+
+	NtUnmapViewOfSection(NtCurrentProcess(), MmpGlobalDataPtr->BaseAddress);
+	MmpGlobalDataPtr = nullptr;
+	return STATUS_SUCCESS;
+}
+
+NTSTATUS NTAPI MmCleanup() {
+	NTSTATUS status;
+	PVOID cookie;
+	LdrLockLoaderLock(LDR_LOCK_LOADER_LOCK_FLAG_RAISE_ON_ERRORS, nullptr, &cookie);
+
+	__try {
+
+		if (MmpGlobalDataPtr == nullptr) {
+			status = STATUS_ACCESS_VIOLATION;
+			__leave;
+		}
+
+		status = CleanupLockHeld();
+	}
+	__finally {
+		LdrUnlockLoaderLock(LDR_UNLOCK_LOADER_LOCK_FLAG_RAISE_ON_ERRORS, cookie);
+	}
+
+	return status;
+}
+
 #ifdef _USRDLL
 extern "C" __declspec(dllexport) BOOL WINAPI ReflectiveMapDll(HMODULE hModule) {
 	PIMAGE_NT_HEADERS headers = RtlImageNtHeader(hModule);
@@ -542,7 +602,8 @@ extern "C" __declspec(dllexport) BOOL WINAPI ReflectiveMapDll(HMODULE hModule) {
 
 BOOL APIENTRY DllMain(HMODULE hModule, DWORD  ul_reason_for_call, LPVOID lpReserved) {
 	if (ul_reason_for_call == DLL_PROCESS_ATTACH) {
-		if (NT_SUCCESS(Initialize())) {
+#ifdef _HAS_AUTO_INITIALIZE
+		if (NT_SUCCESS(MmInitialize())) {
 			if (lpReserved == (PVOID)-1) {
 				if (!ReflectiveMapDll(hModule)) {
 					RtlRaiseStatus(STATUS_NOT_SUPPORTED);
@@ -553,10 +614,13 @@ BOOL APIENTRY DllMain(HMODULE hModule, DWORD  ul_reason_for_call, LPVOID lpReser
 		}
 
 		return FALSE;
+#endif
 	}
 
 	return TRUE;
 }
 #else
-const NTSTATUS Initializer = Initialize();
+#ifdef _HAS_AUTO_INITIALIZE
+const NTSTATUS Initializer = MmInitialize();
+#endif
 #endif

MemoryModule/Initialize.h

@@ -0,0 +1,9 @@
+#pragma once
+
+NTSTATUS NTAPI MmInitialize();
+NTSTATUS NTAPI MmCleanup();
+
+//
+// This function is available only if the MMPP is compiled as a DLL.
+//
+BOOL WINAPI ReflectiveMapDll(HMODULE hModule);

MemoryModule/LoadDllMemoryApi.h

@@ -3,32 +3,37 @@
 
 typedef HMODULE HMEMORYMODULE;
 #include "Loader.h"
+#include "Initialize.h"
 
 #define MemoryModuleToModule(_hMemoryModule_) (_hMemoryModule_)
 
 #ifndef NT_SUCCESS
 #define NT_SUCCESS(Status) (((NTSTATUS)(Status)) >= 0)
 #endif
 
-HMEMORYMODULE WINAPI LoadLibraryMemory(_In_ PVOID BufferAddress);
-
-HMEMORYMODULE WINAPI LoadLibraryMemoryExA(
-	_In_ PVOID BufferAddress,
-	_In_ size_t Reserved,
-	_In_opt_ LPCSTR DllBaseName,
-	_In_opt_ LPCSTR DllFullName,
-	_In_ DWORD Flags
-);
-
-HMEMORYMODULE WINAPI LoadLibraryMemoryExW(
-	_In_ PVOID BufferAddress,
-	_In_ size_t Reserved,
-	_In_opt_ LPCWSTR DllBaseName,
-	_In_opt_ LPCWSTR DllFullName,
-	_In_ DWORD Flags
-);
-
-BOOL WINAPI FreeLibraryMemory(_In_ HMEMORYMODULE hMemoryModule);
+extern "C" {
+
+	HMEMORYMODULE WINAPI LoadLibraryMemory(_In_ PVOID BufferAddress);
+
+	HMEMORYMODULE WINAPI LoadLibraryMemoryExA(
+		_In_ PVOID BufferAddress,
+		_In_ size_t Reserved,
+		_In_opt_ LPCSTR DllBaseName,
+		_In_opt_ LPCSTR DllFullName,
+		_In_ DWORD Flags
+	);
+
+	HMEMORYMODULE WINAPI LoadLibraryMemoryExW(
+		_In_ PVOID BufferAddress,
+		_In_ size_t Reserved,
+		_In_opt_ LPCWSTR DllBaseName,
+		_In_opt_ LPCWSTR DllFullName,
+		_In_ DWORD Flags
+	);
+
+	BOOL WINAPI FreeLibraryMemory(_In_ HMEMORYMODULE hMemoryModule);
+
+}
 
 #define NtLoadDllMemory						LdrLoadDllMemory
 #define NtLoadDllMemoryExA					LdrLoadDllMemoryExA

MemoryModule/Loader.cpp

@@ -61,7 +61,14 @@ NTSTATUS NTAPI LdrLoadDllMemoryExW(
 	__try {
 		*BaseAddress = nullptr;
 		if (LdrEntry)*LdrEntry = nullptr;
-		if (!RtlIsValidImageBuffer(BufferAddress, &BufferSize) && !(dwFlags & LOAD_FLAGS_PASS_IMAGE_CHECK))status = STATUS_INVALID_IMAGE_FORMAT;
+
+		if (!RtlIsValidImageBuffer(BufferAddress, &BufferSize) && !(dwFlags & LOAD_FLAGS_PASS_IMAGE_CHECK)) {
+			status = STATUS_INVALID_IMAGE_FORMAT;
+		}
+
+		if (MmpGlobalDataPtr == nullptr) {
+			status = STATUS_INVALID_PARAMETER;
+		}
 	}
 	__except (EXCEPTION_EXECUTE_HANDLER) {
 		status = GetExceptionCode();
@@ -229,6 +236,11 @@ NTSTATUS NTAPI LdrUnloadDllMemory(_In_ HMEMORYMODULE BaseAddress) {
 			break;
 		}
 
+		if (MmpGlobalDataPtr == nullptr) {
+			status = STATUS_INVALID_PARAMETER;
+			break;
+		}
+
 		//Mapping dll failed
 		if (!module->MappedDll) {
 			module->underUnload = true;

MemoryModule/Loader.h

@@ -9,9 +9,6 @@
 #define MEMORY_FEATURE_LDRP_RELEASE_TLS_ENTRY		0x00000040
 #define MEMORY_FEATURE_ALL                          0x0000007f
 
-//Get the implementation of the currently running operating system.
-NTSTATUS NTAPI LdrQuerySystemMemoryModuleFeatures(_Out_ PDWORD pFeatures);
-
 
 /*
 	LdrLoadDllMemoryEx dwFlags
@@ -47,21 +44,24 @@ NTSTATUS NTAPI LdrQuerySystemMemoryModuleFeatures(_Out_ PDWORD pFeatures);
 //Hook for dotnet dlls
 #define LOAD_FLAGS_HOOK_DOT_NET						0x00000010
 
+extern "C" {
 
-NTSTATUS NTAPI LdrLoadDllMemoryExW(
-	_Out_ HMEMORYMODULE* BaseAddress,		// Output module base address
-	_Out_opt_ PVOID* LdrEntry,				// Receive a pointer to the LDR node of the module
-	_In_ DWORD dwFlags,						// Flags
-	_In_ LPVOID BufferAddress,				// Pointer to the dll file data buffer
-	_In_ size_t Reserved,					// Reserved parameter, must be 0
-	_In_opt_ LPCWSTR DllName,				// Module file name
-	_In_opt_ LPCWSTR DllFullName			// Module file full path
-);
+	//Get the implementation of the currently running operating system.
+	NTSTATUS NTAPI LdrQuerySystemMemoryModuleFeatures(_Out_ PDWORD pFeatures);
+
+	NTSTATUS NTAPI LdrLoadDllMemoryExW(
+		_Out_ HMEMORYMODULE* BaseAddress,		// Output module base address
+		_Out_opt_ PVOID* LdrEntry,				// Receive a pointer to the LDR node of the module
+		_In_ DWORD dwFlags,						// Flags
+		_In_ LPVOID BufferAddress,				// Pointer to the dll file data buffer
+		_In_ size_t Reserved,					// Reserved parameter, must be 0
+		_In_opt_ LPCWSTR DllName,				// Module file name
+		_In_opt_ LPCWSTR DllFullName			// Module file full path
+	);
 
-//Unload modules previously loaded from memory
-NTSTATUS NTAPI LdrUnloadDllMemory(_In_ HMEMORYMODULE BaseAddress);
+	//Unload modules previously loaded from memory
+	NTSTATUS NTAPI LdrUnloadDllMemory(_In_ HMEMORYMODULE BaseAddress);
 
-extern "C" {
 	__declspec(noreturn) VOID NTAPI LdrUnloadDllMemoryAndExitThread(
 		_In_ HMEMORYMODULE BaseAddress,
 		_In_ DWORD dwExitCode