1. Add Win10 forward export support

2. Fixed some bugs

MemoryModule/MemoryModule.cpp

@@ -4,6 +4,7 @@
 #include <tchar.h>
 #include "rtltype.h"
 #include "ntstatus.h"
+#include "Native.h"
 #include <algorithm>
 
 #if _MSC_VER
@@ -134,7 +135,7 @@ static BOOL CheckSize(size_t size, size_t expected) {
 	return TRUE;
 }
 
-static BOOL CopySections(const unsigned char* data, size_t size, PMEMORYMODULE module) {
+static BOOL CopySections(const unsigned char* data, PMEMORYMODULE module) {
 	LPBYTE codeBase = module->codeBase;
 	LPVOID dest;
 	PIMAGE_NT_HEADERS headers = GetImageNtHeaders(module);
@@ -145,7 +146,12 @@ static BOOL CopySections(const unsigned char* data, size_t size, PMEMORYMODULE m
 		alloc_size = headers->OptionalHeader.SectionAlignment;
 		cp = false;
 		if (section->SizeOfRawData) {
-			if (!CheckSize(size, static_cast<size_t>(section->PointerToRawData) + section->SizeOfRawData)) return FALSE;
+			__try {
+				ProbeForRead(data, static_cast<size_t>(section->PointerToRawData) + section->SizeOfRawData);
+			}
+			__except (EXCEPTION_EXECUTE_HANDLER) {
+				return FALSE;
+			}
 			alloc_size = section->SizeOfRawData;
 			cp = true;
 		}
@@ -354,15 +360,23 @@ static BOOL BuildImportTable(PMEMORYMODULE module) {
 	return result;
 }
 
-HMEMORYMODULE MemoryLoadLibrary(const void* data, size_t size) {
+//static BOOL PerformForwardExport(PMEMORYMODULE module) {
+//	PIMAGE_EXPORT_DIRECTORY exports;
+//	PIMAGE_NT_HEADERS headers = GetImageNtHeaders(module);
+//	PIMAGE_DATA_DIRECTORY directory = GET_HEADER_DICTIONARY(headers, IMAGE_DIRECTORY_ENTRY_EXPORT);
+//	if (!directory->Size)return TRUE;
+//	return FALSE;
+//
+//}
+
+HMEMORYMODULE MemoryLoadLibrary(const void* data) {
 	PMEMORYMODULE hMemoryModule = nullptr;
 	PIMAGE_DOS_HEADER dos_header, new_dos_header;
 	PIMAGE_NT_HEADERS old_header, new_header;
-	unsigned char* code;
+	unsigned char* base;
 	ptrdiff_t locationDelta;
-	SYSTEM_INFO sysInfo;
+	static SYSTEM_INFO sysInfo{};
 	PIMAGE_SECTION_HEADER section;
-	DWORD i;
 	size_t optionalSectionSize;
 	size_t lastSectionEnd = 0;
 	size_t alignedImageSize;
@@ -371,40 +385,36 @@ HMEMORYMODULE MemoryLoadLibrary(const void* data, size_t size) {
 	POINTER_LIST* blockedMemory = nullptr;
 #endif
 
-	if (!CheckSize(size, sizeof(IMAGE_DOS_HEADER))) return nullptr;
-	dos_header = (PIMAGE_DOS_HEADER)data;
-	if (dos_header->e_magic != IMAGE_DOS_SIGNATURE) {
-		SetLastError(ERROR_BAD_EXE_FORMAT);
+	__try {
+		ProbeForRead(data, sizeof(IMAGE_DOS_HEADER));
+		dos_header = (PIMAGE_DOS_HEADER)data;
+		if (dos_header->e_magic != IMAGE_DOS_SIGNATURE) {
+			SetLastError(ERROR_BAD_EXE_FORMAT);
+			return nullptr;
+		}
+		ProbeForRead(data, dos_header->e_lfanew + sizeof(IMAGE_NT_HEADERS));
+		old_header = (PIMAGE_NT_HEADERS)((size_t)data + dos_header->e_lfanew);
+		if (old_header->Signature != IMAGE_NT_SIGNATURE ||
+			!ProbeForRead(data, old_header->OptionalHeader.SizeOfHeaders) ||
+			old_header->FileHeader.Machine != HOST_MACHINE ||
+			old_header->OptionalHeader.SectionAlignment & 1) {
+			SetLastError(ERROR_BAD_EXE_FORMAT);
+			return nullptr;
+		}
+		//only dll image support
+		if (!(old_header->FileHeader.Characteristics & IMAGE_FILE_DLL)) {
+			SetLastError(ERROR_NOT_SUPPORTED);
+			return nullptr;
+		}
+	}
+	__except (EXCEPTION_EXECUTE_HANDLER) {
+		SetLastError(ERROR_INVALID_DATA);
 		return nullptr;
 	}
-
-	if (!CheckSize(size, dos_header->e_lfanew + sizeof(IMAGE_NT_HEADERS))) return nullptr;
-	old_header = (PIMAGE_NT_HEADERS) & ((const unsigned char*)(data))[dos_header->e_lfanew];
-	if (old_header->Signature != IMAGE_NT_SIGNATURE) {
-		SetLastError(ERROR_BAD_EXE_FORMAT);
-		return nullptr;
-	}
-
-	if (old_header->FileHeader.Machine != HOST_MACHINE) {
-		SetLastError(ERROR_BAD_EXE_FORMAT);
-		return nullptr;
-	}
-
-	if (old_header->OptionalHeader.SectionAlignment & 1) {
-		// Only support section alignments that are a multiple of 2
-		SetLastError(ERROR_BAD_EXE_FORMAT);
-		return nullptr;
-	}
-
-	//only dll image support
-	if (!(old_header->FileHeader.Characteristics & IMAGE_FILE_DLL)) {
-		SetLastError(ERROR_NOT_SUPPORTED);
-		return nullptr;
-	}
-
+	
 	section = IMAGE_FIRST_SECTION(old_header);
 	optionalSectionSize = old_header->OptionalHeader.SectionAlignment;
-	for (i = 0; i < old_header->FileHeader.NumberOfSections; i++, section++) {
+	for (DWORD i = 0; i < old_header->FileHeader.NumberOfSections; i++, section++) {
 		size_t endOfSection;
 		if (section->SizeOfRawData == 0) {
 			// Section without data in the DLL
@@ -419,7 +429,7 @@ HMEMORYMODULE MemoryLoadLibrary(const void* data, size_t size) {
 		}
 	}
 
-	GetNativeSystemInfo(&sysInfo);
+	if (!sysInfo.dwPageSize)GetNativeSystemInfo(&sysInfo);
 	alignedImageSize = AlignValueUp(old_header->OptionalHeader.SizeOfImage, sysInfo.dwPageSize);
 	if (alignedImageSize != AlignValueUp(lastSectionEnd, sysInfo.dwPageSize)) {
 		SetLastError(ERROR_BAD_EXE_FORMAT);
@@ -430,12 +440,12 @@ HMEMORYMODULE MemoryLoadLibrary(const void* data, size_t size) {
 	// reserve memory for image of library
 	// XXX: is it correct to commit the complete memory region at once?
 	//      calling DllEntry raises an exception if we don't...
-	if (!(code = (LPBYTE)VirtualAlloc((LPVOID)(old_header->OptionalHeader.ImageBase), alignedImageSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE))) {
+	if (!(base = (LPBYTE)VirtualAlloc((LPVOID)(old_header->OptionalHeader.ImageBase), alignedImageSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE))) {
 		if (!(old_header->OptionalHeader.DllCharacteristics & IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE)) {
 			SetLastError(ERROR_BAD_EXE_FORMAT);
 			return nullptr;
 		}
-		if (!(code = (LPBYTE)VirtualAlloc(nullptr, alignedImageSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE))) {
+		if (!(base = (LPBYTE)VirtualAlloc(nullptr, alignedImageSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE))) {
 			SetLastError(ERROR_OUTOFMEMORY);
 			return nullptr;
 		}
@@ -443,20 +453,20 @@ HMEMORYMODULE MemoryLoadLibrary(const void* data, size_t size) {
 
 #ifdef _WIN64
 	// Memory block may not span 4 GB boundaries.
-	while ((((uintptr_t)code) >> 32) < (((uintptr_t)(code + alignedImageSize)) >> 32)) {
+	while ((((uintptr_t)base) >> 32) < (((uintptr_t)(base + alignedImageSize)) >> 32)) {
 		POINTER_LIST* node = new POINTER_LIST;
 		if (!node) {
-			VirtualFree(code, 0, MEM_RELEASE);
+			VirtualFree(base, 0, MEM_RELEASE);
 			FreePointerList(blockedMemory);
 			SetLastError(ERROR_OUTOFMEMORY);
 			return nullptr;
 		}
 
 		node->next = blockedMemory;
-		node->address = code;
+		node->address = base;
 		blockedMemory = node;
 
-		if (!(code = (LPBYTE)VirtualAlloc(nullptr, alignedImageSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE))) {
+		if (!(base = (LPBYTE)VirtualAlloc(nullptr, alignedImageSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE))) {
 			FreePointerList(blockedMemory);
 			SetLastError(ERROR_OUTOFMEMORY);
 			return nullptr;
@@ -464,11 +474,11 @@ HMEMORYMODULE MemoryLoadLibrary(const void* data, size_t size) {
 	}
 #endif
 
-	new_dos_header = (PIMAGE_DOS_HEADER)code;
-	new_header = (PIMAGE_NT_HEADERS)(code + dos_header->e_lfanew);
-	hMemoryModule = (PMEMORYMODULE)(code + old_header->OptionalHeader.SizeOfHeaders);
+	new_dos_header = (PIMAGE_DOS_HEADER)base;
+	new_header = (PIMAGE_NT_HEADERS)(base + dos_header->e_lfanew);
+	hMemoryModule = (PMEMORYMODULE)(base + old_header->OptionalHeader.SizeOfHeaders);
 	RtlZeroMemory(hMemoryModule, sizeof(MEMORYMODULE));
-	hMemoryModule->codeBase = code;
+	hMemoryModule->codeBase = base;
 	hMemoryModule->pageSize = sysInfo.dwPageSize;
 	hMemoryModule->Signature = MEMORY_MODULE_SIGNATURE;
 	hMemoryModule->SizeofHeaders = old_header->OptionalHeader.SizeOfHeaders;
@@ -477,18 +487,14 @@ HMEMORYMODULE MemoryLoadLibrary(const void* data, size_t size) {
 	hMemoryModule->blockedMemory = blockedMemory;
 #endif
 
-	if (!CheckSize(size, old_header->OptionalHeader.SizeOfHeaders)) {
-		goto error;
-	}
-
 	// copy PE header to code
 	memcpy(new_dos_header, dos_header, old_header->OptionalHeader.SizeOfHeaders);
 	new_header->OptionalHeader.SizeOfImage = (DWORD)(alignedImageSize);
-	new_header->OptionalHeader.ImageBase = (size_t)code;
+	new_header->OptionalHeader.ImageBase = (size_t)base;
 	new_header->OptionalHeader.BaseOfCode = headers_align;
 
 	// copy sections from DLL file block to new memory location
-	if (!CopySections((LPBYTE)data, size, hMemoryModule)) goto error;
+	if (!CopySections((LPBYTE)data, hMemoryModule)) goto error;
 
 	// adjust base address of imported data
 	locationDelta = (ptrdiff_t)(hMemoryModule->codeBase - old_header->OptionalHeader.ImageBase);
@@ -509,7 +515,7 @@ HMEMORYMODULE MemoryLoadLibrary(const void* data, size_t size) {
 	if (new_header->OptionalHeader.AddressOfEntryPoint) {
 		__try {
 			// notify library about attaching to process
-			if (!((DllEntryProc)(code + new_header->OptionalHeader.AddressOfEntryPoint))((HINSTANCE)code, DLL_PROCESS_ATTACH, 0)) {
+			if (!((DllEntryProc)(base + new_header->OptionalHeader.AddressOfEntryPoint))((HINSTANCE)base, DLL_PROCESS_ATTACH, 0)) {
 				SetLastError(ERROR_DLL_INIT_FAILED);
 				goto error;
 			}
@@ -521,7 +527,7 @@ HMEMORYMODULE MemoryLoadLibrary(const void* data, size_t size) {
 		hMemoryModule->initialized = TRUE;
 	}
 	
-	return code;
+	return base;
 error:
 	// cleanup
 	MemoryFreeLibrary(hMemoryModule);

MemoryModule/MemoryModule.h

@@ -101,7 +101,7 @@ extern "C" {
      * All dependencies are resolved using default LoadLibrary/GetProcAddress
      * calls through the Windows API.
      */
-    HMEMORYMODULE MemoryLoadLibrary(const void*, size_t);
+    HMEMORYMODULE MemoryLoadLibrary(const void*);
 
     /**
      * Get address of exported method. Supports loading both by name and by

MemoryModule/Native.cpp

@@ -364,3 +364,24 @@ PVOID NTAPI RtlEncodeSystemPointer(PVOID Pointer) {
 PVOID NTAPI RtlDecodeSystemPointer(PVOID Pointer) {
 	return decltype(&RtlDecodeSystemPointer)(RtlGetNtProcAddress("RtlDecodeSystemPointer"))(Pointer);
 }
+
+BOOLEAN NTAPI VirtualAccessCheck(LPCVOID pBuffer, size_t size, ACCESS_MASK protect) {
+	MEMORY_BASIC_INFORMATION mbi{};
+	SIZE_T len = 0;
+	if (!NT_SUCCESS(NtQueryVirtualMemory(NtCurrentProcess(), const_cast<PVOID>(pBuffer), MemoryBasicInformation, &mbi, sizeof(mbi), &len)) ||
+		!(mbi.Protect & protect)) {
+		RaiseException(EXCEPTION_ACCESS_VIOLATION, 0, 0, nullptr);
+		return FALSE;
+	}
+	return TRUE;
+}
+
+NTSTATUS NTAPI LdrLockLoaderLock(size_t Flags, size_t* State, size_t* Cookie) {
+	return (decltype(&LdrLockLoaderLock)(RtlGetNtProcAddress("LdrLockLoaderLock")))(Flags, State, Cookie);
+}
+NTSTATUS NTAPI LdrUnlockLoaderLock(size_t Flags, size_t Cookie) {
+	return (decltype(&LdrUnlockLoaderLock)(RtlGetNtProcAddress("LdrUnlockLoaderLock")))(Flags, Cookie);
+}
+NTSTATUS NTAPI LdrUnloadDll(IN HANDLE ModuleHandle) {
+	return (decltype(&LdrUnloadDll)(RtlGetNtProcAddress("LdrUnloadDll")))(ModuleHandle);
+}

MemoryModule/Native.h

@@ -372,6 +372,60 @@ typedef enum _PROCESSINFOCLASS {
 	ProcessLeapSecondInformation, // PROCESS_LEAP_SECOND_INFORMATION
 	MaxProcessInfoClass
 } PROCESSINFOCLASS;
+typedef enum _THREADINFOCLASS {
+	ThreadBasicInformation, // q: THREAD_BASIC_INFORMATION
+	ThreadTimes, // q: KERNEL_USER_TIMES
+	ThreadPriority, // s: KPRIORITY
+	ThreadBasePriority, // s: LONG
+	ThreadAffinityMask, // s: KAFFINITY
+	ThreadImpersonationToken, // s: HANDLE
+	ThreadDescriptorTableEntry, // q: DESCRIPTOR_TABLE_ENTRY (or WOW64_DESCRIPTOR_TABLE_ENTRY)
+	ThreadEnableAlignmentFaultFixup, // s: BOOLEAN
+	ThreadEventPair,
+	ThreadQuerySetWin32StartAddress, // q: PVOID
+	ThreadZeroTlsCell, // 10
+	ThreadPerformanceCount, // q: LARGE_INTEGER
+	ThreadAmILastThread, // q: ULONG
+	ThreadIdealProcessor, // s: ULONG
+	ThreadPriorityBoost, // qs: ULONG
+	ThreadSetTlsArrayAddress,
+	ThreadIsIoPending, // q: ULONG
+	ThreadHideFromDebugger, // s: void
+	ThreadBreakOnTermination, // qs: ULONG
+	ThreadSwitchLegacyState,
+	ThreadIsTerminated, // q: ULONG // 20
+	ThreadLastSystemCall, // q: THREAD_LAST_SYSCALL_INFORMATION
+	ThreadIoPriority, // qs: IO_PRIORITY_HINT
+	ThreadCycleTime, // q: THREAD_CYCLE_TIME_INFORMATION
+	ThreadPagePriority, // q: ULONG
+	ThreadActualBasePriority,
+	ThreadTebInformation, // q: THREAD_TEB_INFORMATION (requires THREAD_GET_CONTEXT + THREAD_SET_CONTEXT)
+	ThreadCSwitchMon,
+	ThreadCSwitchPmu,
+	ThreadWow64Context, // q: WOW64_CONTEXT
+	ThreadGroupInformation, // q: GROUP_AFFINITY // 30
+	ThreadUmsInformation, // q: THREAD_UMS_INFORMATION
+	ThreadCounterProfiling,
+	ThreadIdealProcessorEx, // q: PROCESSOR_NUMBER
+	ThreadCpuAccountingInformation, // since WIN8
+	ThreadSuspendCount, // since WINBLUE
+	ThreadHeterogeneousCpuPolicy, // q: KHETERO_CPU_POLICY // since THRESHOLD
+	ThreadContainerId, // q: GUID
+	ThreadNameInformation, // qs: THREAD_NAME_INFORMATION
+	ThreadSelectedCpuSets,
+	ThreadSystemThreadInformation, // q: SYSTEM_THREAD_INFORMATION // 40
+	ThreadActualGroupAffinity, // since THRESHOLD2
+	ThreadDynamicCodePolicyInfo,
+	ThreadExplicitCaseSensitivity, // qs: ULONG; s: 0 disables, otherwise enables
+	ThreadWorkOnBehalfTicket,
+	ThreadSubsystemInformation, // q: SUBSYSTEM_INFORMATION_TYPE // since REDSTONE2
+	ThreadDbgkWerReportActive,
+	ThreadAttachContainer,
+	ThreadManageWritesToExecutableMemory, // MANAGE_WRITES_TO_EXECUTABLE_MEMORY // since REDSTONE3
+	ThreadPowerThrottlingState, // THREAD_POWER_THROTTLING_STATE
+	ThreadWorkloadClass, // THREAD_WORKLOAD_CLASS // since REDSTONE5 // 50
+	MaxThreadInfoClass
+} THREADINFOCLASS;
 typedef struct _SYSTEM_PROCESS {
 	ULONG NextEntryOffset;//relative offset
 	ULONG ThreadCount;
@@ -998,6 +1052,13 @@ typedef struct _SECURITY__LOGON_SESSION_DATA {
 	LARGE_INTEGER PasswordMustChange;
 }SECURITY_LOGON_SESSION_DATA, *PSECURITY_LOGON_SESSION_DATA,
 LOGON_SESSION_DATA, *PLOGON_SESSION_DATA;
+typedef struct _INITIAL_TEB {
+	PVOID                StackBase;
+	PVOID                StackLimit;
+	PVOID                StackCommit;
+	PVOID                StackCommitMax;
+	PVOID                StackReserved;
+} INITIAL_TEB, * PINITIAL_TEB;
 
 NTSTATUS NTAPI NtQueryObject(
 	IN HANDLE               ObjectHandle,
@@ -1340,3 +1401,23 @@ typedef struct _UNWIND_INFO {
 NTSTATUS NTAPI NtQuerySystemTime(PLARGE_INTEGER SystemTime);
 PVOID NTAPI RtlEncodeSystemPointer(PVOID Pointer);
 PVOID NTAPI RtlDecodeSystemPointer(PVOID Pointer);
+
+#define NtCurrentProcess()	(HANDLE)-1
+#define NtCurrentThread()	(HANDLE)-2
+
+BOOLEAN NTAPI VirtualAccessCheck(LPCVOID pBuffer, size_t size, ACCESS_MASK protect);
+#define ProbeForRead(pBuffer, size)			VirtualAccessCheck(pBuffer, size, PAGE_READONLY | PAGE_READWRITE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE)
+#define ProbeForWrite(pBuffer, size)		VirtualAccessCheck(pBuffer, size, PAGE_READWRITE | PAGE_EXECUTE_WRITECOPY | PAGE_WRITECOPY | PAGE_EXECUTE_READWRITE)
+#define ProbeForReadWrite(pBuffer, size)	VirtualAccessCheck(pBuffer, size, PAGE_EXECUTE_READWRITE | PAGE_READWRITE)
+#define ProbeForExecute(pBuffer, size)		VirtualAccessCheck(pBuffer, size, PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY)
+
+//Flags
+#define LOCK_RAISE_EXCEPTION 1
+#define LOCK_NO_WAIT_IF_BUSY 2
+//State
+#define LOCK_STATE_NO_ENTER  0
+#define LOCK_STATE_ENTERED   1
+#define LOCK_STATE_LOCK_BUSY 2
+NTSTATUS NTAPI LdrLockLoaderLock(size_t Flags, size_t* State, size_t* Cookie);
+NTSTATUS NTAPI LdrUnlockLoaderLock(size_t Flags, size_t Cookie);
+NTSTATUS NTAPI LdrUnloadDll(IN HANDLE ModuleHandle);

MemoryModule/NativeFunctionsInternal.cpp

@@ -51,26 +51,16 @@ static ULONG NTAPI LdrHashEntry(IN const UNICODE_STRING& str, IN bool _xor = tru
 static HANDLE NTAPI RtlFindtLdrpHeap() {
 	PLIST_ENTRY ListHead, ListEntry;
 	PLDR_DATA_TABLE_ENTRY CurEntry;
+	MEMORY_BASIC_INFORMATION mbi{};
 	static HANDLE result = nullptr;
-	DWORD dwHeaps = 0;
-	HANDLE* hHeaps = nullptr;
-
 	if (result)return result;
-	dwHeaps = GetProcessHeaps(dwHeaps, hHeaps);
-	hHeaps = new HANDLE[dwHeaps];
+
 	ListHead = &NtCurrentPeb()->Ldr->InLoadOrderModuleList;
 	ListEntry = ListHead->Flink;
-	if (ListHead == ListEntry)return nullptr;
+	if (ListHead == ListEntry)return result;
 	CurEntry = CONTAINING_RECORD(ListEntry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
-	GetProcessHeaps(dwHeaps, hHeaps);
-	for (DWORD i = 0; i < dwHeaps; ++i) {
-		if (HeapValidate(hHeaps[i], 0, CurEntry)) {
-			result = hHeaps[i];
-			break;
-		}
-	}
-	delete[]hHeaps;
-	return result;
+	NtQueryVirtualMemory(NtCurrentProcess(), CurEntry, MemoryBasicInformation, &mbi, sizeof(mbi), (PSIZE_T)&ListHead);
+	return result = mbi.AllocationBase;
 }
 static PLDR_DATA_TABLE_ENTRY NTAPI RtlFindNtdllLdrEntry() {
 	PLIST_ENTRY ListHead, ListEntry;
@@ -242,7 +232,7 @@ static PRTL_BALANCED_NODE NTAPI RtlFindLdrpModuleBaseAddressIndex() {
 	return LdrpModuleBaseAddressIndex;
 }
 static NTSTATUS NTAPI NtInsertModuleBaseAddressIndexNode(IN PLDR_DATA_TABLE_ENTRY DataTableEntry, IN PVOID BaseAddress) {
-	auto LdrpModuleBaseAddressIndex = RtlFindLdrpModuleBaseAddressIndex();
+	static auto LdrpModuleBaseAddressIndex = RtlFindLdrpModuleBaseAddressIndex();
 	if (!LdrpModuleBaseAddressIndex)return STATUS_UNSUCCESSFUL;
 
 	PLDR_DATA_TABLE_ENTRY_WIN8 LdrNode = decltype(LdrNode)((size_t)LdrpModuleBaseAddressIndex - offsetof(LDR_DATA_TABLE_ENTRY_WIN8, BaseAddressIndexNode));
@@ -274,20 +264,82 @@ static NTSTATUS NTAPI NtInsertModuleBaseAddressIndexNode(IN PLDR_DATA_TABLE_ENTR
 	return STATUS_SUCCESS;
 }
 static NTSTATUS NTAPI NtRemoveModuleBaseAddressIndexNode(IN PLDR_DATA_TABLE_ENTRY DataTableEntry) {
-	RTL_RB_TREE tree{ RtlFindLdrpModuleBaseAddressIndex() };
+	static RTL_RB_TREE tree{ RtlFindLdrpModuleBaseAddressIndex() };
 	if (!tree.Root)return STATUS_UNSUCCESSFUL;
 
 	RtlRbRemoveNode(&tree, &PLDR_DATA_TABLE_ENTRY_WIN8(DataTableEntry)->BaseAddressIndexNode);
 	return STATUS_SUCCESS;
 }
 
+static NTSTATUS NTAPI NtFreeDependencies(IN PLDR_DATA_TABLE_ENTRY_WIN8 LdrEntry) {
+	_LDR_DDAG_NODE* DependentDdgeNode = nullptr;
+	PLDR_DATA_TABLE_ENTRY_WIN8 ModuleEntry = nullptr;
+	_LDRP_CSLIST* head = (decltype(head))LdrEntry->DdagNode->Dependencies, *entry = head;
+	if (!LdrEntry->DdagNode->Dependencies)return STATUS_SUCCESS;
+
+	//find all dependencies and free
+	do {
+		DependentDdgeNode = entry->Dependent.DependentDdagNode;
+		if (DependentDdgeNode->Modules.Flink->Flink != &DependentDdgeNode->Modules) RaiseException(-1, EXCEPTION_NONCONTINUABLE, 0, nullptr);
+		ModuleEntry = decltype(ModuleEntry)((size_t)DependentDdgeNode->Modules.Flink - offsetof(_LDR_DATA_TABLE_ENTRY_WIN8, NodeModuleLink));
+		if (ModuleEntry->DdagNode != DependentDdgeNode) RaiseException(-1, EXCEPTION_NONCONTINUABLE, 0, nullptr);
+		if (!DependentDdgeNode->IncomingDependencies) RaiseException(-1, EXCEPTION_NONCONTINUABLE, 0, nullptr);
+		_LDRP_CSLIST::_LDRP_CSLIST_INCOMMING* _last = DependentDdgeNode->IncomingDependencies, *_entry = _last;
+		_LDR_DDAG_NODE* CurrentDdagNode;
+		size_t State = 0, Cookies;
+
+		//Acquire LoaderLock
+		do {
+			if (!NT_SUCCESS(LdrLockLoaderLock(LOCK_NO_WAIT_IF_BUSY, &State, &Cookies)))
+				RaiseException(-1, EXCEPTION_NONCONTINUABLE, 0, nullptr);
+		} while (State != LOCK_STATE_ENTERED);
+		do {
+			CurrentDdagNode = (decltype(CurrentDdagNode))((size_t)_entry->IncommingDdagNode & ~1);
+			if (CurrentDdagNode == LdrEntry->DdagNode) {
+				//node is head
+				if (_entry == DependentDdgeNode->IncomingDependencies) {
+					//only one node in list
+					if (_entry->NextIncommingEntry == (PSINGLE_LIST_ENTRY)DependentDdgeNode->IncomingDependencies) {
+						DependentDdgeNode->IncomingDependencies = nullptr;
+					}
+					else {
+						//find the last node in the list
+						PSINGLE_LIST_ENTRY i = _entry->NextIncommingEntry;
+						while (i->Next != (PSINGLE_LIST_ENTRY)_entry)i = i->Next;
+						i->Next = _entry->NextIncommingEntry;
+						DependentDdgeNode->IncomingDependencies = (_LDRP_CSLIST::_LDRP_CSLIST_INCOMMING*)_entry->NextIncommingEntry;
+					}
+				}
+				//node is not head
+				else {
+					_last->NextIncommingEntry = _entry->NextIncommingEntry;
+				}
+				break;
+			}
+			//save the last entry
+			if (_last != _entry)_last = (decltype(_last))_last->NextIncommingEntry;
+			_entry = (decltype(_entry))_entry->NextIncommingEntry;
+		} while (_entry != _last);
+		//free LoaderLock
+		LdrUnlockLoaderLock(0, Cookies);
+
+		//free it
+		LdrUnloadDll(ModuleEntry->DllBase);
+		NtFreeLdrpHeap(LdrEntry->DdagNode->Dependencies);
+
+		//lookup next dependent.
+		entry = (decltype(entry))entry->Dependent.NextDependentEntry;
+		LdrEntry->DdagNode->Dependencies = (_LDRP_CSLIST::_LDRP_CSLIST_DEPENDENT*)(entry == head ? nullptr : entry);
+	} while (entry != head);
+
+	return STATUS_SUCCESS;
+}
 static bool NTAPI NtInitializeLdrDataTableEntry(
 	OUT PLDR_DATA_TABLE_ENTRY LdrEntry,
 	IN DWORD dwFlags,
 	IN PVOID BaseAddress,
 	IN UNICODE_STRING &DllBaseName,
 	IN UNICODE_STRING &DllFullName) {
-	UNREFERENCED_PARAMETER(dwFlags);
 	RtlZeroMemory(LdrEntry, NtLdrDataTableEntrySize());
 	PIMAGE_NT_HEADERS headers = RtlImageNtHeader(BaseAddress);
 	if (!headers)return false;
@@ -303,7 +355,7 @@ static bool NTAPI NtInitializeLdrDataTableEntry(
 	case win8:
 	case win8_1: {
 		auto entry = (PLDR_DATA_TABLE_ENTRY_WIN8)LdrEntry;
-		
+		NtQuerySystemTime(&entry->LoadTime);
 		entry->OriginalBase = headers->OptionalHeader.ImageBase;
 		entry->BaseNameHashValue = LdrHashEntry(DllBaseName, false);
 		entry->LoadReason = LoadReasonDynamicLoad;
@@ -311,14 +363,17 @@ static bool NTAPI NtInitializeLdrDataTableEntry(
 		if (!(entry->DdagNode = (decltype(entry->DdagNode))NtAllocateLdrpHeap(sizeof(_LDR_DDAG_NODE))))return false;
 		//NtInitializeListEntry(&entry->NodeModuleLink);
 		//NtInitializeListEntry(&entry->DdagNode->Modules);
+		//NtInitializeSingleEntry(&entry->DdagNode->CondenseLink);
 		entry->NodeModuleLink.Flink = &entry->DdagNode->Modules;
 		entry->NodeModuleLink.Blink = &entry->DdagNode->Modules;
 		entry->DdagNode->Modules.Flink = &entry->NodeModuleLink;
 		entry->DdagNode->Modules.Blink = &entry->NodeModuleLink;
 		entry->DdagNode->State = LdrModulesReadyToRun;
-		entry->DdagNode->LoadCount = 0;
-		
-		NtInitializeSingleEntry(&entry->DdagNode->CondenseLink);
+		entry->DdagNode->LoadCount = 1;
+		entry->ImageDll = entry->LoadNotificationsSent = entry->EntryProcessed =
+			entry->InLegacyLists = entry->InIndexes = entry->ProcessAttachCalled = true;
+		entry->InExceptionTable = !(dwFlags & LOAD_FLAGS_NOT_ADD_INVERTED_FUNCTION);
+		FlagsProcessed = true;
 	}
 
 	case win7: {
@@ -344,6 +399,7 @@ static bool NTAPI NtInitializeLdrDataTableEntry(
 		LdrEntry->BaseDllName = DllBaseName;
 		LdrEntry->FullDllName = DllFullName;
 		LdrEntry->EntryPoint = (PVOID)((size_t)BaseAddress + headers->OptionalHeader.AddressOfEntryPoint);
+		LdrEntry->LoadCount = 1;
 		if (!FlagsProcessed) LdrEntry->Flags = LDRP_IMAGE_DLL | LDRP_ENTRY_INSERTED | LDRP_ENTRY_PROCESSED | LDRP_PROCESS_ATTACH_CALLED;
 		NtInitializeListEntry(&LdrEntry->HashLinks);
 		return true;
@@ -355,12 +411,12 @@ static bool NTAPI NtFreeLdrDataTableEntry(IN PLDR_DATA_TABLE_ENTRY LdrEntry) {
 	switch (NtWindowsVersion()) {
 	case win10:
 	case win10_1:
-	case win10_2: {
-		auto entry = (PLDR_DATA_TABLE_ENTRY_WIN10)LdrEntry;
-		NtFreeLdrpHeap(entry->DdagNode);
-	}
+	case win10_2:
 	case win8:
 	case win8_1: {
+		auto entry = (PLDR_DATA_TABLE_ENTRY_WIN8)LdrEntry;
+		NtFreeDependencies(entry);
+		NtFreeLdrpHeap(entry->DdagNode);
 		NtRemoveModuleBaseAddressIndexNode(LdrEntry);
 	}
 	case win7:
@@ -556,7 +612,7 @@ static NTSTATUS NTAPI NtMapDllMemory(IN HMEMORYMODULE ViewBase, IN DWORD dwFlags
 }
 
 NTSTATUS NTAPI NtLoadDllMemory(OUT HMEMORYMODULE* BaseAddress, IN LPVOID BufferAddress, IN size_t BufferSize) {
-	return NtLoadDllMemoryExW(BaseAddress, nullptr, 0, BufferAddress, BufferSize, nullptr, nullptr);
+	return NtLoadDllMemoryExW(BaseAddress, nullptr, LOAD_FLAGS_NOT_FAIL_IF_HANDLE_TLS, BufferAddress, BufferSize, nullptr, nullptr);
 }
 
 NTSTATUS NTAPI NtLoadDllMemoryExW(
@@ -570,6 +626,8 @@ NTSTATUS NTAPI NtLoadDllMemoryExW(
 	PMEMORYMODULE module = nullptr;
 	NTSTATUS status = STATUS_SUCCESS;
 	PLDR_DATA_TABLE_ENTRY ModuleEntry = nullptr;
+	PIMAGE_NT_HEADERS headers = nullptr;
+	UNREFERENCED_PARAMETER(BufferSize);
 
 	__try {
 		if (IsBadReadPtr(BufferAddress, BufferSize))status = STATUS_ACCESS_VIOLATION;
@@ -616,7 +674,7 @@ NTSTATUS NTAPI NtLoadDllMemoryExW(
 		}
 	}
 
-	if (!(*BaseAddress = MemoryLoadLibrary(BufferAddress, BufferSize))) {
+	if (!(*BaseAddress = MemoryLoadLibrary(BufferAddress))) {
 		switch (GetLastError()) {
 		case ERROR_BAD_EXE_FORMAT:
 			return STATUS_INVALID_IMAGE_FORMAT;
@@ -633,7 +691,9 @@ NTSTATUS NTAPI NtLoadDllMemoryExW(
 		return STATUS_INVALID_ADDRESS;
 	}
 	module->loadFromNtLoadDllMemory = true;
-	if (dwFlags & LOAD_FLAGS_NOT_MAP_DLL) return STATUS_SUCCESS;
+	headers = RtlImageNtHeader(*BaseAddress);
+	if (headers->OptionalHeader.DllCharacteristics & IMAGE_DLLCHARACTERISTICS_NO_SEH)dwFlags |= LOAD_FLAGS_NOT_ADD_INVERTED_FUNCTION;
+	if (dwFlags & LOAD_FLAGS_NOT_MAP_DLL) return status;
 
 	status = NtMapDllMemory(*BaseAddress, dwFlags, DllName, DllFullName, &ModuleEntry);
 	if (!NT_SUCCESS(status)) {
@@ -647,26 +707,34 @@ NTSTATUS NTAPI NtLoadDllMemoryExW(
 
 	if (!(dwFlags & LOAD_FLAGS_NOT_USE_REFERENCE_COUNT))module->UseReferenceCount = true;
 
-	if (dwFlags & LOAD_FLAGS_NOT_ADD_INVERTED_FUNCTION)return STATUS_SUCCESS;
-	status = RtlInsertInvertedFunctionTable((PVOID)module->codeBase, RtlImageNtHeader(*BaseAddress)->OptionalHeader.SizeOfImage);
-	if (!NT_SUCCESS(status)) {
-		NtUnloadDllMemory(*BaseAddress);
-		*BaseAddress = nullptr;
-		if (LdrEntry)*LdrEntry = nullptr;
-		return status;
-	}
-	module->InsertInvertedFunctionTableEntry = true;
-
-	if (dwFlags & LOAD_FLAGS_NOT_HANDLE_TLS)return STATUS_SUCCESS;
-	status = LdrpHandleTlsData(ModuleEntry);
-	if (!NT_SUCCESS(status)) {
-		NtUnloadDllMemory(*BaseAddress);
-		*BaseAddress = nullptr;
-		if (LdrEntry)*LdrEntry = nullptr;
-		return status;
+	if (!(dwFlags & LOAD_FLAGS_NOT_ADD_INVERTED_FUNCTION)) {
+		status = RtlInsertInvertedFunctionTable((PVOID)module->codeBase, headers->OptionalHeader.SizeOfImage);
+		if (!NT_SUCCESS(status)) {
+			NtUnloadDllMemory(*BaseAddress);
+			*BaseAddress = nullptr;
+			if (LdrEntry)*LdrEntry = nullptr;
+			return status;
+		}
+		module->InsertInvertedFunctionTableEntry = true;
 	}
 
-	return STATUS_SUCCESS;
+	if (!(dwFlags & LOAD_FLAGS_NOT_HANDLE_TLS)) {
+		status = LdrpHandleTlsData(ModuleEntry);
+		if (!NT_SUCCESS(status)) {
+			do {
+				if (dwFlags & LOAD_FLAGS_NOT_FAIL_IF_HANDLE_TLS) {
+					status = 0x7fffffff;
+					break;
+				}
+				NtUnloadDllMemory(*BaseAddress);
+				*BaseAddress = nullptr;
+				if (LdrEntry)*LdrEntry = nullptr;
+				return status;
+			} while (false);
+		}
+	}
+
+	return status;
 }
 
 NTSTATUS NTAPI NtLoadDllMemoryExA(
@@ -705,22 +773,28 @@ NTSTATUS NTAPI NtUnloadDllMemory(IN HMEMORYMODULE BaseAddress) {
 	NTSTATUS status = STATUS_SUCCESS;
 	PMEMORYMODULE module = MapMemoryModuleHandle(BaseAddress);
 
+	//Not a memory module loaded via NtLoadDllMemory
 	if (!module || !module->loadFromNtLoadDllMemory)return STATUS_INVALID_HANDLE;
+
+	//Mapping dll failed
+	if (module->loadFromNtLoadDllMemory && !module->MappedDll) {
+		module->underUnload = true;
+		MemoryFreeLibrary(BaseAddress);
+		return STATUS_SUCCESS;
+	}
+
 	ListHead = &NtCurrentPeb()->Ldr->InLoadOrderModuleList;
 	ListEntry = ListHead->Flink;
 	while (ListEntry != ListHead) {
 		CurEntry = CONTAINING_RECORD(ListEntry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
 		ListEntry = ListEntry->Flink;
-		/* Check if it's being unloaded */
-		if (!CurEntry->InMemoryOrderLinks.Flink) continue;
-		/* Check if name matches */
 		if (CurEntry->DllBase == BaseAddress) {
 			if (RtlImageNtHeader(BaseAddress)->OptionalHeader.SizeOfImage == CurEntry->SizeOfImage) {
 				if (module->UseReferenceCount) {
 					status = NtGetReferenceCount(CurEntry, &count);
 					if (!NT_SUCCESS(status))return status;
 				}
-				if (!count) {
+				if (!(count & ~1)) {
 					module->underUnload = true;
 					if (module->MappedDll) {
 						if (module->InsertInvertedFunctionTableEntry) {
@@ -774,7 +848,10 @@ static VOID NTAPI RtlpInsertInvertedFunctionTable(IN PRTL_INVERTED_FUNCTION_TABL
 	if (CurrentSize != InvertedTable->MaxCount) {
 		//if (need)_InterlockedIncrement(&InvertedTable->Epoch);
 		if (CurrentSize != 0) {
-			while (Index < CurrentSize)if (ImageBase < InvertedTable->Entries[Index].ImageBase)break;
+			while (Index < CurrentSize) {
+				if (ImageBase < InvertedTable->Entries[Index].ImageBase)break;
+				++Index;
+			}
 
 			if (Index != CurrentSize) {
 				RtlMoveMemory(&InvertedTable->Entries[Index + 1],
@@ -873,13 +950,13 @@ static VOID NTAPI RtlpRemoveInvertedFunctionTable(IN PRTL_INVERTED_FUNCTION_TABL
 #else
 			if (IsWin10) {
 				RtlMoveMemory(&InvertedTable->Entries[Index], &InvertedTable->Entries[Index + 1],
-					(CurrentSize - Index) * sizeof(PRTL_INVERTED_FUNCTION_TABLE_ENTRY));
+					(CurrentSize - Index) * sizeof(RTL_INVERTED_FUNCTION_TABLE_ENTRY));
 			}
 			else {
 				RtlMoveMemory(
 					Index ? &InvertedTable->Entries[Index - 1].NextEntrySEHandlerTableEncoded : (PVOID)&InvertedTable->NextEntrySEHandlerTableEncoded,
 					&InvertedTable->Entries[Index].NextEntrySEHandlerTableEncoded,
-					(CurrentSize - Index) * sizeof(PRTL_INVERTED_FUNCTION_TABLE_ENTRY));
+					(CurrentSize - Index) * sizeof(RTL_INVERTED_FUNCTION_TABLE_ENTRY));
 			}
 #endif
 		}

MemoryModule/NativeFunctionsInternal.h

@@ -81,7 +81,14 @@ struct _LDR_SERVICE_TAG_RECORD {
 };
 //0x8 bytes (sizeof)
 struct _LDRP_CSLIST {
-	_SINGLE_LIST_ENTRY* Tail;                                        //0x0
+	struct _LDRP_CSLIST_DEPENDENT {
+		_SINGLE_LIST_ENTRY* NextDependentEntry;                                        //0x0
+		struct _LDR_DDAG_NODE* DependentDdagNode;
+	}Dependent;
+	struct _LDRP_CSLIST_INCOMMING {
+		_SINGLE_LIST_ENTRY* NextIncommingEntry;
+		struct _LDR_DDAG_NODE* IncommingDdagNode;
+	}Incomming;
 };
 //0x4 bytes (sizeof)
 enum _LDR_DDAG_STATE {
@@ -108,8 +115,8 @@ struct _LDR_DDAG_NODE {
 	ULONG LoadCount;                                                        //0x18
 	ULONG LoadWhileUnloadingCount;                                          //0x1c
 	ULONG LowestLink;                                                       //0x20
-	_LDRP_CSLIST Dependencies;												//0x28
-	_LDRP_CSLIST IncomingDependencies;										//0x30
+	_LDRP_CSLIST::_LDRP_CSLIST_DEPENDENT* Dependencies;						//0x28
+	_LDRP_CSLIST::_LDRP_CSLIST_INCOMMING* IncomingDependencies;				//0x30
 	_LDR_DDAG_STATE State;													//0x38
 	_SINGLE_LIST_ENTRY CondenseLink;										//0x40
 	ULONG PreorderNumber;                                                   //0x48
@@ -398,6 +405,9 @@ NTSTATUS NTAPI NtLoadDllMemory(
 //Also, will be incompatible with Win32 API.
 #define LOAD_FLAGS_NOT_MAP_DLL						0x10000000
 
+//If this flag is specified, this routine will not fail even if the call to LdrpTlsData fails.
+#define LOAD_FLAGS_NOT_FAIL_IF_HANDLE_TLS			0x20000000
+
 //If this flag is specified, exception handling will not be supported.
 #define LOAD_FLAGS_NOT_ADD_INVERTED_FUNCTION		0x00000001