Add MmpTlsFiber

MemoryModule/MmpTlsFiber.cpp

@@ -0,0 +1,100 @@
+#include "stdafx.h"
+#include "MmpTlsp.h"
+#include "MmpTlsFiber.h"
+
+#include <vector>
+
+typedef struct _MMP_POSTPONED_TLS {
+
+	HANDLE hThread;
+	PMMP_TLSP_RECORD lpTlsRecord;
+	PTLS_VECTOR lpOldTlsVector;
+
+}MMP_POSTPONED_TLS, * PMMP_POSTPONED_TLS;
+
+std::vector<MMP_POSTPONED_TLS>MmpPostponedTlsList;
+
+HANDLE MmpPostponedTlsEvent;
+CRITICAL_SECTION MmpPostponedTlsLock;
+
+DWORD WINAPI MmpReleasePostponedTlsWorker(PVOID) {
+
+	DWORD code;
+
+	while (true) {
+		WaitForSingleObject(MmpPostponedTlsEvent, INFINITE);
+
+		EnterCriticalSection(&MmpPostponedTlsLock);
+
+		auto iter = MmpPostponedTlsList.begin();
+
+		while (iter != MmpPostponedTlsList.end()) {
+			const auto& item = *iter;
+			GetExitCodeThread(item.hThread, &code);
+
+			if (code == STILL_ACTIVE) {
+				++iter;
+			}
+			else {
+
+				RtlAcquireSRWLockExclusive(&MmpGlobalDataPtr->MmpTls->MmpTlsListLock);
+
+				auto TlspMmpBlock = (PVOID*)item.lpOldTlsVector->ModuleTlsData;
+				auto entry = MmpGlobalDataPtr->MmpTls->MmpTlsList.Flink;
+				while (entry != &MmpGlobalDataPtr->MmpTls->MmpTlsList) {
+
+					auto p = CONTAINING_RECORD(entry, TLS_ENTRY, TlsEntryLinks);
+					RtlFreeHeap(RtlProcessHeap(), 0, TlspMmpBlock[p->TlsDirectory.Characteristics]);
+
+					entry = entry->Flink;
+				}
+
+				RtlFreeHeap(RtlProcessHeap(), 0, CONTAINING_RECORD(item.lpTlsRecord->TlspLdrBlock, TLS_VECTOR, TLS_VECTOR::ModuleTlsData));
+				RtlFreeHeap(RtlProcessHeap(), 0, item.lpTlsRecord);
+				RtlFreeHeap(RtlProcessHeap(), 0, item.lpOldTlsVector);
+
+				RtlReleaseSRWLockExclusive(&MmpGlobalDataPtr->MmpTls->MmpTlsListLock);
+
+				CloseHandle(item.hThread);
+				iter = MmpPostponedTlsList.erase(iter);
+			}
+
+		}
+
+		LeaveCriticalSection(&MmpPostponedTlsLock);
+	}
+
+	return 0;
+}
+
+VOID WINAPI MmpQueuePostponedTls(PMMP_TLSP_RECORD record) {
+	MMP_POSTPONED_TLS item;
+
+	item.hThread = OpenThread(
+		THREAD_QUERY_INFORMATION,
+		FALSE,
+		(DWORD)(ULONG_PTR)NtCurrentThreadId()
+	);
+
+	item.lpOldTlsVector = MmpAllocateTlsp();
+
+	item.lpTlsRecord = record;
+
+	RtlCopyMemory(
+		item.lpOldTlsVector->ModuleTlsData,
+		record->TlspMmpBlock,
+		sizeof(PVOID) * MMP_MAXIMUM_TLS_INDEX
+	);
+
+	EnterCriticalSection(&MmpPostponedTlsLock);
+	
+	MmpPostponedTlsList.push_back(item);
+	SetEvent(MmpPostponedTlsEvent);
+	
+	LeaveCriticalSection(&MmpPostponedTlsLock);
+}
+
+VOID MmpTlsFiberInitialize() {
+	InitializeCriticalSection(&MmpPostponedTlsLock);
+	MmpPostponedTlsEvent = CreateEvent(nullptr, FALSE, FALSE, nullptr);
+}