Chain multiple hooks to the same function in a single transaction

Source/Demo/ChainedHook.c

@@ -0,0 +1,169 @@
+/*
+ * This demo chains three hooks on user32.dll!EqualRect within a single transaction, and removes
+ * them within a single transaction. The detour attached last is called first.
+ *
+ * Run "Demo.exe -Run ChainedHook".
+ */
+
+#include "Demo.h"
+
+#define CHAINED_HOOK_COUNT 3
+
+/* Covers the longest patch SlimDetours writes, the bytes past it are the same in both snapshots. */
+#define CHAINED_HOOK_CODE_SIZE 32
+
+static typeof(&EqualRect) g_pfnChainedEqualRect[CHAINED_HOOK_COUNT] = { NULL };
+static LONG volatile g_lChainedCallCount = 0;
+static LONG g_lChainedCallOrder[CHAINED_HOOK_COUNT] = { 0 };
+
+static
+BOOL
+WINAPI
+Hooked_ChainedEqualRect1(
+    _In_ CONST RECT *lprc1,
+    _In_ CONST RECT *lprc2)
+{
+    g_lChainedCallOrder[0] = _InterlockedIncrement(&g_lChainedCallCount);
+    return g_pfnChainedEqualRect[0](lprc1, lprc2);
+}
+
+static
+BOOL
+WINAPI
+Hooked_ChainedEqualRect2(
+    _In_ CONST RECT *lprc1,
+    _In_ CONST RECT *lprc2)
+{
+    g_lChainedCallOrder[1] = _InterlockedIncrement(&g_lChainedCallCount);
+    return g_pfnChainedEqualRect[1](lprc1, lprc2);
+}
+
+static
+BOOL
+WINAPI
+Hooked_ChainedEqualRect3(
+    _In_ CONST RECT *lprc1,
+    _In_ CONST RECT *lprc2)
+{
+    g_lChainedCallOrder[2] = _InterlockedIncrement(&g_lChainedCallCount);
+    return g_pfnChainedEqualRect[2](lprc1, lprc2);
+}
+
+static CONST PVOID g_pChainedDetours[CHAINED_HOOK_COUNT] = {
+    (PVOID)Hooked_ChainedEqualRect1,
+    (PVOID)Hooked_ChainedEqualRect2,
+    (PVOID)Hooked_ChainedEqualRect3
+};
+
+static
+NTSTATUS
+GetPageProtect(
+    _In_ PVOID Address,
+    _Out_ PULONG Protect)
+{
+    NTSTATUS Status;
+    MEMORY_BASIC_INFORMATION mbi;
+
+    Status = NtQueryVirtualMemory(NtCurrentProcess(),
+                                  Address,
+                                  MemoryBasicInformation,
+                                  &mbi,
+                                  sizeof(mbi),
+                                  NULL);
+    if (NT_SUCCESS(Status))
+    {
+        *Protect = mbi.Protect;
+    }
+    return Status;
+}
+
+TEST_FUNC(ChainedHook)
+{
+    NTSTATUS Status;
+    HRESULT hr;
+    RECT rc = { 0 };
+    PVOID pCode;
+    BYTE OriginalCode[CHAINED_HOOK_CODE_SIZE];
+    ULONG OriginalProtect, Protect;
+    ULONG i;
+
+    Status = LoadEqualRect();
+    if (!NT_SUCCESS(Status))
+    {
+        TEST_SKIP("Load user32.dll!EqualRect failed with 0x%08lX\n", Status);
+        return;
+    }
+
+    pCode = SlimDetoursCodeFromPointer(g_pfnEqualRect);
+    RtlCopyMemory(OriginalCode, pCode, sizeof(OriginalCode));
+    Status = GetPageProtect(pCode, &OriginalProtect);
+    if (!NT_SUCCESS(Status))
+    {
+        TEST_SKIP("Query page protection failed with 0x%08lX\n", Status);
+        return;
+    }
+
+    hr = SlimDetoursTransactionBegin();
+    if (FAILED(hr))
+    {
+        TEST_FAIL("SlimDetoursTransactionBegin failed with 0x%08lX\n", hr);
+        return;
+    }
+    for (i = 0; i < CHAINED_HOOK_COUNT; i++)
+    {
+        g_pfnChainedEqualRect[i] = g_pfnEqualRect;
+        hr = SlimDetoursAttach((PVOID*)&g_pfnChainedEqualRect[i], g_pChainedDetours[i]);
+        if (FAILED(hr))
+        {
+            SlimDetoursTransactionAbort();
+            TEST_FAIL("SlimDetoursAttach #%lu failed with 0x%08lX\n", i, hr);
+            return;
+        }
+    }
+    hr = SlimDetoursTransactionCommit();
+    if (FAILED(hr))
+    {
+        TEST_FAIL("SlimDetoursTransactionCommit failed with 0x%08lX\n", hr);
+        return;
+    }
+
+    TEST_OK(g_pfnEqualRect(&rc, &rc) != FALSE);
+    TEST_OK(g_lChainedCallCount == CHAINED_HOOK_COUNT);
+    TEST_OK(g_lChainedCallOrder[2] == 1);
+    TEST_OK(g_lChainedCallOrder[1] == 2);
+    TEST_OK(g_lChainedCallOrder[0] == 3);
+
+    /* The page keeps the protection it had before the transaction. */
+    TEST_OK(NT_SUCCESS(GetPageProtect(pCode, &Protect)) && Protect == OriginalProtect);
+
+    /* Detaching the innermost hook first makes the removal take several passes to unwind. */
+    hr = SlimDetoursTransactionBegin();
+    if (FAILED(hr))
+    {
+        TEST_FAIL("SlimDetoursTransactionBegin failed with 0x%08lX\n", hr);
+        return;
+    }
+    for (i = 0; i < CHAINED_HOOK_COUNT; i++)
+    {
+        hr = SlimDetoursDetach((PVOID*)&g_pfnChainedEqualRect[i], g_pChainedDetours[i]);
+        if (FAILED(hr))
+        {
+            SlimDetoursTransactionAbort();
+            TEST_FAIL("SlimDetoursDetach #%lu failed with 0x%08lX\n", i, hr);
+            return;
+        }
+    }
+    hr = SlimDetoursTransactionCommit();
+    if (FAILED(hr))
+    {
+        TEST_FAIL("SlimDetoursTransactionCommit failed with 0x%08lX\n", hr);
+        return;
+    }
+
+    TEST_OK(RtlEqualMemory(pCode, OriginalCode, sizeof(OriginalCode)));
+    TEST_OK(NT_SUCCESS(GetPageProtect(pCode, &Protect)) && Protect == OriginalProtect);
+
+    g_lChainedCallCount = 0;
+    TEST_OK(g_pfnEqualRect(&rc, &rc) != FALSE);
+    TEST_OK(g_lChainedCallCount == 0);
+}

Source/Demo/Demo.vcxproj

@@ -190,6 +190,7 @@
     </Link>
   </ItemDefinitionGroup>
   <ItemGroup>
+    <ClCompile Include="ChainedHook.c" />
     <ClCompile Include="COMHook.c" />
     <ClCompile Include="DeadLock.c" />
     <ClCompile Include="DelayHook.c" />

Source/Demo/Demo.vcxproj.filters

@@ -5,6 +5,7 @@
     <ClCompile Include="DeadLock.c" />
     <ClCompile Include="DelayHook.c" />
     <ClCompile Include="TwiceSimpleHook.c" />
+    <ClCompile Include="ChainedHook.c" />
     <ClCompile Include="Instruction.c" />
     <ClCompile Include="COMHook.c" />
     <ClCompile Include="OutputDebugStringHook.c" />

Source/Demo/Main.c

@@ -9,6 +9,7 @@ TEST_DECL_FUNC(OutputDebugStringHook);
 
 /* Auto tests */
 TEST_DECL_FUNC(TwiceSimpleHook);
+TEST_DECL_FUNC(ChainedHook);
 TEST_DECL_FUNC(Instruction);
 #if _WIN32_WINNT >= _WIN32_WINNT_WIN6
 TEST_DECL_FUNC(DelayHook);
@@ -22,6 +23,7 @@ CONST UNITTEST_ENTRY UnitTestList[] = {
 #endif
 
     TEST_DECL_ENTRY(TwiceSimpleHook),
+    TEST_DECL_ENTRY(ChainedHook),
     TEST_DECL_ENTRY(Instruction),
 #if _WIN32_WINNT >= _WIN32_WINNT_WIN6
     TEST_DECL_ENTRY(DelayHook),

Source/KNSoft.SlimDetours/Instruction.c

@@ -153,6 +153,29 @@ detour_gen_jmp_indirect(
     return pbCode + sizeof(INT32);
 }
 
+#if defined(_M_X64)
+
+_Ret_notnull_
+PBYTE
+detour_gen_jmp_aligned_literal(
+    _In_ PBYTE pbCode,
+    _In_ PBYTE pbJmpVal)
+{
+    // The destination is stored inline, aligned and behind the jump, so it can be anywhere.
+    DETOUR_ASSERT(((ULONG_PTR)pbCode & (sizeof(PBYTE) - 1)) == 0);
+
+    *pbCode++ = 0xff;   // jmp [+imm32]
+    *pbCode++ = 0x25;
+    *((INT32*)pbCode) = 2;
+    pbCode += sizeof(INT32);
+    *pbCode++ = 0xcc;   // brk;
+    *pbCode++ = 0xcc;   // brk;
+    *((PBYTE*)pbCode) = pbJmpVal;
+    return pbCode + sizeof(PBYTE);
+}
+
+#endif
+
 BOOL
 detour_is_jmp_indirect_to(
     _In_ PBYTE pbCode,

Source/KNSoft.SlimDetours/SlimDetours.inl

@@ -10,9 +10,11 @@
 #if _DEBUG
 #define DETOUR_TRACE DbgPrint
 #define DETOUR_BREAK() __debugbreak()
+#define DETOUR_ASSERT(Expression) ((Expression) ? (VOID)0 : __debugbreak())
 #else
 #define DETOUR_TRACE(Format, ...)
 #define DETOUR_BREAK()
+#define DETOUR_ASSERT(Expression) ((VOID)0)
 #endif
 
 EXTERN_C_START
@@ -65,7 +67,7 @@ typedef struct _DETOUR_TRAMPOLINE
 #elif defined(_M_IX86) || defined(_M_X64)
     BYTE            rbCode[30];         // target code + jmp to pbRemain.
 #endif
-    BYTE            cbCode;             // size of moved target code.
+    BYTE            cbCode;             // size of the code in rbCode.
 #if defined(_M_ARM64) || defined(_M_ARM64EC)
     BYTE            cbCodeBreak[3];     // padding to make debugging easier.
 #elif defined(_M_IX86) || defined(_M_X64)
@@ -109,6 +111,7 @@ struct _DETOUR_OPERATION
     PDETOUR_OPERATION pNext;
     BOOL fIsAdd : 1;
     BOOL fIsRemove : 1;
+    BOOL fIsRestored : 1;   // removal only: the target code was put back.
 #if defined(_M_ARM64EC)
     BOOL fTargetArm64Ec : 1;
 #endif
@@ -221,6 +224,16 @@ detour_is_jmp_indirect_to(
     _In_ PBYTE pbCode,
     _In_ PBYTE* ppbJmpVal);
 
+#if defined(_M_X64)
+
+_Ret_notnull_
+PBYTE
+detour_gen_jmp_aligned_literal(
+    _In_ PBYTE pbCode,
+    _In_ PBYTE pbJmpVal);
+
+#endif
+
 #endif
 
 #if defined(_M_ARM64) || defined(_M_ARM64EC)