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 support for thread pool(ntdll!TppWorkerThread)

Browse Source
This commit is contained in:
Boring
2023-03-08 13:36:50 +08:00
parent a32f925e8c
commit 58379acaf4
4 changed files with 108 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
MemoryModule/Loader.cpp
+8 -4
View File
@@ -64,26 +64,30 @@ NTSTATUS NTAPI LdrLoadDllMemoryExW(
if (dwFlags & LOAD_FLAGS_USE_DLL_NAME && (!DllName || !DllFullName))return STATUS_INVALID_PARAMETER_3;
if (DllName) {
PLIST_ENTRY ListHead, ListEntry;
PLDR_DATA_TABLE_ENTRY CurEntry;
PLIST_ENTRY ListHead = &NtCurrentPeb()->Ldr->InLoadOrderModuleList, ListEntry = ListHead->Flink;
PIMAGE_NT_HEADERS h1 = RtlImageNtHeader(BufferAddress), h2 = nullptr;
if (!h1)return STATUS_INVALID_IMAGE_FORMAT;
ListEntry = (ListHead = &NtCurrentPeb()->Ldr->InLoadOrderModuleList)->Flink;
while (ListEntry != ListHead) {
CurEntry = CONTAINING_RECORD(ListEntry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
PLDR_DATA_TABLE_ENTRY 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 (!_wcsnicmp(DllName, CurEntry->BaseDllName.Buffer, (CurEntry->BaseDllName.Length / sizeof(wchar_t)) - 4) ||
!_wcsnicmp(DllName, CurEntry->BaseDllName.Buffer, CurEntry->BaseDllName.Length / sizeof(wchar_t))) {
/* Let's compare their headers */
if (!(h2 = RtlImageNtHeader(CurEntry->DllBase)))continue;
if (!(module = MapMemoryModuleHandle((HMEMORYMODULE)CurEntry->DllBase)))continue;
if ((h1->OptionalHeader.SizeOfCode == h2->OptionalHeader.SizeOfCode) &&
(h1->OptionalHeader.SizeOfHeaders == h2->OptionalHeader.SizeOfHeaders)) {
/* This is our entry!, update load count and return success */
if (!module->UseReferenceCount || dwFlags & LOAD_FLAGS_NOT_USE_REFERENCE_COUNT)return STATUS_INVALID_PARAMETER_3;
RtlUpdateReferenceCount(module, FLAG_REFERENCE);
*BaseAddress = (HMEMORYMODULE)CurEntry->DllBase;
if (LdrEntry)*LdrEntry = CurEntry;
MemoryModule/MmpGlobalData.h
+4 -3
View File
@@ -29,10 +29,11 @@ typedef struct _MMP_TLS_DATA {
DWORD MmpActiveThreadCount;
struct {
decltype(&NtCreateThread) OriginNtCreateThread;
decltype(&NtCreateThreadEx) OriginNtCreateThreadEx;
PVOID HookReserved1;
PVOID HookReserved2;
decltype(&NtSetInformationProcess) OriginNtSetInformationProcess;
decltype(&LdrShutdownThread) OriginLdrShutdownThread;
decltype(&RtlUserThreadStart) OriginRtlUserThreadStart;
}Hooks;
}MMP_TLS_DATA, * PMMP_TLS_DATA;
@@ -77,7 +78,7 @@ typedef enum class _WINDOWS_VERSION :BYTE {
}WINDOWS_VERSION;
#define MEMORY_MODULE_MAJOR_VERSION 1
#define MEMORY_MODULE_MINOR_VERSION 1
#define MEMORY_MODULE_MINOR_VERSION 2
typedef struct _MMP_GLOBAL_DATA {
MemoryModule/MmpTls.cpp
+74 -112
View File
@@ -114,7 +114,37 @@ DWORD NTAPI MmpGetThreadCount() {
while (true) {
if (p->UniqueProcessId == pid) {
result = p->NumberOfThreads;
OBJECT_ATTRIBUTES oa{};
InitializeObjectAttributes(&oa, nullptr, 0, nullptr, nullptr);
THREAD_BASIC_INFORMATION tbi{};
NTSTATUS status;
for (ULONG i = 0; i < p->NumberOfThreads; ++i) {
HANDLE hThread;
status = NtOpenThread(
&hThread,
THREAD_QUERY_INFORMATION,
&oa,
&p->Threads[i].ClientId
);
if (NT_SUCCESS(status)) {
status = NtQueryInformationThread(
hThread,
ThreadBasicInformation,
&tbi,
sizeof(tbi),
nullptr
);
if (NT_SUCCESS(status) && !!tbi.TebBaseAddress->ThreadLocalStoragePointer) {
++result;
}
NtClose(hThread);
}
}
break;
}
@@ -128,6 +158,22 @@ DWORD NTAPI MmpGetThreadCount() {
return result;
}
PMMP_TLSP_RECORD MmpFindTlspRecordLockHeld() {
PLIST_ENTRY entry = MmpGlobalDataPtr->MmpTls->MmpThreadLocalStoragePointer.Flink;
while (entry != &MmpGlobalDataPtr->MmpTls->MmpThreadLocalStoragePointer) {
auto p = CONTAINING_RECORD(entry, MMP_TLSP_RECORD, InMmpThreadLocalStoragePointer);
if (p->UniqueThread == NtCurrentThreadId()) {
assert(p->TlspMmpBlock == NtCurrentTeb()->ThreadLocalStoragePointer);
return p;
}
entry = entry->Flink;
}
return nullptr;
}
DWORD NTAPI MmpUserThreadStart(LPVOID lpThreadParameter) {
THREAD_CONTEXT Context;
@@ -140,8 +186,6 @@ DWORD NTAPI MmpUserThreadStart(LPVOID lpThreadParameter) {
lpThreadParameter,
sizeof(Context)
);
RtlFreeHeap(RtlProcessHeap(), 0, lpThreadParameter);
}
__except (EXCEPTION_EXECUTE_HANDLER) {
return GetExceptionCode();
@@ -151,6 +195,15 @@ DWORD NTAPI MmpUserThreadStart(LPVOID lpThreadParameter) {
goto __skip_tls;
}
//
// Check if we have already initialized
//
EnterCriticalSection(&MmpGlobalDataPtr->MmpTls->MmpTlspLock);
record = MmpFindTlspRecordLockHeld();
LeaveCriticalSection(&MmpGlobalDataPtr->MmpTls->MmpTlspLock);
if (!!record)goto __skip_tls;
//
// Allocate and replace ThreadLocalStoragePointer for new thread
//
@@ -162,7 +215,6 @@ DWORD NTAPI MmpUserThreadStart(LPVOID lpThreadParameter) {
record->TlspMmpBlock = (PVOID*)MmpAllocateTlsp();
record->UniqueThread = NtCurrentThreadId();
if (record->TlspMmpBlock) {
auto size = CONTAINING_RECORD(record->TlspLdrBlock, TLS_VECTOR, ModuleTlsData)->Length;
if ((HANDLE)(ULONG_PTR)size != NtCurrentThreadId()) {
RtlCopyMemory(
@@ -230,91 +282,14 @@ __skip_tls:
return Context.ThreadStartRoutine(Context.ThreadParameter);
}
NTSTATUS NTAPI HookNtCreateThread(
_Out_ PHANDLE ThreadHandle,
_In_ ACCESS_MASK DesiredAccess,
_In_opt_ POBJECT_ATTRIBUTES ObjectAttributes,
_In_ HANDLE ProcessHandle,
_Out_ PCLIENT_ID ClientId,
_In_ PCONTEXT ThreadContext,
_In_ PVOID InitialTeb,
_In_ BOOLEAN CreateSuspended) {
CONTEXT Context = *ThreadContext;
PTHREAD_CONTEXT _Context = PTHREAD_CONTEXT(RtlAllocateHeap(RtlProcessHeap(), 0, sizeof(*_Context)));
NTSTATUS status;
VOID NTAPI HookRtlUserThreadStart(
_In_ PTHREAD_START_ROUTINE Function,
_In_ PVOID Parameter) {
THREAD_CONTEXT Context;
Context.ThreadStartRoutine = PTHREAD_START_ROUTINE(Function);
Context.ThreadParameter = Parameter;
if (!_Context)return STATUS_NO_MEMORY;
#ifndef _WIN64
_Context->ThreadStartRoutine = PTHREAD_START_ROUTINE(Context.Eax);
_Context->ThreadParameter = LPVOID(Context.Ebx);
Context.Eax = DWORD(MmpUserThreadStart);
Context.Ebx = DWORD(_Context);
#else
_Context->ThreadStartRoutine = PTHREAD_START_ROUTINE(Context.Rcx);
_Context->ThreadParameter = LPVOID(Context.Rdx);
Context.Rcx = ULONG64(MmpUserThreadStart);
Context.Rdx = ULONG64(_Context);
#endif
status = MmpGlobalDataPtr->MmpTls->Hooks.OriginNtCreateThread(
ThreadHandle,
DesiredAccess,
ObjectAttributes,
ProcessHandle,
ClientId,
&Context,
(PINITIAL_TEB)InitialTeb,
CreateSuspended
);
if (!NT_SUCCESS(status)) {
RtlFreeHeap(RtlProcessHeap(), 0, _Context);
}
return status;
}
NTSTATUS NTAPI HookNtCreateThreadEx(
_Out_ PHANDLE ThreadHandle,
_In_ ACCESS_MASK DesiredAccess,
_In_opt_ POBJECT_ATTRIBUTES ObjectAttributes,
_In_ HANDLE ProcessHandle,
_In_ PVOID StartRoutine,
_In_opt_ PVOID Argument,
_In_ ULONG CreateFlags,
_In_ SIZE_T ZeroBits,
_In_ SIZE_T StackSize,
_In_ SIZE_T MaximumStackSize,
_In_opt_ PVOID AttributeList) {
PTHREAD_CONTEXT Context = PTHREAD_CONTEXT(RtlAllocateHeap(RtlProcessHeap(), 0, sizeof(*Context)));
if (!Context) {
return STATUS_NO_MEMORY;
}
Context->ThreadStartRoutine = PTHREAD_START_ROUTINE(StartRoutine);
Context->ThreadParameter = Argument;
NTSTATUS status = MmpGlobalDataPtr->MmpTls->Hooks.OriginNtCreateThreadEx(
ThreadHandle,
DesiredAccess,
ObjectAttributes,
ProcessHandle,
MmpUserThreadStart,
Context,
CreateFlags,
ZeroBits,
StackSize,
MaximumStackSize,
(PPS_ATTRIBUTE_LIST)AttributeList
);
if (!NT_SUCCESS(status)) {
RtlFreeHeap(RtlProcessHeap(), 0, Context);
}
return status;
return MmpGlobalDataPtr->MmpTls->Hooks.OriginRtlUserThreadStart(MmpUserThreadStart, &Context);
}
VOID NTAPI HookLdrShutdownThread(VOID) {
@@ -327,27 +302,16 @@ VOID NTAPI HookLdrShutdownThread(VOID) {
//
EnterCriticalSection(&MmpGlobalDataPtr->MmpTls->MmpTlspLock);
entry = MmpGlobalDataPtr->MmpTls->MmpThreadLocalStoragePointer.Flink;
while (entry != &MmpGlobalDataPtr->MmpTls->MmpThreadLocalStoragePointer) {
auto p = CONTAINING_RECORD(entry, MMP_TLSP_RECORD, InMmpThreadLocalStoragePointer);
if (p->UniqueThread == NtCurrentThreadId()) {
assert(p->TlspMmpBlock == NtCurrentTeb()->ThreadLocalStoragePointer);
//
// Restore tlsp
//
NtCurrentTeb()->ThreadLocalStoragePointer = p->TlspLdrBlock;
RemoveEntryList(&p->InMmpThreadLocalStoragePointer);
record = p;
break;
}
entry = entry->Flink;
}
record = MmpFindTlspRecordLockHeld();
if (record) {
//
// Restore tlsp
//
NtCurrentTeb()->ThreadLocalStoragePointer = record->TlspLdrBlock;
RemoveEntryList(&record->InMmpThreadLocalStoragePointer);
--MmpGlobalDataPtr->MmpTls->MmpActiveThreadCount;
}
@@ -829,17 +793,15 @@ BOOL NTAPI MmpTlsInitialize() {
// Hook functions
//
MmpGlobalDataPtr->MmpTls->Hooks.OriginNtCreateThread = NtCreateThread;
MmpGlobalDataPtr->MmpTls->Hooks.OriginNtCreateThreadEx = NtCreateThreadEx;
MmpGlobalDataPtr->MmpTls->Hooks.OriginLdrShutdownThread = LdrShutdownThread;
MmpGlobalDataPtr->MmpTls->Hooks.OriginNtSetInformationProcess = NtSetInformationProcess;
MmpGlobalDataPtr->MmpTls->Hooks.OriginRtlUserThreadStart = (decltype(&RtlUserThreadStart))GetProcAddress((HMODULE)MmpGlobalDataPtr->MmpBaseAddressIndex->NtdllLdrEntry->DllBase, "RtlUserThreadStart");
DetourTransactionBegin();
DetourUpdateThread(NtCurrentThread());
DetourAttach((PVOID*)&MmpGlobalDataPtr->MmpTls->Hooks.OriginNtCreateThread, HookNtCreateThread);
DetourAttach((PVOID*)&MmpGlobalDataPtr->MmpTls->Hooks.OriginNtCreateThreadEx, HookNtCreateThreadEx);
DetourAttach((PVOID*)&MmpGlobalDataPtr->MmpTls->Hooks.OriginLdrShutdownThread, HookLdrShutdownThread);
DetourAttach((PVOID*)&MmpGlobalDataPtr->MmpTls->Hooks.OriginNtSetInformationProcess, HookNtSetInformationProcess);
DetourAttach((PVOID*)&MmpGlobalDataPtr->MmpTls->Hooks.OriginRtlUserThreadStart, HookRtlUserThreadStart);
DetourTransactionCommit();
return TRUE;