Skip to content

Commit 826747e

Browse files
bb107bb107
committed
Add MmpTlsFiber
1 parent 5fb7f59 commit 826747e

6 files changed

Lines changed: 216 additions & 81 deletions

File tree

‎MemoryModule/MemoryModule.vcxproj‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
<ClCompile Include="LoadDllMemoryApi.cpp" />
5050
<ClCompile Include="MemoryModule.cpp" />
5151
<ClCompile Include="MmpDotNet.cpp" />
52+
<ClCompile Include="MmpTlsFiber.cpp" />
5253
<ClCompile Include="MmpLdrpTls.cpp" />
5354
<ClCompile Include="MmpTls.cpp" />
5455
<ClCompile Include="Loader.cpp" />
@@ -109,12 +110,14 @@
109110
<ClInclude Include="LoaderPrivate.h" />
110111
<ClInclude Include="MemoryModule.h" />
111112
<ClInclude Include="MmpDotNet.h" />
113+
<ClInclude Include="MmpTlsFiber.h" />
112114
<ClInclude Include="MmpGlobalData.h" />
113115
<ClInclude Include="MmpTls.h" />
114116
<ClInclude Include="Loader.h" />
115117
<ClInclude Include="BaseAddressIndex.h" />
116118
<ClInclude Include="InvertedFunctionTable.h" />
117119
<ClInclude Include="LdrEntry.h" />
120+
<ClInclude Include="MmpTlsp.h" />
118121
<ClInclude Include="ReflectiveDLLInjection.h" />
119122
<ClInclude Include="ReflectiveLoader.h" />
120123
<ClInclude Include="stdafx.h" />

‎MemoryModule/MemoryModule.vcxproj.filters‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,9 @@
102102
<ClCompile Include="ReflectiveLoader.c">
103103
<Filter>Source Files\3rdparty\ReflectiveLoader</Filter>
104104
</ClCompile>
105+
<ClCompile Include="MmpTlsFiber.cpp">
106+
<Filter>Source Files</Filter>
107+
</ClCompile>
105108
</ItemGroup>
106109
<ItemGroup>
107110
<ClInclude Include="MemoryModule.h">
@@ -245,6 +248,12 @@
245248
<ClInclude Include="LoaderPrivate.h">
246249
<Filter>Header Files</Filter>
247250
</ClInclude>
251+
<ClInclude Include="MmpTlsFiber.h">
252+
<Filter>Header Files</Filter>
253+
</ClInclude>
254+
<ClInclude Include="MmpTlsp.h">
255+
<Filter>Header Files</Filter>
256+
</ClInclude>
248257
</ItemGroup>
249258
<ItemGroup>
250259
<None Include="..\README.md">

‎MemoryModule/MmpTls.cpp‎

Lines changed: 42 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -1,66 +1,13 @@
11
#include "stdafx.h"
22

33
#if (MMPP_USE_TLS)
4+
#include "MmpTlsp.h"
5+
#include "MmpTlsFiber.h"
6+
47
#include <cassert>
58
#include <algorithm>
69
#include <3rdparty/Detours/detours.h>
710

8-
//
9-
// ThreadLocalStoragePointer Tls indexs
10-
// [0, MMP_START_TLS_INDEX) Reserved for ntdll loader
11-
// [MMP_START_TLS_INDEX, MMP_MAXIMUM_TLS_INDEX) Reserved for MemoryModule
12-
//
13-
14-
#define MMP_START_TLS_INDEX 0x80 //128
15-
16-
#define MMP_MAXIMUM_TLS_INDEX 0x100 //256
17-
18-
#define MMP_TLSP_INDEX_BUFFER_SIZE (MMP_MAXIMUM_TLS_INDEX / 8) //32
19-
20-
#if (((MMP_START_TLS_INDEX | MMP_MAXIMUM_TLS_INDEX) & 7) || (MMP_START_TLS_INDEX >= MMP_MAXIMUM_TLS_INDEX))
21-
#error "MMP_START_TLS_INDEX must be smaller than MMP_MAXIMUM_TLS_INDEX, and both are 8-bit aligned."
22-
#endif
23-
24-
#define MmpAllocateTlsp() (RtlAllocateHeap(\
25-
RtlProcessHeap(),\
26-
HEAP_ZERO_MEMORY,\
27-
sizeof(PVOID)* MMP_MAXIMUM_TLS_INDEX\
28-
))
29-
30-
typedef struct _TLS_VECTOR {
31-
union
32-
{
33-
ULONG Length;
34-
HANDLE ThreadId;
35-
};
36-
37-
struct _TLS_VECTOR* PreviousDeferredTlsVector;
38-
PVOID ModuleTlsData[ANYSIZE_ARRAY];
39-
} TLS_VECTOR, * PTLS_VECTOR;
40-
41-
typedef struct _TLS_ENTRY {
42-
LIST_ENTRY TlsEntryLinks;
43-
IMAGE_TLS_DIRECTORY TlsDirectory;
44-
PLDR_DATA_TABLE_ENTRY ModuleEntry;
45-
} TLS_ENTRY, * PTLS_ENTRY;
46-
47-
typedef struct _MMP_TLSP_RECORD {
48-
49-
LIST_ENTRY InMmpThreadLocalStoragePointer;
50-
51-
HANDLE UniqueThread;
52-
53-
// PEB->ThreadLocalStoragePointer allocated by ntdll!Ldr
54-
PVOID* TlspLdrBlock;
55-
56-
// PEB->ThreadLocalStoragePointer allocated by MemoryModulePP
57-
PVOID* TlspMmpBlock;
58-
}MMP_TLSP_RECORD, * PMMP_TLSP_RECORD;
59-
60-
typedef struct _THREAD_CONTEXT {
61-
PTHREAD_START_ROUTINE ThreadStartRoutine;
62-
LPVOID ThreadParameter;
63-
}THREAD_CONTEXT, * PTHREAD_CONTEXT;
6411

6512
PVOID NTAPI MmpQuerySystemInformation(
6613
_In_ SYSTEM_INFORMATION_CLASS SystemInformationClass,
@@ -217,6 +164,8 @@ DWORD NTAPI MmpUserThreadStart(LPVOID lpThreadParameter) {
217164
record->TlspMmpBlock = (PVOID*)MmpAllocateTlsp();
218165
record->UniqueThread = NtCurrentThreadId();
219166
if (record->TlspMmpBlock) {
167+
record->TlspMmpBlock = ((PTLS_VECTOR)record->TlspMmpBlock)->ModuleTlsData;
168+
220169
auto size = CONTAINING_RECORD(record->TlspLdrBlock, TLS_VECTOR, ModuleTlsData)->Length;
221170
if ((HANDLE)(ULONG_PTR)size != NtCurrentThreadId()) {
222171
RtlCopyMemory(
@@ -327,6 +276,7 @@ VOID NTAPI HookLdrShutdownThread(VOID) {
327276

328277
PLIST_ENTRY entry;
329278
PMMP_TLSP_RECORD record = nullptr;
279+
BOOL postpone = IsThreadAFiber();
330280

331281
//
332282
// Find our tlsp record
@@ -335,48 +285,54 @@ VOID NTAPI HookLdrShutdownThread(VOID) {
335285

336286
record = MmpFindTlspRecordLockHeld();
337287
if (record) {
338-
339-
//
340-
// Restore tlsp
341-
//
342-
343-
NtCurrentTeb()->ThreadLocalStoragePointer = record->TlspLdrBlock;
344288
RemoveEntryList(&record->InMmpThreadLocalStoragePointer);
345-
346289
--MmpGlobalDataPtr->MmpTls->MmpActiveThreadCount;
347290
}
348291

349292
assert(0 < (int)MmpGlobalDataPtr->MmpTls->MmpActiveThreadCount);
350293

351294
LeaveCriticalSection(&MmpGlobalDataPtr->MmpTls->MmpTlspLock);
352295

353-
//
354-
// Free MemoryModule Tls data
355-
//
356-
RtlAcquireSRWLockExclusive(&MmpGlobalDataPtr->MmpTls->MmpTlsListLock);
357-
358296
if (record) {
359-
auto TlspMmpBlock = (PVOID*)record->TlspMmpBlock;
360-
entry = MmpGlobalDataPtr->MmpTls->MmpTlsList.Flink;
361-
while (entry != &MmpGlobalDataPtr->MmpTls->MmpTlsList) {
297+
if (postpone) {
298+
299+
//
300+
// Free MemoryModule Tls data after terminated
301+
//
302+
303+
MmpQueuePostponedTls(record);
304+
}
305+
else {
362306

363-
auto p = CONTAINING_RECORD(entry, TLS_ENTRY, TlsEntryLinks);
364-
RtlFreeHeap(RtlProcessHeap(), 0, TlspMmpBlock[p->TlsDirectory.Characteristics]);
307+
//
308+
// Free MemoryModule Tls data
309+
//
365310

366-
entry = entry->Flink;
367-
}
311+
RtlAcquireSRWLockExclusive(&MmpGlobalDataPtr->MmpTls->MmpTlsListLock);
312+
313+
auto TlspMmpBlock = (PVOID*)record->TlspMmpBlock;
314+
entry = MmpGlobalDataPtr->MmpTls->MmpTlsList.Flink;
315+
while (entry != &MmpGlobalDataPtr->MmpTls->MmpTlsList) {
316+
317+
auto p = CONTAINING_RECORD(entry, TLS_ENTRY, TlsEntryLinks);
318+
RtlFreeHeap(RtlProcessHeap(), 0, TlspMmpBlock[p->TlsDirectory.Characteristics]);
319+
TlspMmpBlock[p->TlsDirectory.Characteristics] = nullptr;
368320

369-
RtlFreeHeap(RtlProcessHeap(), 0, TlspMmpBlock);
370-
RtlFreeHeap(RtlProcessHeap(), 0, record);
321+
entry = entry->Flink;
322+
}
323+
324+
RtlFreeHeap(RtlProcessHeap(), 0, CONTAINING_RECORD(record->TlspLdrBlock, TLS_VECTOR, TLS_VECTOR::ModuleTlsData));
325+
RtlFreeHeap(RtlProcessHeap(), 0, record);
326+
327+
RtlReleaseSRWLockExclusive(&MmpGlobalDataPtr->MmpTls->MmpTlsListLock);
328+
}
371329
}
372330
else {
373331
if (MmpGlobalDataPtr->MmpTls->MmpTlsList.Flink != &MmpGlobalDataPtr->MmpTls->MmpTlsList) {
374332
assert(NtCurrentTeb()->ThreadLocalStoragePointer == nullptr);
375333
}
376334
}
377335

378-
RtlReleaseSRWLockExclusive(&MmpGlobalDataPtr->MmpTls->MmpTlsListLock);
379-
380336
//
381337
// Call the original function
382338
//
@@ -405,9 +361,12 @@ BOOL NTAPI PreHookNtSetInformationProcess() {
405361
for (DWORD i = 0; i < CurrentThreadCount; ++i) {
406362
auto& current = ProcessTlsInformation->ThreadData[i];
407363
current.TlsVector = (PVOID*)MmpAllocateTlsp();
408-
if (!current.TlsVector) {
364+
if (current.TlsVector) {
365+
current.TlsVector = ((PTLS_VECTOR)current.TlsVector)->ModuleTlsData;
366+
}
367+
else {
409368
for (DWORD j = 0; j < i; ++j) {
410-
RtlFreeHeap(RtlProcessHeap(), 0, ProcessTlsInformation->ThreadData[j].TlsVector);
369+
RtlFreeHeap(RtlProcessHeap(), 0, CONTAINING_RECORD(ProcessTlsInformation->ThreadData[j].TlsVector, TLS_VECTOR, TLS_VECTOR::ModuleTlsData));
411370
}
412371

413372
success = FALSE;
@@ -848,6 +807,8 @@ BOOL NTAPI MmpTlsInitialize() {
848807
DetourAttach((PVOID*)&MmpGlobalDataPtr->MmpTls->Hooks.OriginRtlUserThreadStart, HookRtlUserThreadStart);
849808
DetourTransactionCommit();
850809

810+
MmpTlsFiberInitialize();
811+
851812
return TRUE;
852813
}
853814

‎MemoryModule/MmpTlsFiber.cpp‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
#include "stdafx.h"
2+
#include "MmpTlsp.h"
3+
#include "MmpTlsFiber.h"
4+
5+
#include <vector>
6+
7+
typedef struct _MMP_POSTPONED_TLS {
8+
9+
HANDLE hThread;
10+
PMMP_TLSP_RECORD lpTlsRecord;
11+
PTLS_VECTOR lpOldTlsVector;
12+
13+
}MMP_POSTPONED_TLS, * PMMP_POSTPONED_TLS;
14+
15+
std::vector<MMP_POSTPONED_TLS>MmpPostponedTlsList;
16+
17+
HANDLE MmpPostponedTlsEvent;
18+
CRITICAL_SECTION MmpPostponedTlsLock;
19+
20+
DWORD WINAPI MmpReleasePostponedTlsWorker(PVOID) {
21+
22+
DWORD code;
23+
24+
while (true) {
25+
WaitForSingleObject(MmpPostponedTlsEvent, INFINITE);
26+
27+
EnterCriticalSection(&MmpPostponedTlsLock);
28+
29+
auto iter = MmpPostponedTlsList.begin();
30+
31+
while (iter != MmpPostponedTlsList.end()) {
32+
const auto& item = *iter;
33+
GetExitCodeThread(item.hThread, &code);
34+
35+
if (code == STILL_ACTIVE) {
36+
++iter;
37+
}
38+
else {
39+
40+
RtlAcquireSRWLockExclusive(&MmpGlobalDataPtr->MmpTls->MmpTlsListLock);
41+
42+
auto TlspMmpBlock = (PVOID*)item.lpOldTlsVector->ModuleTlsData;
43+
auto entry = MmpGlobalDataPtr->MmpTls->MmpTlsList.Flink;
44+
while (entry != &MmpGlobalDataPtr->MmpTls->MmpTlsList) {
45+
46+
auto p = CONTAINING_RECORD(entry, TLS_ENTRY, TlsEntryLinks);
47+
RtlFreeHeap(RtlProcessHeap(), 0, TlspMmpBlock[p->TlsDirectory.Characteristics]);
48+
49+
entry = entry->Flink;
50+
}
51+
52+
RtlFreeHeap(RtlProcessHeap(), 0, CONTAINING_RECORD(item.lpTlsRecord->TlspLdrBlock, TLS_VECTOR, TLS_VECTOR::ModuleTlsData));
53+
RtlFreeHeap(RtlProcessHeap(), 0, item.lpTlsRecord);
54+
RtlFreeHeap(RtlProcessHeap(), 0, item.lpOldTlsVector);
55+
56+
RtlReleaseSRWLockExclusive(&MmpGlobalDataPtr->MmpTls->MmpTlsListLock);
57+
58+
CloseHandle(item.hThread);
59+
iter = MmpPostponedTlsList.erase(iter);
60+
}
61+
62+
}
63+
64+
LeaveCriticalSection(&MmpPostponedTlsLock);
65+
}
66+
67+
return 0;
68+
}
69+
70+
VOID WINAPI MmpQueuePostponedTls(PMMP_TLSP_RECORD record) {
71+
MMP_POSTPONED_TLS item;
72+
73+
item.hThread = OpenThread(
74+
THREAD_QUERY_INFORMATION,
75+
FALSE,
76+
(DWORD)(ULONG_PTR)NtCurrentThreadId()
77+
);
78+
79+
item.lpOldTlsVector = MmpAllocateTlsp();
80+
81+
item.lpTlsRecord = record;
82+
83+
RtlCopyMemory(
84+
item.lpOldTlsVector->ModuleTlsData,
85+
record->TlspMmpBlock,
86+
sizeof(PVOID) * MMP_MAXIMUM_TLS_INDEX
87+
);
88+
89+
EnterCriticalSection(&MmpPostponedTlsLock);
90+
91+
MmpPostponedTlsList.push_back(item);
92+
SetEvent(MmpPostponedTlsEvent);
93+
94+
LeaveCriticalSection(&MmpPostponedTlsLock);
95+
}
96+
97+
VOID MmpTlsFiberInitialize() {
98+
InitializeCriticalSection(&MmpPostponedTlsLock);
99+
MmpPostponedTlsEvent = CreateEvent(nullptr, FALSE, FALSE, nullptr);
100+
}

‎MemoryModule/MmpTlsFiber.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
#pragma once
2+
3+
VOID WINAPI MmpQueuePostponedTls(PMMP_TLSP_RECORD record);
4+
VOID MmpTlsFiberInitialize();

‎MemoryModule/MmpTlsp.h‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
#pragma once
2+
3+
//
4+
// ThreadLocalStoragePointer Tls indexs
5+
// [0, MMP_START_TLS_INDEX) Reserved for ntdll loader
6+
// [MMP_START_TLS_INDEX, MMP_MAXIMUM_TLS_INDEX) Reserved for MemoryModule
7+
//
8+
9+
#define MMP_START_TLS_INDEX 0x80 //128
10+
11+
#define MMP_MAXIMUM_TLS_INDEX 0x100 //256
12+
13+
#define MMP_TLSP_INDEX_BUFFER_SIZE (MMP_MAXIMUM_TLS_INDEX / 8) //32
14+
15+
#if (((MMP_START_TLS_INDEX | MMP_MAXIMUM_TLS_INDEX) & 7) || (MMP_START_TLS_INDEX >= MMP_MAXIMUM_TLS_INDEX))
16+
#error "MMP_START_TLS_INDEX must be smaller than MMP_MAXIMUM_TLS_INDEX, and both are 8-bit aligned."
17+
#endif
18+
19+
#define MmpAllocateTlsp() (PTLS_VECTOR)(RtlAllocateHeap(\
20+
RtlProcessHeap(),\
21+
HEAP_ZERO_MEMORY,\
22+
sizeof(TLS_VECTOR) + sizeof(PVOID)* MMP_MAXIMUM_TLS_INDEX\
23+
))
24+
25+
typedef struct _TLS_VECTOR {
26+
union
27+
{
28+
ULONG Length;
29+
HANDLE ThreadId;
30+
};
31+
32+
struct _TLS_VECTOR* PreviousDeferredTlsVector;
33+
PVOID ModuleTlsData[ANYSIZE_ARRAY];
34+
} TLS_VECTOR, * PTLS_VECTOR;
35+
36+
typedef struct _TLS_ENTRY {
37+
LIST_ENTRY TlsEntryLinks;
38+
IMAGE_TLS_DIRECTORY TlsDirectory;
39+
PLDR_DATA_TABLE_ENTRY ModuleEntry;
40+
} TLS_ENTRY, * PTLS_ENTRY;
41+
42+
typedef struct _MMP_TLSP_RECORD {
43+
44+
LIST_ENTRY InMmpThreadLocalStoragePointer;
45+
46+
HANDLE UniqueThread;
47+
48+
// PEB->ThreadLocalStoragePointer allocated by ntdll!Ldr
49+
PVOID* TlspLdrBlock;
50+
51+
// PEB->ThreadLocalStoragePointer allocated by MemoryModulePP
52+
PVOID* TlspMmpBlock;
53+
}MMP_TLSP_RECORD, * PMMP_TLSP_RECORD;
54+
55+
typedef struct _THREAD_CONTEXT {
56+
PTHREAD_START_ROUTINE ThreadStartRoutine;
57+
LPVOID ThreadParameter;
58+
}THREAD_CONTEXT, * PTHREAD_CONTEXT;

0 commit comments

Comments
 (0)