From 58379acaf4bb22c4705ab0171d232b62b2f668fb Mon Sep 17 00:00:00 2001 From: Boring <1079299053@qq.com> Date: Wed, 8 Mar 2023 13:36:50 +0800 Subject: [PATCH] Add support for thread pool(ntdll!TppWorkerThread) --- MemoryModule/Loader.cpp | 12 ++- MemoryModule/MmpGlobalData.h | 7 +- MemoryModule/MmpTls.cpp | 186 ++++++++++++++--------------------- test/test.cpp | 23 ++++- 4 files changed, 108 insertions(+), 120 deletions(-) diff --git a/MemoryModule/Loader.cpp b/MemoryModule/Loader.cpp index cd8c927..efc5c46 100644 --- a/MemoryModule/Loader.cpp +++ b/MemoryModule/Loader.cpp @@ -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; diff --git a/MemoryModule/MmpGlobalData.h b/MemoryModule/MmpGlobalData.h index e0fc984..e93adf1 100644 --- a/MemoryModule/MmpGlobalData.h +++ b/MemoryModule/MmpGlobalData.h @@ -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 { diff --git a/MemoryModule/MmpTls.cpp b/MemoryModule/MmpTls.cpp index 6757ae6..7a94a29 100644 --- a/MemoryModule/MmpTls.cpp +++ b/MemoryModule/MmpTls.cpp @@ -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; diff --git a/test/test.cpp b/test/test.cpp index 16f085b..c54455f 100644 --- a/test/test.cpp +++ b/test/test.cpp @@ -2,6 +2,8 @@ #include "../MemoryModule/LoadDllMemoryApi.h" #include +//PMMP_GLOBAL_DATA MmpGlobalDataPtr = *(PMMP_GLOBAL_DATA*)GetProcAddress(GetModuleHandleA("MemoryModule.dll"), "MmpGlobalDataPtr"); + static PVOID ReadDllFile(LPCSTR FileName) { LPVOID buffer; size_t size; @@ -130,8 +132,27 @@ void test_uef() { return; } +void Tp() { + auto pool = CreateThreadpool(nullptr); + if (pool) { + + SetThreadpoolThreadMaximum(pool, 1); + SetThreadpoolThreadMinimum(pool, 1); + + Sleep(1000); + + CloseThreadpool(pool); + } +} + int main() { - test_uef(); + + DisplayStatus(); + test(); + + Tp(); + + WaitForSingleObject(NtCurrentProcess(), INFINITE); return 0; }