From 956932cbdf3e13847a7a23f1c1d6b3a2a3cfaf26 Mon Sep 17 00:00:00 2001 From: Trung <57174311+trungnt2910@users.noreply.github.com> Date: Mon, 4 Dec 2023 22:35:40 +0700 Subject: [PATCH] feat: Implemented MaRegisterPicoProvider Implemented `MaRegisterPicoProvider`. Non-WSL Pico process handlers can use this function to register with Monika. Monika will manage all Pico processes created by these handlers and forward requests to and from the correct driver. Currently, Monika allows 16 handlers. If the code is correctly implemented, these 16 drivers can work harmoniously with each other and with WSL. --- lxmonika/Locker.h | 55 ++++ lxmonika/ProcessMap.cpp | 237 ++++++++++++++ lxmonika/ProcessMap.h | 92 ++++++ lxmonika/driver.cpp | 1 - lxmonika/lxmonika.vcxproj | 5 + lxmonika/lxmonika.vcxproj.filters | 6 + lxmonika/monika.cpp | 496 +++++++++++++++++++++++++++++- lxmonika/monika.h | 7 + lxmonika/monika_providers.cpp | 23 ++ lxmonika/os.h | 32 +- 10 files changed, 939 insertions(+), 15 deletions(-) create mode 100644 lxmonika/Locker.h create mode 100644 lxmonika/ProcessMap.cpp create mode 100644 lxmonika/ProcessMap.h create mode 100644 lxmonika/monika_providers.cpp diff --git a/lxmonika/Locker.h b/lxmonika/Locker.h new file mode 100644 index 0000000..831c4aa --- /dev/null +++ b/lxmonika/Locker.h @@ -0,0 +1,55 @@ +#pragma once + +#include + +template +concept Lockable = requires(T t) +{ + t.Lock(); + t.Unlock(); +}; + +template +class Locker +{ +private: + T* m_obj = NULL; + +public: + Locker(T* obj) : m_obj(obj) + { + m_obj->Lock(); + } + + Locker(const Locker&) = delete; + + Locker(const Locker&& other) + { + m_obj = NULL; + InterlockedExchangePointer((PVOID*)&other.m_obj, m_obj); + } + + Locker& operator=(const Locker&) = delete; + + Locker& operator=(Locker&& other) + { + if (this == &other) + { + return *this; + } + this->~Locker(); + InterlockedExchangePointer((PVOID*)&other.m_obj, m_obj); + return *this; + } + + ~Locker() + { + T* obj = NULL; + obj = (T*)InterlockedExchangePointer((PVOID*)&m_obj, obj); + + if (obj != NULL) + { + obj->Unlock(); + } + } +}; diff --git a/lxmonika/ProcessMap.cpp b/lxmonika/ProcessMap.cpp new file mode 100644 index 0000000..5930ed0 --- /dev/null +++ b/lxmonika/ProcessMap.cpp @@ -0,0 +1,237 @@ +#include "ProcessMap.h" + +VOID +ProcessMap::Initialize() +{ + RtlInitializeGenericTableAvl(&m_table, &_CompareRoutine, &_AllocateRoutine, &_FreeRoutine, + NULL); + ExInitializeFastMutex(&m_mutex); +} + +VOID +ProcessMap::Clear() +{ + PVOID pCurrentElement = NULL; + while ((pCurrentElement = RtlEnumerateGenericTableAvl(&m_table, TRUE)) != NULL) + { + RtlDeleteElementGenericTableAvl(&m_table, pCurrentElement); + } +} + +BOOLEAN +ProcessMap::ProcessBelongsToHandler( + _In_ PEPROCESS process, + _In_ DWORD handler +) +{ + PPROCESS_HANDLER_INFORMATION pInfo = (*this)[process]; + + if (pInfo == NULL) + { + return FALSE; + } + + return pInfo->Handler == handler; +} + +NTSTATUS +ProcessMap::GetProcessHandler( + _In_ PEPROCESS process, + _Out_ DWORD* handler +) +{ + if (handler == NULL) + { + return STATUS_INVALID_PARAMETER; + } + + PPROCESS_HANDLER_INFORMATION pInfo = (*this)[process]; + + if (pInfo == NULL) + { + return STATUS_NOT_FOUND; + } + + *handler = pInfo->Handler; + + return STATUS_SUCCESS; +} + +NTSTATUS +ProcessMap::GetProcessHandler( + _In_ PEPROCESS process, + _Out_ PPROCESS_HANDLER_INFORMATION handlerInformation +) +{ + if (handlerInformation == NULL) + { + return STATUS_INVALID_PARAMETER; + } + + PPROCESS_HANDLER_INFORMATION pInfo = (*this)[process]; + + if (pInfo == NULL) + { + return STATUS_NOT_FOUND; + } + + *handlerInformation = *pInfo; + + return STATUS_SUCCESS; +} + +NTSTATUS +ProcessMap::RegisterProcessHandler( + _In_ PEPROCESS process, + _In_ DWORD handler +) +{ + _NODE node; + memset(&node, 0, sizeof(node)); + + node.Key = process; + node.Value.Handler = handler; + + BOOLEAN bNewEntry = FALSE; + + _PNODE pInsertedNode = (_PNODE) + RtlInsertElementGenericTableAvl(&m_table, &node, sizeof(node), &bNewEntry); + + if (pInsertedNode == FALSE) + { + return STATUS_INSUFFICIENT_RESOURCES; + } + + if (!bNewEntry) + { + return STATUS_ALREADY_REGISTERED; + } + + ASSERT(pInsertedNode->Value.Handler == handler); + + return STATUS_SUCCESS; +} + +NTSTATUS +ProcessMap::SwitchProcessHandler( + _In_ PEPROCESS process, + _In_ DWORD newHandler +) +{ + PPROCESS_HANDLER_INFORMATION pInfo = (*this)[process]; + BOOLEAN bHasInternalProvider = TRUE; + + if (pInfo == NULL) + { + NTSTATUS status = RegisterProcessHandler(process, (DWORD)-1); + if (!NT_SUCCESS(status)) + { + return status; + } + + pInfo = (*this)[process]; + bHasInternalProvider = FALSE; + + ASSERT(pInfo != NULL); + } + + // TODO: This is mainly a Monika restriction, not the data structure's; + // move this check to monika.cpp instead? + if (pInfo->HasParentHandler) + { + return STATUS_NOT_IMPLEMENTED; + } + + pInfo->HasParentHandler = TRUE; + pInfo->HasInternalParentHandler = bHasInternalProvider; + pInfo->ParentHandler = pInfo->Handler; + pInfo->Handler = newHandler; + + return STATUS_SUCCESS; +} + +NTSTATUS +ProcessMap::UnregisterProcess( + _In_ PEPROCESS process +) +{ + _NODE node; + + // Only this member is needed. + node.Key = process; + + if (RtlDeleteElementGenericTableAvl(&m_table, &node) == FALSE) + { + return STATUS_NOT_FOUND; + } + + return STATUS_SUCCESS; +} + +PPROCESS_HANDLER_INFORMATION +ProcessMap::operator[]( + _In_ PEPROCESS key +) +{ + _NODE node + { + .Key = key + }; + + PVOID pBuffer = RtlLookupElementGenericTableAvl(&m_table, &node); + + if (pBuffer == NULL) + { + return NULL; + } + + return &((_PNODE)pBuffer)->Value; +} + +RTL_GENERIC_COMPARE_RESULTS +ProcessMap::_CompareRoutine( + _In_ PRTL_AVL_TABLE Table, + _In_ PVOID FirstStruct, + _In_ PVOID SecondStruct +) +{ + UNREFERENCED_PARAMETER(Table); + + _PNODE FirstNode = (_PNODE)FirstStruct; + _PNODE SecondNode = (_PNODE)SecondStruct; + + if (FirstNode->Key == SecondNode->Key) + { + return GenericEqual; + } + else if ((ULONG_PTR)FirstNode->Key > (ULONG_PTR)SecondNode->Key) + { + return GenericGreaterThan; + } + else + { + return GenericLessThan; + } +} + +PVOID +ProcessMap::_AllocateRoutine( + _In_ PRTL_AVL_TABLE Table, + _In_ CLONG ByteSize +) +{ + UNREFERENCED_PARAMETER(Table); + + return ExAllocatePoolZero(PagedPool, ByteSize, 'PMAP'); +} + +VOID +ProcessMap::_FreeRoutine( + _In_ PRTL_AVL_TABLE Table, + _In_ PVOID Buffer +) +{ + UNREFERENCED_PARAMETER(Table); + + return ExFreePoolWithTag(Buffer, 'PMAP'); +} diff --git a/lxmonika/ProcessMap.h b/lxmonika/ProcessMap.h new file mode 100644 index 0000000..71d4024 --- /dev/null +++ b/lxmonika/ProcessMap.h @@ -0,0 +1,92 @@ +#pragma once + +#include +#include + +#include "Locker.h" + +#pragma warning(disable: 4201) + +typedef union _PROCESS_HANDLER_INFORMATION { + ULONG64 Value; + struct { + DWORD Handler; + DWORD HasParentHandler : 1; + DWORD HasInternalParentHandler : 1; + DWORD ParentHandler : 30; + } DUMMYSTRUCTNAME; +} PROCESS_HANDLER_INFORMATION, *PPROCESS_HANDLER_INFORMATION; + +class ProcessMap +{ +private: + RTL_AVL_TABLE m_table; + FAST_MUTEX m_mutex; + + typedef struct __NODE { + PEPROCESS Key; + PROCESS_HANDLER_INFORMATION Value; + } _NODE, *_PNODE; + + PPROCESS_HANDLER_INFORMATION + operator[]( + _In_ PEPROCESS key + ); + + static RTL_GENERIC_COMPARE_RESULTS + _CompareRoutine( + _In_ PRTL_AVL_TABLE Table, + _In_ PVOID FirstStruct, + _In_ PVOID SecondStruct + ); + static PVOID + _AllocateRoutine( + _In_ PRTL_AVL_TABLE Table, + _In_ CLONG ByteSize + ); + static VOID + _FreeRoutine( + _In_ PRTL_AVL_TABLE Table, + _In_ PVOID Buffer + ); +public: + ProcessMap() = default; + + VOID + Initialize(); + + VOID + Clear(); + + BOOLEAN + ProcessBelongsToHandler( + _In_ PEPROCESS process, + _In_ DWORD handler + ); + NTSTATUS + GetProcessHandler( + _In_ PEPROCESS process, + _Out_ DWORD* handler + ); + NTSTATUS + GetProcessHandler( + _In_ PEPROCESS process, + _Out_ PPROCESS_HANDLER_INFORMATION handlerInformation + ); + NTSTATUS + RegisterProcessHandler( + _In_ PEPROCESS process, + _In_ DWORD handler + ); + NTSTATUS + SwitchProcessHandler( + _In_ PEPROCESS process, + _In_ DWORD newHandler + ); + NTSTATUS + UnregisterProcess( + _In_ PEPROCESS process + ); + void Lock() { ExAcquireFastMutex(&m_mutex); } + void Unlock() { ExReleaseFastMutex(&m_mutex); } +}; diff --git a/lxmonika/driver.cpp b/lxmonika/driver.cpp index 178d943..5657327 100644 --- a/lxmonika/driver.cpp +++ b/lxmonika/driver.cpp @@ -1,7 +1,6 @@ #include "driver.h" #include -#include #include "Logger.h" diff --git a/lxmonika/lxmonika.vcxproj b/lxmonika/lxmonika.vcxproj index 35c08cb..5aa7d01 100644 --- a/lxmonika/lxmonika.vcxproj +++ b/lxmonika/lxmonika.vcxproj @@ -97,6 +97,7 @@ MONIKA_KERNEL_TIMESTAMP="$([System.DateTime]::UtcNow.ToString(ddd MMM dd HH:mm:ss UTC yyyy))";%(PreprocessorDefinitions) + stdcpp20 @@ -110,11 +111,14 @@ + + + @@ -123,6 +127,7 @@ + diff --git a/lxmonika/lxmonika.vcxproj.filters b/lxmonika/lxmonika.vcxproj.filters index b90743f..db85e2f 100644 --- a/lxmonika/lxmonika.vcxproj.filters +++ b/lxmonika/lxmonika.vcxproj.filters @@ -42,6 +42,9 @@ Source Files + + Source Files + @@ -71,6 +74,9 @@ Header Files + + Header Files + diff --git a/lxmonika/monika.cpp b/lxmonika/monika.cpp index fed4559..88e0dc1 100644 --- a/lxmonika/monika.cpp +++ b/lxmonika/monika.cpp @@ -1,8 +1,10 @@ #include "monika.h" +#include "os.h" #include "picosupport.h" #include "Logger.h" +#include "ProcessMap.h" // // Monika macros @@ -20,6 +22,33 @@ #define MONIKA_KERNEL_BUILD_NUMBER "1" #endif +// +// Monika definitions +// + +enum +{ +#define MONIKA_PROVIDER(index) MaPicoProvider##index = index, +#include "monika_providers.cpp" +#undef MONIKA_PROVIDER + MaPicoProviderMaxCount +}; + +#define MONIKA_PROVIDER(index) \ + static PS_PICO_CREATE_PROCESS MaPicoCreateProcess##index; \ + static PS_PICO_CREATE_THREAD MaPicoCreateThread##index; \ + static PS_PICO_GET_PROCESS_CONTEXT MaPicoGetProcessContext##index; \ + static PS_PICO_GET_THREAD_CONTEXT MaPicoGetThreadContext##index; \ + static PS_PICO_SET_THREAD_DESCRIPTOR_BASE MaPicoSetThreadDescriptorBase##index; \ + static PS_PICO_TERMINATE_PROCESS MaPicoTerminateProcess##index; \ + static PS_SET_CONTEXT_THREAD_INTERNAL MaPicoSetContextThreadInternal##index; \ + static PS_GET_CONTEXT_THREAD_INTERNAL MaPicoGetContextThreadInternal##index; \ + static PS_TERMINATE_THREAD MaPicoTerminateThread##index; \ + static PS_SUSPEND_THREAD MaPicoSuspendThread##index; \ + static PS_RESUME_THREAD MaPicoResumeThread##index; +#include "monika_providers.cpp" +#undef MONIKA_PROVIDER + // // Monika data // @@ -27,6 +56,12 @@ static PS_PICO_PROVIDER_ROUTINES MaOriginalProviderRoutines; static PS_PICO_ROUTINES MaOriginalRoutines; +static SIZE_T MaProvidersCount = 0; +static PS_PICO_PROVIDER_ROUTINES MaProviderRoutines[MaPicoProviderMaxCount]; +static PS_PICO_ROUTINES MaRoutines[MaPicoProviderMaxCount]; + +static ProcessMap MaProcessMap; + // // Monika lifetime functions // @@ -81,6 +116,29 @@ MapInitialize() Logger::LogTrace("Successfully patched provider routines."); + // Initialize pico routines +#define MONIKA_PROVIDER(index) \ + MaRoutines[MaPicoProvider##index] = \ + { \ + .Size = sizeof(PS_PICO_ROUTINES), \ + .CreateProcess = MaPicoCreateProcess##index, \ + .CreateThread = MaPicoCreateThread##index, \ + .GetProcessContext = MaPicoGetProcessContext##index, \ + .GetThreadContext = MaPicoGetThreadContext##index, \ + .GetContextThreadInternal = MaPicoGetContextThreadInternal##index, \ + .SetContextThreadInternal = MaPicoSetContextThreadInternal##index, \ + .TerminateThread = MaPicoTerminateThread##index, \ + .ResumeThread = MaPicoResumeThread##index, \ + .SetThreadDescriptorBase = MaPicoSetThreadDescriptorBase##index, \ + .SuspendThread = MaPicoSuspendThread##index, \ + .TerminateProcess = MaPicoTerminateProcess##index \ + }; +#include "monika_providers.cpp" +#undef MONIKA_PROVIDER + + // Initialize process map + MaProcessMap.Initialize(); + return STATUS_SUCCESS; } @@ -93,18 +151,89 @@ MapCleanup() { memcpy(pRoutines, &MaOriginalProviderRoutines, sizeof(MaOriginalProviderRoutines)); } + MaProcessMap.Clear(); +} + + +// +// Pico provider registration +// + +extern "C" +__declspec(dllexport) +NTSTATUS NTAPI +MaRegisterPicoProvider( + _In_ PPS_PICO_PROVIDER_ROUTINES ProviderRoutines, + _Inout_ PPS_PICO_ROUTINES PicoRoutines +) +{ + if (ProviderRoutines->Size != sizeof(PS_PICO_PROVIDER_ROUTINES) + || PicoRoutines->Size != sizeof(PS_PICO_ROUTINES)) + { + return STATUS_INFO_LENGTH_MISMATCH; + } + + const ACCESS_MASK uPicoMaximumRights = STANDARD_RIGHTS_ALL | SPECIFIC_RIGHTS_ALL; + if ((ProviderRoutines->OpenProcess & ~uPicoMaximumRights) != 0 + || (ProviderRoutines->OpenThread & ~uPicoMaximumRights) != 0) + { + return STATUS_INVALID_PARAMETER; + } + + // Acquire an index for the provider. + // Once a provider has been registered, it cannot be unregistered. + // This is the same as the NT kernel as PsUnregisterPicoProvider does not exist. + SIZE_T uProviderIndex = InterlockedIncrementSizeT(&MaProvidersCount) - 1; + if (uProviderIndex >= MaPicoProviderMaxCount) + { + // Prevent SIZE_T overflow (quite hard on 64-bit systems). + InterlockedDecrementSizeT(&MaProvidersCount); + + // PsRegisterPicoProvider would return STATUS_TOO_LATE here. + return STATUS_INSUFFICIENT_RESOURCES; + } + + MaProviderRoutines[uProviderIndex] = *ProviderRoutines; + MaRoutines[uProviderIndex] = *PicoRoutines; + + return STATUS_SUCCESS; } // -// Pico handlers +// Pico provider dispatchers // +#define MA_LOCK_PROCESS_MAP auto _ = Locker(&MaProcessMap) + +#define MA_TRY_DISPATCH_TO_HANDLER(process, function, ...) \ + do \ + { \ + \ + BOOLEAN bHasHandler = FALSE; \ + DWORD dwHandler; \ + \ + { \ + MA_LOCK_PROCESS_MAP; \ + if (NT_SUCCESS(MaProcessMap.GetProcessHandler((process), &dwHandler))) \ + { \ + bHasHandler = TRUE; \ + } \ + } \ + \ + if (bHasHandler) \ + { \ + return MaProviderRoutines[dwHandler].function(__VA_ARGS__); \ + } \ + } while (false) + extern "C" VOID MapSystemCallDispatch( _In_ PPS_PICO_SYSTEM_CALL_INFORMATION SystemCall ) { + MA_TRY_DISPATCH_TO_HANDLER(PsGetCurrentProcess(), DispatchSystemCall, SystemCall); + bool bIsUname = false; struct old_utsname { char sysname[65]; @@ -165,7 +294,9 @@ MapThreadExit( _In_ PETHREAD Thread ) { - MaOriginalProviderRoutines.ExitThread(Thread); + MA_TRY_DISPATCH_TO_HANDLER(PsGetThreadProcess(Thread), ExitThread, Thread); + + return MaOriginalProviderRoutines.ExitThread(Thread); } extern "C" @@ -174,7 +305,32 @@ MapProcessExit( _In_ PEPROCESS Process ) { - MaOriginalProviderRoutines.ExitProcess(Process); + BOOLEAN bHasHandler = FALSE; + PROCESS_HANDLER_INFORMATION info; + + { + MA_LOCK_PROCESS_MAP; + if (NT_SUCCESS(MaProcessMap.GetProcessHandler(Process, &info))) + { + bHasHandler = TRUE; + MaProcessMap.UnregisterProcess(Process); + } + } + + if (bHasHandler) + { + MaProviderRoutines[info.Handler].ExitProcess(Process); + } + + if (bHasHandler && info.HasInternalParentHandler) + { + MaProviderRoutines[info.ParentHandler].ExitProcess(Process); + } + + if (!bHasHandler || !info.HasInternalParentHandler) + { + MaOriginalProviderRoutines.ExitProcess(Process); + } } extern "C" @@ -187,6 +343,14 @@ MapDispatchException( _In_ KPROCESSOR_MODE PreviousMode ) { + MA_TRY_DISPATCH_TO_HANDLER(PsGetCurrentProcess(), DispatchException, + ExceptionRecord, + ExceptionFrame, + TrapFrame, + Chance, + PreviousMode + ); + return MaOriginalProviderRoutines.DispatchException( ExceptionRecord, ExceptionFrame, @@ -203,6 +367,8 @@ MapTerminateProcess( _In_ NTSTATUS TerminateStatus ) { + MA_TRY_DISPATCH_TO_HANDLER(Process, TerminateProcess, Process, TerminateStatus); + return MaOriginalProviderRoutines.TerminateProcess( Process, TerminateStatus @@ -218,9 +384,333 @@ MapWalkUserStack( _In_ ULONG FrameCount ) { + MA_TRY_DISPATCH_TO_HANDLER(PsGetCurrentProcess(), WalkUserStack, + TrapFrame, + Callers, + FrameCount + ); + return MaOriginalProviderRoutines.WalkUserStack( TrapFrame, Callers, FrameCount ); } + +#define MA_PROCESS_BELONGS_TO_HANDLER(process, index) \ + (([](PEPROCESS p, DWORD i) \ + { \ + MA_LOCK_PROCESS_MAP; \ + return MaProcessMap.ProcessBelongsToHandler(p, i); \ + })((process), (index))) + +static +NTSTATUS +MapCreateProcess( + _In_ DWORD HandlerIndex, + _In_ PPS_PICO_PROCESS_ATTRIBUTES ProcessAttributes, + _Outptr_ PHANDLE ProcessHandle +) +{ + NTSTATUS status = MaOriginalRoutines.CreateProcess(ProcessAttributes, ProcessHandle); + + // TODO: Check if this "HANDLE" is actually the same as a `PEPROCESS` struct? + if (NT_SUCCESS(status)) + { + { + MA_LOCK_PROCESS_MAP; + status = MaProcessMap.RegisterProcessHandler((PEPROCESS)*ProcessHandle, HandlerIndex); + } + + if (!NT_SUCCESS(status)) + { + MaOriginalRoutines.TerminateProcess((PEPROCESS)*ProcessHandle, status); + } + } + + return status; +} + +static +NTSTATUS +MapCreateThread( + _In_ DWORD HandlerIndex, + _In_ PPS_PICO_THREAD_ATTRIBUTES ThreadAttributes, + _Outptr_ PHANDLE ThreadHandle +) +{ + if (ThreadAttributes == NULL) + { + return STATUS_INVALID_PARAMETER; + } + + if (!MA_PROCESS_BELONGS_TO_HANDLER((PEPROCESS)ThreadAttributes->Process, HandlerIndex)) + { + return STATUS_INVALID_PARAMETER; + } + + return MaOriginalRoutines.CreateThread(ThreadAttributes, ThreadHandle); +} + +static +PVOID +MapGetProcessContext( + _In_ DWORD HandlerIndex, + _In_ PEPROCESS Process +) +{ + if (!MA_PROCESS_BELONGS_TO_HANDLER(Process, HandlerIndex)) + { + // TODO: Return what on error? + return NULL; + } + + return MaOriginalRoutines.GetProcessContext(Process); +} + +static +PVOID +MapGetThreadContext( + _In_ DWORD HandlerIndex, + _In_ PETHREAD Thread +) +{ + if (!MA_PROCESS_BELONGS_TO_HANDLER(PsGetThreadProcess(Thread), HandlerIndex)) + { + // TODO: Return what on error? + return NULL; + } + + return MaOriginalRoutines.GetThreadContext(Thread); +} + +static +VOID +MapSetThreadDescriptorBase( + _In_ DWORD HandlerIndex, + _In_ PS_PICO_THREAD_DESCRIPTOR_TYPE Type, + _In_ ULONG_PTR Base +) +{ +#if !DBG + UNREFERENCED_PARAMETER(HandlerIndex); +#endif + ASSERT(MA_PROCESS_BELONGS_TO_HANDLER(PsGetCurrentProcess(), HandlerIndex)); + + return MaOriginalRoutines.SetThreadDescriptorBase(Type, Base); +} + +static +NTSTATUS +MapTerminateProcess( + _In_ DWORD HandlerIndex, + _Inout_ PEPROCESS Process, + _In_ NTSTATUS ExitStatus +) +{ + if (!MA_PROCESS_BELONGS_TO_HANDLER(Process, HandlerIndex)) + { + return STATUS_INVALID_PARAMETER; + } + + // TODO: Why __inout for the Process parameter? + return MaOriginalRoutines.TerminateProcess(Process, ExitStatus); +} + +static +NTSTATUS +MapSetContextThreadInternal( + _In_ DWORD HandlerIndex, + _In_ PETHREAD Thread, + _In_ PCONTEXT ThreadContext, + _In_ KPROCESSOR_MODE ProbeMode, + _In_ KPROCESSOR_MODE CtxMode, + _In_ BOOLEAN PerformUnwind +) +{ + if (!MA_PROCESS_BELONGS_TO_HANDLER(PsGetThreadProcess(Thread), HandlerIndex)) + { + return STATUS_INVALID_PARAMETER; + } + + return MaOriginalRoutines.SetContextThreadInternal( + Thread, ThreadContext, ProbeMode, CtxMode, PerformUnwind); +} + +static +NTSTATUS +MapGetContextThreadInternal( + _In_ DWORD HandlerIndex, + _In_ PETHREAD Thread, + _Inout_ PCONTEXT ThreadContext, + _In_ KPROCESSOR_MODE ProbeMode, + _In_ KPROCESSOR_MODE CtxMode, + _In_ BOOLEAN PerformUnwind +) +{ + if (!MA_PROCESS_BELONGS_TO_HANDLER(PsGetThreadProcess(Thread), HandlerIndex)) + { + return STATUS_INVALID_PARAMETER; + } + + return MaOriginalRoutines.GetContextThreadInternal( + Thread, ThreadContext, ProbeMode, CtxMode, PerformUnwind); +} + +static +NTSTATUS +MapTerminateThread( + _In_ DWORD HandlerIndex, + _Inout_ PETHREAD Thread, + _In_ NTSTATUS ExitStatus, + _In_ BOOLEAN DirectTerminate +) +{ + if (!MA_PROCESS_BELONGS_TO_HANDLER(PsGetThreadProcess(Thread), HandlerIndex)) + { + return STATUS_INVALID_PARAMETER; + } + + return MaOriginalRoutines.TerminateThread(Thread, ExitStatus, DirectTerminate); +} + +static +NTSTATUS +MapSuspendThread( + _In_ DWORD HandlerIndex, + _Inout_ PETHREAD Thread, + _Out_opt_ PULONG PreviousSuspendCount +) +{ + if (!MA_PROCESS_BELONGS_TO_HANDLER(PsGetThreadProcess(Thread), HandlerIndex)) + { + return STATUS_INVALID_PARAMETER; + } + + return MaOriginalRoutines.SuspendThread(Thread, PreviousSuspendCount); +} + +static +NTSTATUS +MapResumeThread( + _In_ DWORD HandlerIndex, + _Inout_ PETHREAD Thread, + _Out_opt_ PULONG PreviousSuspendCount +) +{ + if (!MA_PROCESS_BELONGS_TO_HANDLER(PsGetThreadProcess(Thread), HandlerIndex)) + { + return STATUS_INVALID_PARAMETER; + } + + return MaOriginalRoutines.ResumeThread(Thread, PreviousSuspendCount); +} + +#define MONIKA_PROVIDER(index) \ + static NTSTATUS \ + MaPicoCreateProcess##index( \ + _In_ PPS_PICO_PROCESS_ATTRIBUTES ProcessAttributes, \ + _Outptr_ PHANDLE ProcessHandle \ + ) \ + { \ + return MapCreateProcess(index, ProcessAttributes, ProcessHandle); \ + } \ + \ + static NTSTATUS \ + MaPicoCreateThread##index( \ + _In_ PPS_PICO_THREAD_ATTRIBUTES ThreadAttributes, \ + _Outptr_ PHANDLE ThreadHandle \ + ) \ + { \ + return MapCreateThread(index, ThreadAttributes, ThreadHandle); \ + } \ + \ + static PVOID \ + MaPicoGetProcessContext##index( \ + _In_ PEPROCESS Process \ + ) \ + { \ + return MapGetProcessContext(index, Process); \ + } \ + \ + static PVOID \ + MaPicoGetThreadContext##index( \ + _In_ PETHREAD Thread \ + ) \ + { \ + return MapGetThreadContext(index, Thread); \ + } \ + \ + static VOID \ + MaPicoSetThreadDescriptorBase##index( \ + _In_ PS_PICO_THREAD_DESCRIPTOR_TYPE Type, \ + _In_ ULONG_PTR Base \ + ) \ + { \ + return MapSetThreadDescriptorBase(index, Type, Base); \ + } \ + \ + static NTSTATUS \ + MaPicoTerminateProcess##index( \ + __inout PEPROCESS Process, \ + __in NTSTATUS ExitStatus \ + ) \ + { \ + return MapTerminateProcess(index, Process, ExitStatus); \ + } \ + \ + static NTSTATUS \ + MaPicoSetContextThreadInternal##index( \ + __in PETHREAD Thread, \ + __in PCONTEXT ThreadContext, \ + __in KPROCESSOR_MODE ProbeMode, \ + __in KPROCESSOR_MODE CtxMode, \ + __in BOOLEAN PerformUnwind \ + ) \ + { \ + return MapSetContextThreadInternal(index, \ + Thread, ThreadContext, ProbeMode, CtxMode, PerformUnwind); \ + } \ + \ + static NTSTATUS \ + MaPicoGetContextThreadInternal##index( \ + __in PETHREAD Thread, \ + __inout PCONTEXT ThreadContext, \ + __in KPROCESSOR_MODE ProbeMode, \ + __in KPROCESSOR_MODE CtxMode, \ + __in BOOLEAN PerformUnwind \ + ) \ + { \ + return MapGetContextThreadInternal(index, \ + Thread, ThreadContext, ProbeMode, CtxMode, PerformUnwind); \ + } \ + \ + static NTSTATUS \ + MaPicoTerminateThread##index( \ + __inout PETHREAD Thread, \ + __in NTSTATUS ExitStatus, \ + __in BOOLEAN DirectTerminate \ + ) \ + { \ + return MapTerminateThread(index, Thread, ExitStatus, DirectTerminate); \ + } \ + \ + static NTSTATUS \ + MaPicoSuspendThread##index( \ + _In_ PETHREAD Thread, \ + _Out_opt_ PULONG PreviousSuspendCount \ + ) \ + { \ + return MapSuspendThread(index, Thread, PreviousSuspendCount); \ + } \ + \ + static NTSTATUS \ + MaPicoResumeThread##index( \ + _In_ PETHREAD Thread, \ + _Out_opt_ PULONG PreviousSuspendCount \ + ) \ + { \ + return MapResumeThread(index, Thread, PreviousSuspendCount); \ + } +#include "monika_providers.cpp" +#undef MONIKA_PROVIDER diff --git a/lxmonika/monika.h b/lxmonika/monika.h index 0fa06d0..54c2f48 100644 --- a/lxmonika/monika.h +++ b/lxmonika/monika.h @@ -13,6 +13,13 @@ extern "C" { #endif +__declspec(dllexport) +NTSTATUS NTAPI + MaRegisterPicoProvider( + _In_ PPS_PICO_PROVIDER_ROUTINES ProviderRoutines, + _Inout_ PPS_PICO_ROUTINES PicoRoutines + ); + NTSTATUS MapInitialize(); diff --git a/lxmonika/monika_providers.cpp b/lxmonika/monika_providers.cpp new file mode 100644 index 0000000..2783fc1 --- /dev/null +++ b/lxmonika/monika_providers.cpp @@ -0,0 +1,23 @@ +#ifndef MONIKA_PROVIDER +#define MONIKA_PROVIDER(index) +#endif + +// Currently, I cannot find a use case for running more than 16 different types of kernel ABI at +// the same time. + +MONIKA_PROVIDER(0) +MONIKA_PROVIDER(1) +MONIKA_PROVIDER(2) +MONIKA_PROVIDER(3) +MONIKA_PROVIDER(4) +MONIKA_PROVIDER(5) +MONIKA_PROVIDER(6) +MONIKA_PROVIDER(7) +MONIKA_PROVIDER(8) +MONIKA_PROVIDER(9) +MONIKA_PROVIDER(10) +MONIKA_PROVIDER(11) +MONIKA_PROVIDER(12) +MONIKA_PROVIDER(13) +MONIKA_PROVIDER(14) +MONIKA_PROVIDER(15) diff --git a/lxmonika/os.h b/lxmonika/os.h index 26414d9..f155437 100644 --- a/lxmonika/os.h +++ b/lxmonika/os.h @@ -67,20 +67,30 @@ typedef struct _RTL_PROCESS_MODULES { RTL_PROCESS_MODULE_INFORMATION Modules[1]; } RTL_PROCESS_MODULES, *PRTL_PROCESS_MODULES; +__declspec(dllimport) NTSTATUS -ZwQuerySection( - IN HANDLE SectionHandle, - IN SECTION_INFORMATION_CLASS InformationClass, - OUT PVOID InformationBuffer, - IN ULONG InformationBufferSize, - OUT PULONG ResultLength OPTIONAL); + ZwQuerySection( + _In_ HANDLE SectionHandle, + _In_ SECTION_INFORMATION_CLASS InformationClass, + _Out_ PVOID InformationBuffer, + _In_ ULONG InformationBufferSize, + _Out_opt_ PULONG ResultLength + ); +__declspec(dllimport) NTSTATUS -ZwQuerySystemInformation( - IN SYSTEM_INFORMATION_CLASS SystemInformationClass, - IN PVOID SystemInformation, - IN ULONG SystemInformationLength, - OUT PULONG ReturnLength); + ZwQuerySystemInformation( + _In_ SYSTEM_INFORMATION_CLASS SystemInformationClass, + _In_ PVOID SystemInformation, + _In_ ULONG SystemInformationLength, + _Out_ PULONG ReturnLength + ); + +__declspec(dllimport) +PEPROCESS + PsGetThreadProcess( + _In_ PETHREAD Thread + ); #ifdef __cplusplus }