From 826747e27cf18422cee41d99f6a62729cf507f19 Mon Sep 17 00:00:00 2001
From: Boring <1079299053@qq.com>
Date: Sun, 24 Sep 2023 20:34:19 +0800
Subject: [PATCH] Add MmpTlsFiber
---
MemoryModule/MemoryModule.vcxproj | 3 +
MemoryModule/MemoryModule.vcxproj.filters | 9 ++
MemoryModule/MmpTls.cpp | 125 ++++++++--------------
MemoryModule/MmpTlsFiber.cpp | 100 +++++++++++++++++
MemoryModule/MmpTlsFiber.h | 4 +
MemoryModule/MmpTlsp.h | 58 ++++++++++
6 files changed, 217 insertions(+), 82 deletions(-)
create mode 100644 MemoryModule/MmpTlsFiber.cpp
create mode 100644 MemoryModule/MmpTlsFiber.h
create mode 100644 MemoryModule/MmpTlsp.h
diff --git a/MemoryModule/MemoryModule.vcxproj b/MemoryModule/MemoryModule.vcxproj
index 7d73247..22bee5d 100644
--- a/MemoryModule/MemoryModule.vcxproj
+++ b/MemoryModule/MemoryModule.vcxproj
@@ -49,6 +49,7 @@
+
@@ -109,12 +110,14 @@
+
+
diff --git a/MemoryModule/MemoryModule.vcxproj.filters b/MemoryModule/MemoryModule.vcxproj.filters
index 2fba3b2..4364a80 100644
--- a/MemoryModule/MemoryModule.vcxproj.filters
+++ b/MemoryModule/MemoryModule.vcxproj.filters
@@ -102,6 +102,9 @@
Source Files\3rdparty\ReflectiveLoader
+
+ Source Files
+
@@ -245,6 +248,12 @@
Header Files
+
+ Header Files
+
+
+ Header Files
+
diff --git a/MemoryModule/MmpTls.cpp b/MemoryModule/MmpTls.cpp
index f574816..4fb6220 100644
--- a/MemoryModule/MmpTls.cpp
+++ b/MemoryModule/MmpTls.cpp
@@ -1,66 +1,13 @@
#include "stdafx.h"
#if (MMPP_USE_TLS)
+#include "MmpTlsp.h"
+#include "MmpTlsFiber.h"
+
#include
#include
#include <3rdparty/Detours/detours.h>
-//
-// ThreadLocalStoragePointer Tls indexs
-// [0, MMP_START_TLS_INDEX) Reserved for ntdll loader
-// [MMP_START_TLS_INDEX, MMP_MAXIMUM_TLS_INDEX) Reserved for MemoryModule
-//
-
-#define MMP_START_TLS_INDEX 0x80 //128
-
-#define MMP_MAXIMUM_TLS_INDEX 0x100 //256
-
-#define MMP_TLSP_INDEX_BUFFER_SIZE (MMP_MAXIMUM_TLS_INDEX / 8) //32
-
-#if (((MMP_START_TLS_INDEX | MMP_MAXIMUM_TLS_INDEX) & 7) || (MMP_START_TLS_INDEX >= MMP_MAXIMUM_TLS_INDEX))
-#error "MMP_START_TLS_INDEX must be smaller than MMP_MAXIMUM_TLS_INDEX, and both are 8-bit aligned."
-#endif
-
-#define MmpAllocateTlsp() (RtlAllocateHeap(\
- RtlProcessHeap(),\
- HEAP_ZERO_MEMORY,\
- sizeof(PVOID)* MMP_MAXIMUM_TLS_INDEX\
- ))
-
-typedef struct _TLS_VECTOR {
- union
- {
- ULONG Length;
- HANDLE ThreadId;
- };
-
- struct _TLS_VECTOR* PreviousDeferredTlsVector;
- PVOID ModuleTlsData[ANYSIZE_ARRAY];
-} TLS_VECTOR, * PTLS_VECTOR;
-
-typedef struct _TLS_ENTRY {
- LIST_ENTRY TlsEntryLinks;
- IMAGE_TLS_DIRECTORY TlsDirectory;
- PLDR_DATA_TABLE_ENTRY ModuleEntry;
-} TLS_ENTRY, * PTLS_ENTRY;
-
-typedef struct _MMP_TLSP_RECORD {
-
- LIST_ENTRY InMmpThreadLocalStoragePointer;
-
- HANDLE UniqueThread;
-
- // PEB->ThreadLocalStoragePointer allocated by ntdll!Ldr
- PVOID* TlspLdrBlock;
-
- // PEB->ThreadLocalStoragePointer allocated by MemoryModulePP
- PVOID* TlspMmpBlock;
-}MMP_TLSP_RECORD, * PMMP_TLSP_RECORD;
-
-typedef struct _THREAD_CONTEXT {
- PTHREAD_START_ROUTINE ThreadStartRoutine;
- LPVOID ThreadParameter;
-}THREAD_CONTEXT, * PTHREAD_CONTEXT;
PVOID NTAPI MmpQuerySystemInformation(
_In_ SYSTEM_INFORMATION_CLASS SystemInformationClass,
@@ -217,6 +164,8 @@ DWORD NTAPI MmpUserThreadStart(LPVOID lpThreadParameter) {
record->TlspMmpBlock = (PVOID*)MmpAllocateTlsp();
record->UniqueThread = NtCurrentThreadId();
if (record->TlspMmpBlock) {
+ record->TlspMmpBlock = ((PTLS_VECTOR)record->TlspMmpBlock)->ModuleTlsData;
+
auto size = CONTAINING_RECORD(record->TlspLdrBlock, TLS_VECTOR, ModuleTlsData)->Length;
if ((HANDLE)(ULONG_PTR)size != NtCurrentThreadId()) {
RtlCopyMemory(
@@ -327,6 +276,7 @@ VOID NTAPI HookLdrShutdownThread(VOID) {
PLIST_ENTRY entry;
PMMP_TLSP_RECORD record = nullptr;
+ BOOL postpone = IsThreadAFiber();
//
// Find our tlsp record
@@ -335,14 +285,7 @@ VOID NTAPI HookLdrShutdownThread(VOID) {
record = MmpFindTlspRecordLockHeld();
if (record) {
-
- //
- // Restore tlsp
- //
-
- NtCurrentTeb()->ThreadLocalStoragePointer = record->TlspLdrBlock;
RemoveEntryList(&record->InMmpThreadLocalStoragePointer);
-
--MmpGlobalDataPtr->MmpTls->MmpActiveThreadCount;
}
@@ -350,24 +293,39 @@ VOID NTAPI HookLdrShutdownThread(VOID) {
LeaveCriticalSection(&MmpGlobalDataPtr->MmpTls->MmpTlspLock);
- //
- // Free MemoryModule Tls data
- //
- RtlAcquireSRWLockExclusive(&MmpGlobalDataPtr->MmpTls->MmpTlsListLock);
-
if (record) {
- auto TlspMmpBlock = (PVOID*)record->TlspMmpBlock;
- 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;
+ if (postpone) {
+
+ //
+ // Free MemoryModule Tls data after terminated
+ //
+
+ MmpQueuePostponedTls(record);
}
+ else {
- RtlFreeHeap(RtlProcessHeap(), 0, TlspMmpBlock);
- RtlFreeHeap(RtlProcessHeap(), 0, record);
+ //
+ // Free MemoryModule Tls data
+ //
+
+ RtlAcquireSRWLockExclusive(&MmpGlobalDataPtr->MmpTls->MmpTlsListLock);
+
+ auto TlspMmpBlock = (PVOID*)record->TlspMmpBlock;
+ 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]);
+ TlspMmpBlock[p->TlsDirectory.Characteristics] = nullptr;
+
+ entry = entry->Flink;
+ }
+
+ RtlFreeHeap(RtlProcessHeap(), 0, CONTAINING_RECORD(record->TlspLdrBlock, TLS_VECTOR, TLS_VECTOR::ModuleTlsData));
+ RtlFreeHeap(RtlProcessHeap(), 0, record);
+
+ RtlReleaseSRWLockExclusive(&MmpGlobalDataPtr->MmpTls->MmpTlsListLock);
+ }
}
else {
if (MmpGlobalDataPtr->MmpTls->MmpTlsList.Flink != &MmpGlobalDataPtr->MmpTls->MmpTlsList) {
@@ -375,8 +333,6 @@ VOID NTAPI HookLdrShutdownThread(VOID) {
}
}
- RtlReleaseSRWLockExclusive(&MmpGlobalDataPtr->MmpTls->MmpTlsListLock);
-
//
// Call the original function
//
@@ -405,9 +361,12 @@ BOOL NTAPI PreHookNtSetInformationProcess() {
for (DWORD i = 0; i < CurrentThreadCount; ++i) {
auto& current = ProcessTlsInformation->ThreadData[i];
current.TlsVector = (PVOID*)MmpAllocateTlsp();
- if (!current.TlsVector) {
+ if (current.TlsVector) {
+ current.TlsVector = ((PTLS_VECTOR)current.TlsVector)->ModuleTlsData;
+ }
+ else {
for (DWORD j = 0; j < i; ++j) {
- RtlFreeHeap(RtlProcessHeap(), 0, ProcessTlsInformation->ThreadData[j].TlsVector);
+ RtlFreeHeap(RtlProcessHeap(), 0, CONTAINING_RECORD(ProcessTlsInformation->ThreadData[j].TlsVector, TLS_VECTOR, TLS_VECTOR::ModuleTlsData));
}
success = FALSE;
@@ -848,6 +807,8 @@ BOOL NTAPI MmpTlsInitialize() {
DetourAttach((PVOID*)&MmpGlobalDataPtr->MmpTls->Hooks.OriginRtlUserThreadStart, HookRtlUserThreadStart);
DetourTransactionCommit();
+ MmpTlsFiberInitialize();
+
return TRUE;
}
diff --git a/MemoryModule/MmpTlsFiber.cpp b/MemoryModule/MmpTlsFiber.cpp
new file mode 100644
index 0000000..3366542
--- /dev/null
+++ b/MemoryModule/MmpTlsFiber.cpp
@@ -0,0 +1,100 @@
+#include "stdafx.h"
+#include "MmpTlsp.h"
+#include "MmpTlsFiber.h"
+
+#include
+
+typedef struct _MMP_POSTPONED_TLS {
+
+ HANDLE hThread;
+ PMMP_TLSP_RECORD lpTlsRecord;
+ PTLS_VECTOR lpOldTlsVector;
+
+}MMP_POSTPONED_TLS, * PMMP_POSTPONED_TLS;
+
+std::vectorMmpPostponedTlsList;
+
+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);
+}
diff --git a/MemoryModule/MmpTlsFiber.h b/MemoryModule/MmpTlsFiber.h
new file mode 100644
index 0000000..954990a
--- /dev/null
+++ b/MemoryModule/MmpTlsFiber.h
@@ -0,0 +1,4 @@
+#pragma once
+
+VOID WINAPI MmpQueuePostponedTls(PMMP_TLSP_RECORD record);
+VOID MmpTlsFiberInitialize();
diff --git a/MemoryModule/MmpTlsp.h b/MemoryModule/MmpTlsp.h
new file mode 100644
index 0000000..6a05ba8
--- /dev/null
+++ b/MemoryModule/MmpTlsp.h
@@ -0,0 +1,58 @@
+#pragma once
+
+//
+// ThreadLocalStoragePointer Tls indexs
+// [0, MMP_START_TLS_INDEX) Reserved for ntdll loader
+// [MMP_START_TLS_INDEX, MMP_MAXIMUM_TLS_INDEX) Reserved for MemoryModule
+//
+
+#define MMP_START_TLS_INDEX 0x80 //128
+
+#define MMP_MAXIMUM_TLS_INDEX 0x100 //256
+
+#define MMP_TLSP_INDEX_BUFFER_SIZE (MMP_MAXIMUM_TLS_INDEX / 8) //32
+
+#if (((MMP_START_TLS_INDEX | MMP_MAXIMUM_TLS_INDEX) & 7) || (MMP_START_TLS_INDEX >= MMP_MAXIMUM_TLS_INDEX))
+#error "MMP_START_TLS_INDEX must be smaller than MMP_MAXIMUM_TLS_INDEX, and both are 8-bit aligned."
+#endif
+
+#define MmpAllocateTlsp() (PTLS_VECTOR)(RtlAllocateHeap(\
+ RtlProcessHeap(),\
+ HEAP_ZERO_MEMORY,\
+ sizeof(TLS_VECTOR) + sizeof(PVOID)* MMP_MAXIMUM_TLS_INDEX\
+ ))
+
+typedef struct _TLS_VECTOR {
+ union
+ {
+ ULONG Length;
+ HANDLE ThreadId;
+ };
+
+ struct _TLS_VECTOR* PreviousDeferredTlsVector;
+ PVOID ModuleTlsData[ANYSIZE_ARRAY];
+} TLS_VECTOR, * PTLS_VECTOR;
+
+typedef struct _TLS_ENTRY {
+ LIST_ENTRY TlsEntryLinks;
+ IMAGE_TLS_DIRECTORY TlsDirectory;
+ PLDR_DATA_TABLE_ENTRY ModuleEntry;
+} TLS_ENTRY, * PTLS_ENTRY;
+
+typedef struct _MMP_TLSP_RECORD {
+
+ LIST_ENTRY InMmpThreadLocalStoragePointer;
+
+ HANDLE UniqueThread;
+
+ // PEB->ThreadLocalStoragePointer allocated by ntdll!Ldr
+ PVOID* TlspLdrBlock;
+
+ // PEB->ThreadLocalStoragePointer allocated by MemoryModulePP
+ PVOID* TlspMmpBlock;
+}MMP_TLSP_RECORD, * PMMP_TLSP_RECORD;
+
+typedef struct _THREAD_CONTEXT {
+ PTHREAD_START_ROUTINE ThreadStartRoutine;
+ LPVOID ThreadParameter;
+}THREAD_CONTEXT, * PTHREAD_CONTEXT;