admin/bb107-MemoryModulePP
1
0
Fork 0
mirror of https://github.com/bb107/MemoryModulePP synced 2026-06-08 13:15:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add MmpTlsFiber

Browse Source
This commit is contained in:
Boring
2023-09-24 20:34:19 +08:00
parent 5fb7f591d2
commit 826747e27c
6 changed files with 217 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
MemoryModule/MemoryModule.vcxproj
+3
View File
@@ -49,6 +49,7 @@
<ClCompile Include="LoadDllMemoryApi.cpp" />
<ClCompile Include="MemoryModule.cpp" />
<ClCompile Include="MmpDotNet.cpp" />
<ClCompile Include="MmpTlsFiber.cpp" />
<ClCompile Include="MmpLdrpTls.cpp" />
<ClCompile Include="MmpTls.cpp" />
<ClCompile Include="Loader.cpp" />
@@ -109,12 +110,14 @@
<ClInclude Include="LoaderPrivate.h" />
<ClInclude Include="MemoryModule.h" />
<ClInclude Include="MmpDotNet.h" />
<ClInclude Include="MmpTlsFiber.h" />
<ClInclude Include="MmpGlobalData.h" />
<ClInclude Include="MmpTls.h" />
<ClInclude Include="Loader.h" />
<ClInclude Include="BaseAddressIndex.h" />
<ClInclude Include="InvertedFunctionTable.h" />
<ClInclude Include="LdrEntry.h" />
<ClInclude Include="MmpTlsp.h" />
<ClInclude Include="ReflectiveDLLInjection.h" />
<ClInclude Include="ReflectiveLoader.h" />
<ClInclude Include="stdafx.h" />
MemoryModule/MemoryModule.vcxproj.filters
+9
View File
@@ -102,6 +102,9 @@
<ClCompile Include="ReflectiveLoader.c">
<Filter>Source Files\3rdparty\ReflectiveLoader</Filter>
</ClCompile>
<ClCompile Include="MmpTlsFiber.cpp">
<Filter>Source Files</Filter>
</ClCompile>
</ItemGroup>
<ItemGroup>
<ClInclude Include="MemoryModule.h">
@@ -245,6 +248,12 @@
<ClInclude Include="LoaderPrivate.h">
<Filter>Header Files</Filter>
</ClInclude>
<ClInclude Include="MmpTlsFiber.h">
<Filter>Header Files</Filter>
</ClInclude>
<ClInclude Include="MmpTlsp.h">
<Filter>Header Files</Filter>
</ClInclude>
</ItemGroup>
<ItemGroup>
<None Include="..\README.md">
MemoryModule/MmpTls.cpp
+43 -82
View File
@@ -1,66 +1,13 @@
#include "stdafx.h"
#if (MMPP_USE_TLS)
#include "MmpTlsp.h"
#include "MmpTlsFiber.h"
#include <cassert>
#include <algorithm>
#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;
}
MemoryModule/MmpTlsFiber.cpp
+100
View File
@@ -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);
}
MemoryModule/MmpTlsFiber.h
+4
View File
@@ -0,0 +1,4 @@
#pragma once
VOID WINAPI MmpQueuePostponedTls(PMMP_TLSP_RECORD record);
VOID MmpTlsFiberInitialize();
MemoryModule/MmpTlsp.h
+58
View File
@@ -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;