diff --git a/Source/KNSoft.SlimDetours/SlimDetours.h b/Source/KNSoft.SlimDetours/SlimDetours.h index 3d757c52..5530c36f 100644 --- a/Source/KNSoft.SlimDetours/SlimDetours.h +++ b/Source/KNSoft.SlimDetours/SlimDetours.h @@ -62,11 +62,35 @@ SlimDetoursAttach( _Inout_ PVOID* ppPointer, _In_ PVOID pDetour); +typedef struct _DETOUR_DETACH_OPTIONS +{ + PVOID *ppTrampolineToFreeManually; +} DETOUR_DETACH_OPTIONS, *PDETOUR_DETACH_OPTIONS; + +typedef const DETOUR_DETACH_OPTIONS* PCDETOUR_DETACH_OPTIONS; + HRESULT NTAPI +SlimDetoursDetachEx( + _Inout_ PVOID* ppPointer, + _In_ PVOID pDetour, + _In_ PCDETOUR_DETACH_OPTIONS pOptions); + +FORCEINLINE +HRESULT SlimDetoursDetach( _Inout_ PVOID* ppPointer, - _In_ PVOID pDetour); + _In_ PVOID pDetour) +{ + DETOUR_DETACH_OPTIONS Options; + Options.ppTrampolineToFreeManually = NULL; + return SlimDetoursDetachEx(ppPointer, pDetour, &Options); +} + +HRESULT +NTAPI +SlimDetoursFreeTrampoline( + _In_ PVOID pTrampoline); PVOID NTAPI diff --git a/Source/KNSoft.SlimDetours/SlimDetours.inl b/Source/KNSoft.SlimDetours/SlimDetours.inl index 6f0f44ce..db870993 100644 --- a/Source/KNSoft.SlimDetours/SlimDetours.inl +++ b/Source/KNSoft.SlimDetours/SlimDetours.inl @@ -96,17 +96,24 @@ _STATIC_ASSERT(sizeof(DETOUR_TRAMPOLINE) == 96); _STATIC_ASSERT(sizeof(DETOUR_TRAMPOLINE) == 184); #endif +enum +{ + DETOUR_OPERATION_NONE = 0, + DETOUR_OPERATION_ADD, + DETOUR_OPERATION_REMOVE, +}; + typedef struct _DETOUR_OPERATION DETOUR_OPERATION, *PDETOUR_OPERATION; struct _DETOUR_OPERATION { PDETOUR_OPERATION pNext; - BOOL fIsAdd : 1; - BOOL fIsRemove : 1; + DWORD dwOperation; PBYTE* ppbPointer; PBYTE pbTarget; PDETOUR_TRAMPOLINE pTrampoline; ULONG dwPerm; + PVOID* ppTrampolineToFreeManually; }; /* Memory management */ diff --git a/Source/KNSoft.SlimDetours/Thread.c b/Source/KNSoft.SlimDetours/Thread.c index e18335a3..cc64d809 100644 --- a/Source/KNSoft.SlimDetours/Thread.c +++ b/Source/KNSoft.SlimDetours/Thread.c @@ -315,7 +315,7 @@ detour_thread_update( bUpdateContext = FALSE; for (PDETOUR_OPERATION o = PendingOperations; o != NULL && !bUpdateContext; o = o->pNext) { - if (o->fIsRemove) + if (o->dwOperation == DETOUR_OPERATION_REMOVE) { if (cxt.CONTEXT_PC >= (ULONG_PTR)o->pTrampoline->rbCode && cxt.CONTEXT_PC < ((ULONG_PTR)o->pTrampoline->rbCode + RTL_FIELD_SIZE(DETOUR_TRAMPOLINE, rbCode))) @@ -331,7 +331,7 @@ detour_thread_update( bUpdateContext = TRUE; } #endif - } else if (o->fIsAdd) + } else if (o->dwOperation == DETOUR_OPERATION_ADD) { if (cxt.CONTEXT_PC >= (ULONG_PTR)o->pbTarget && cxt.CONTEXT_PC < ((ULONG_PTR)o->pbTarget + o->pTrampoline->cbRestore)) diff --git a/Source/KNSoft.SlimDetours/Transaction.c b/Source/KNSoft.SlimDetours/Transaction.c index ab0e0700..691f6ba8 100644 --- a/Source/KNSoft.SlimDetours/Transaction.c +++ b/Source/KNSoft.SlimDetours/Transaction.c @@ -89,7 +89,7 @@ SlimDetoursTransactionAbort(VOID) pMem = o->pbTarget; sMem = o->pTrampoline->cbRestore; NtProtectVirtualMemory(NtCurrentProcess(), &pMem, &sMem, o->dwPerm, &dwOld); - if (o->fIsAdd) + if (o->dwOperation == DETOUR_OPERATION_ADD) { detour_free_trampoline(o->pTrampoline); o->pTrampoline = NULL; @@ -146,7 +146,7 @@ SlimDetoursTransactionCommit(VOID) o = s_pPendingOperations; do { - if (o->fIsRemove) + if (o->dwOperation == DETOUR_OPERATION_REMOVE) { // Check if the jmps still points where we expect, otherwise someone might have hooked us. BOOL hookIsStillThere = @@ -163,14 +163,16 @@ SlimDetoursTransactionCommit(VOID) NtFlushInstructionCache(NtCurrentProcess(), o->pbTarget, o->pTrampoline->cbRestore); } else { - // Don't remove in this case, put in bypass mode and leak trampoline. - o->fIsRemove = FALSE; - o->pTrampoline->pbDetour = o->pTrampoline->rbCode; + // Don't remove and leak trampoline in this case. + o->dwOperation = DETOUR_OPERATION_NONE; DETOUR_TRACE("detours: Leaked hook on pbTarget=%p due to external hooking\n", o->pbTarget); } + // Put hook in bypass mode. + o->pTrampoline->pbDetour = o->pTrampoline->rbCode; + *o->ppbPointer = o->pbTarget; - } else if (o->fIsAdd) + } else if (o->dwOperation == DETOUR_OPERATION_ADD) { DETOUR_TRACE("detours: pbTramp =%p, pbRemain=%p, pbDetour=%p, cbRestore=%u\n", o->pTrampoline, @@ -237,11 +239,18 @@ SlimDetoursTransactionCommit(VOID) pMem = o->pbTarget; sMem = o->pTrampoline->cbRestore; NtProtectVirtualMemory(NtCurrentProcess(), &pMem, &sMem, o->dwPerm, &dwOld); - if (o->fIsRemove) + if (o->dwOperation == DETOUR_OPERATION_REMOVE) { - detour_free_trampoline(o->pTrampoline); + if (!o->ppTrampolineToFreeManually) + { + detour_free_trampoline(o->pTrampoline); + freed = TRUE; + } else + { + // The caller is responsible for freeing the trampoline. + *o->ppTrampolineToFreeManually = o->pTrampoline; + } o->pTrampoline = NULL; - freed = TRUE; } n = o->pNext; @@ -461,8 +470,7 @@ SlimDetoursAttach( pTrampoline->rbCode[8], pTrampoline->rbCode[9], pTrampoline->rbCode[10], pTrampoline->rbCode[11]); - o->fIsAdd = TRUE; - o->fIsRemove = FALSE; + o->dwOperation = DETOUR_OPERATION_ADD; o->ppbPointer = (PBYTE*)ppPointer; o->pTrampoline = pTrampoline; o->pbTarget = pbTarget; @@ -475,9 +483,10 @@ SlimDetoursAttach( HRESULT NTAPI -SlimDetoursDetach( +SlimDetoursDetachEx( _Inout_ PVOID* ppPointer, - _In_ PVOID pDetour) + _In_ PVOID pDetour, + _In_ PCDETOUR_DETACH_OPTIONS pOptions) { NTSTATUS Status; PVOID pMem; @@ -525,18 +534,72 @@ SlimDetoursDetach( goto fail; } - o->fIsAdd = FALSE; - o->fIsRemove = TRUE; + o->dwOperation = DETOUR_OPERATION_REMOVE; o->ppbPointer = (PBYTE*)ppPointer; o->pTrampoline = pTrampoline; o->pbTarget = pbTarget; o->dwPerm = dwOld; + o->ppTrampolineToFreeManually = pOptions->ppTrampolineToFreeManually; o->pNext = s_pPendingOperations; s_pPendingOperations = o; return HRESULT_FROM_NT(STATUS_SUCCESS); } +HRESULT +NTAPI +SlimDetoursFreeTrampoline( + _In_ PVOID pTrampoline) +{ + if (pTrampoline == NULL) + { + return HRESULT_FROM_NT(STATUS_SUCCESS); + } + + // This function can be called as part of a transaction or outside of a transaction. + HANDLE nPrevPendingThreadId = _InterlockedCompareExchangePointer(&s_nPendingThreadId, NtCurrentThreadId(), NULL); + BOOL bInTransaction = nPrevPendingThreadId != NULL; + if (bInTransaction && nPrevPendingThreadId != NtCurrentThreadId()) + { + return HRESULT_FROM_NT(STATUS_TRANSACTIONAL_CONFLICT); + } + + NTSTATUS Status; + + if (!bInTransaction) + { + // Make sure the trampoline pages are writable. + Status = detour_writable_trampoline_regions(); + if (!NT_SUCCESS(Status)) + { + goto fail; + } + } + + detour_free_trampoline((PDETOUR_TRAMPOLINE)pTrampoline); + detour_free_trampoline_region_if_unused((PDETOUR_TRAMPOLINE)pTrampoline); + + if (!bInTransaction) + { + detour_runnable_trampoline_regions(); + } + + Status = STATUS_SUCCESS; + +fail: + if (!bInTransaction) + { +#ifdef _MSC_VER +#pragma warning(disable: __WARNING_INTERLOCKED_ACCESS) +#endif + s_nPendingThreadId = NULL; +#ifdef _MSC_VER +#pragma warning(default: __WARNING_INTERLOCKED_ACCESS) +#endif + } + return HRESULT_FROM_NT(Status); +} + HRESULT NTAPI SlimDetoursUninitialize(VOID)