[SLIMDETOURS] Small refactors

- [HOOK] COM initialization should be done by the caller - [HOOK] Improve COM hook initialization - [SLIMDETOURS] **Refactor some APIs and add some comments**
KNSoft · Oct 29, 2024 · 831c15d · 831c15d
1 parent d87330b
commit 831c15d
Show file tree

Hide file tree

Showing 15 changed files with 447 additions and 242 deletions.
diff --git a/README.md b/README.md
@@ -52,9 +52,11 @@ If your project configuration name is neither "Release" nor "Debug", [MSBuild sh
 
 The usage has been simplified, e.g. the hook only needs one line:
 ```C
-SlimDetoursSetHook((PVOID*)&g_pfnXxx, Hooked_Xxx);
+SlimDetoursInlineHook(TRUE, (PVOID*)&g_pfnXxx, Hooked_Xxx);  // Hook
+...
+SlimDetoursInlineHook(FALSE, (PVOID*)&g_pfnXxx, Hooked_Xxx); // Unhook
 ```
-For more simplified API see [Wrapper.c](https://github.com/KNSoft/KNSoft.SlimDetours/blob/main/Source/SlimDetours/Wrapper.c).
+For more simplified API see [InlineHook.c](https://github.com/KNSoft/KNSoft.SlimDetours/blob/main/Source/SlimDetours/InlineHook.c).
 
 ### Details
 

diff --git a/README.zh-CN.md b/README.zh-CN.md
@@ -52,9 +52,11 @@ NuGet包[KNSoft.SlimDetours](https://www.nuget.org/packages/KNSoft.SlimDetours)
 
 用法已进行了简化，例如挂钩仅需一行：
 ```C
-SlimDetoursSetHook((PVOID*)&g_pfnXxx, Hooked_Xxx);
+SlimDetoursInlineHook(TRUE, (PVOID*)&g_pfnXxx, Hooked_Xxx);  // 挂钩
+...
+SlimDetoursInlineHook(FALSE, (PVOID*)&g_pfnXxx, Hooked_Xxx); // 脱钩
 ```
-更多简化的API参考[Wrapper.c](https://github.com/KNSoft/KNSoft.SlimDetours/blob/main/Source/SlimDetours/Wrapper.c)。
+更多简化的API参考[InlineHook.c](https://github.com/KNSoft/KNSoft.SlimDetours/blob/main/Source/SlimDetours/InlineHook.c)。
 
 ### 详细说明
 

diff --git a/Source/Demo/COMHook.c b/Source/Demo/COMHook.c
@@ -59,7 +59,6 @@ TEST_FUNC(COMHook)
     HRESULT hr;
     IOpenControlPanel* pocp1;
     IOpenControlPanel* pocp2;
-    IOpenControlPanel* pocp3;
 
     hr = CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);
     if (FAILED(hr))
@@ -68,21 +67,22 @@ TEST_FUNC(COMHook)
         return;
     }
 
-    hr = CoCreateInstance(&CLSID_OpenControlPanel, NULL, CLSCTX_INPROC_SERVER, &IID_IOpenControlPanel, &pocp1);
+    hr = SlimDetoursCOMHook(&CLSID_OpenControlPanel,
+                            &IID_IOpenControlPanel,
+                            FIELD_OFFSET(IOpenControlPanelVtbl, GetPath),
+                            (PVOID*)&g_pfnIOpenControlPanel_GetPath,
+                            Hooked_IOpenControlPanel_GetPath);
     if (FAILED(hr))
     {
-        TEST_SKIP("CoCreateInstance failed with: 0x%08lX\n", hr);
+        TEST_SKIP("SlimDetoursSetCOMHook failed with: 0x%08lX\n", hr);
         goto _Exit_0;
     }
 
-    hr = SlimDetoursSetTableHook((PVOID*)pocp1->lpVtbl,
-                                 FIELD_OFFSET(IOpenControlPanelVtbl, GetPath),
-                                 (PVOID*)&g_pfnIOpenControlPanel_GetPath,
-                                 Hooked_IOpenControlPanel_GetPath);
+    hr = CoCreateInstance(&CLSID_OpenControlPanel, NULL, CLSCTX_INPROC_SERVER, &IID_IOpenControlPanel, &pocp1);
     if (FAILED(hr))
     {
-        TEST_SKIP("SlimDetoursSetVTableHook failed with: 0x%08lX\n", hr);
-        goto _Exit_1;
+        TEST_SKIP("CoCreateInstance failed with: 0x%08lX\n", hr);
+        goto _Exit_0;
     }
     TEST_OK(Test_IOpenControlPanel_GetPath(pocp1));
 
@@ -94,16 +94,6 @@ TEST_FUNC(COMHook)
     }
     TEST_OK(Test_IOpenControlPanel_GetPath(pocp2));
 
-    hr = CoCreateInstance(&CLSID_OpenControlPanel, NULL, CLSCTX_INPROC_SERVER, &IID_IOpenControlPanel, &pocp3);
-    if (FAILED(hr))
-    {
-        TEST_SKIP("CoCreateInstance failed with: 0x%08lX\n", hr);
-        goto _Exit_2;
-    }
-    TEST_OK(Test_IOpenControlPanel_GetPath(pocp3));
-
-    pocp2->lpVtbl->Release(pocp3);
-_Exit_2:
     pocp2->lpVtbl->Release(pocp2);
 _Exit_1:
     pocp1->lpVtbl->Release(pocp1);

diff --git a/Source/Demo/DelayHook.c b/Source/Demo/DelayHook.c
@@ -11,6 +11,7 @@
 static BOOL g_bDelayAttach = FALSE;
 
 static
+_Function_class_(DETOUR_DELAY_ATTACH_CALLBACK_FN)
 VOID
 CALLBACK
 DelayAttachCallback(

diff --git a/Source/Demo/Main.c b/Source/Demo/Main.c
@@ -80,7 +80,7 @@ Hooked_EqualRect(
 {
     BOOL Ret;
 
-    InterlockedIncrement(&g_lEqualRectCount);
+    _InterlockedIncrement(&g_lEqualRectCount);
     UnitTest_FormatMessage("Hooked EqualRect enter: lprc1 = (%ld, %ld, %ld, %ld), lprc2 = (%ld, %ld, %ld, %ld)\n",
                            lprc1->top,
                            lprc1->right,
@@ -90,9 +90,9 @@ Hooked_EqualRect(
                            lprc2->right,
                            lprc2->bottom,
                            lprc2->left);
-    InterlockedIncrement(&g_lEqualRectRefCount);
+    _InterlockedIncrement(&g_lEqualRectRefCount);
     Ret = g_pfnEqualRect(lprc1, lprc2);
-    InterlockedDecrement(&g_lEqualRectRefCount);
+    _InterlockedDecrement(&g_lEqualRectRefCount);
     UnitTest_FormatMessage("Hooked EqualRect leave with return value: %ld\n", Ret);
     return Ret;
 }

diff --git a/Source/Demo/TwiceSimpleHook.c b/Source/Demo/TwiceSimpleHook.c
@@ -10,6 +10,7 @@ static LONG volatile g_lEqualRect = 0;
 static FN_EqualRect* g_pfnEqualRect1 = NULL;
 static FN_EqualRect* g_pfnEqualRect2 = NULL;
 
+static
 BOOL
 WINAPI
 Hooked_EqualRect1(
@@ -21,6 +22,7 @@ Hooked_EqualRect1(
     return g_pfnEqualRect1(lprc1, lprc2);
 }
 
+static
 BOOL
 WINAPI
 Hooked_EqualRect2(
@@ -46,16 +48,16 @@ TEST_FUNC(TwiceSimpleHook)
     }
 
     g_pfnEqualRect1 = g_pfnEqualRect2 = g_pfnEqualRect;
-    hr = SlimDetoursSetHook((PVOID)&g_pfnEqualRect1, Hooked_EqualRect1);
+    hr = SlimDetoursInlineHook(TRUE, (PVOID)&g_pfnEqualRect1, Hooked_EqualRect1);
     if (FAILED(hr))
     {
-        TEST_SKIP("1st SlimDetoursSetHook failed with 0x%08lX\n", hr);
+        TEST_SKIP("1st SlimDetoursInlineHook failed with 0x%08lX\n", hr);
         return;
     }
-    hr = SlimDetoursSetHook((PVOID)&g_pfnEqualRect2, Hooked_EqualRect2);
+    hr = SlimDetoursInlineHook(TRUE, (PVOID)&g_pfnEqualRect2, Hooked_EqualRect2);
     if (FAILED(hr))
     {
-        TEST_SKIP("2nd SlimDetoursSetHook failed with 0x%08lX\n", hr);
+        TEST_SKIP("2nd SlimDetoursInlineHook failed with 0x%08lX\n", hr);
         return;
     }
 

diff --git a/Source/KNSoft.SlimDetours.vcxproj b/Source/KNSoft.SlimDetours.vcxproj
@@ -207,9 +207,10 @@
     <ClCompile Include="SlimDetours\Thread.c" />
     <ClCompile Include="SlimDetours\Trampoline.c" />
     <ClCompile Include="SlimDetours\Transaction.c" />
-    <ClCompile Include="SlimDetours\Wrapper.c" />
+    <ClCompile Include="SlimDetours\InlineHook.c" />
     <None Include="packages.config" />
     <None Include="SlimDetours\SlimDetours.NDK.inl" />
+    <None Include="SlimDetours\Utils.inl" />
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="Detours\src\detours.h" />

diff --git a/Source/KNSoft.SlimDetours.vcxproj.filters b/Source/KNSoft.SlimDetours.vcxproj.filters
@@ -27,7 +27,7 @@
     <ClCompile Include="SlimDetours\Transaction.c">
       <Filter>SlimDetours</Filter>
     </ClCompile>
-    <ClCompile Include="SlimDetours\Wrapper.c">
+    <ClCompile Include="SlimDetours\InlineHook.c">
       <Filter>SlimDetours</Filter>
     </ClCompile>
     <ClCompile Include="Detours\src\creatwth.cpp">
@@ -86,5 +86,8 @@
       <Filter>SlimDetours</Filter>
     </None>
     <None Include="packages.config" />
+    <None Include="SlimDetours\Utils.inl">
+      <Filter>SlimDetours</Filter>
+    </None>
   </ItemGroup>
 </Project>
diff --git a/Source/SlimDetours/FuncTableHook.c b/Source/SlimDetours/FuncTableHook.c
@@ -1,8 +1,8 @@
 /*
  * KNSoft.SlimDetours (https://github.com/KNSoft/KNSoft.SlimDetours) Function Table Hook Implementation
- * 
- * Hook function address in a read-only table, used by COM/IAT/EAT/... hooking.
- * 
+ *
+ * Hook function address in a read-only table, used by COM/IAT/... hooking.
+ *
  * Copyright (c) KNSoft.org (https://github.com/KNSoft). All rights reserved.
  * Licensed under the MIT license.
  */
@@ -75,11 +75,11 @@ detour_hook_table_funcs(
         Method = Add2Ptr(pFuncTable, pHooks[i].ulOffset);
         if (bEnable)
         {
-            *pHooks[i].ppOriginal = *Method;
-            *Method = pHooks[i].pDetour;
+            *pHooks[i].ppOldFunc = *Method;
+            *Method = pHooks[i].pNewFunc;
         } else
         {
-            *Method = *pHooks[i].ppOriginal;
+            *Method = *pHooks[i].ppOldFunc;
         }
     }
 
@@ -89,32 +89,105 @@ detour_hook_table_funcs(
 
 HRESULT
 NTAPI
-SlimDetoursSetTableHook(
+SlimDetoursFuncTableHook(
     _In_ PVOID* pFuncTable,
     _In_ ULONG ulOffset,
-    _Out_ PVOID* ppOriginal,
-    _In_ PVOID pDetour)
+    _Out_opt_ PVOID* ppOldFunc,
+    _In_ PVOID pNewFunc)
 {
-    return HRESULT_FROM_NT(detour_hook_table_func(pFuncTable, ulOffset, ppOriginal, pDetour));
+    return HRESULT_FROM_NT(detour_hook_table_func(pFuncTable, ulOffset, ppOldFunc, pNewFunc));
 }
 
 HRESULT
 NTAPI
-SlimDetoursUnsetTableHook(
+SlimDetoursFuncTableHooks(
+    _In_ BOOL bEnable,
     _In_ PVOID* pFuncTable,
-    _In_ ULONG ulOffset,
-    _In_ PVOID pOriginal)
+    _In_ ULONG ulCount,
+    _Inout_updates_(ulCount) PDETOUR_FUNC_TABLE_HOOK pHooks)
 {
-    return HRESULT_FROM_NT(detour_hook_table_func(pFuncTable, ulOffset, NULL, pOriginal));
+    return HRESULT_FROM_NT(detour_hook_table_funcs(bEnable, pFuncTable, ulCount, pHooks));
 }
 
+/* COM Hook */
+
+typedef
+_Check_return_
+HRESULT
+STDAPICALLTYPE
+FN_CoCreateInstanceEx(
+    _In_ REFCLSID Clsid,
+    _In_opt_ IUnknown* punkOuter,
+    _In_ DWORD dwClsCtx,
+    _In_opt_ COSERVERINFO* pServerInfo,
+    _In_ DWORD dwCount,
+    _Inout_updates_(dwCount) MULTI_QI* pResults);
+
+static PVOID g_hComBase = NULL;
+
+static CONST UNICODE_STRING g_usCombaseDllName = RTL_CONSTANT_STRING(L"combase.dll");
+static CONST ANSI_STRING g_asCoCreateInstanceEx = RTL_CONSTANT_STRING("CoCreateInstanceEx");
+static FN_CoCreateInstanceEx* g_pfnCoCreateInstanceEx = NULL;
+
+static PS_RUNONCE g_stRunOnceCombaseInit = PS_RUNONCE_INIT;
+static NTSTATUS g_lCombaseInitStatus = STATUS_UNSUCCESSFUL;
+
 HRESULT
 NTAPI
-SlimDetoursEnableTableHooks(
+SlimDetoursCOMHooks(
     _In_ BOOL bEnable,
-    _In_ PVOID* pFuncTable,
+    _In_ REFCLSID rCLSID,
+    _In_ REFCLSID rIID,
     _In_ ULONG ulCount,
     _Inout_updates_(ulCount) PDETOUR_FUNC_TABLE_HOOK pHooks)
 {
-    return HRESULT_FROM_NT(detour_hook_table_funcs(bEnable, pFuncTable, ulCount, pHooks));
+    NTSTATUS Status;
+    HRESULT hr;
+    MULTI_QI MQI = { rIID };
+
+    /* Initialize combase.dll */
+    if (PS_RunOnceBegin(&g_stRunOnceCombaseInit))
+    {
+        Status = LdrLoadDll(NULL, NULL, (PUNICODE_STRING)&g_usCombaseDllName, &g_hComBase);
+        if (!NT_SUCCESS(Status))
+        {
+            goto _Init_Exit;
+        }
+        Status = LdrGetProcedureAddress(g_hComBase, (PANSI_STRING)&g_asCoCreateInstanceEx, 0, (PVOID*)&g_pfnCoCreateInstanceEx);
+        if (!NT_SUCCESS(Status))
+        {
+            LdrUnloadDll(g_hComBase);
+            g_hComBase = NULL;
+            goto _Init_Exit;
+        }
+        Status = STATUS_SUCCESS;
+_Init_Exit:
+        g_lCombaseInitStatus = Status;
+        PS_RunOnceEnd(&g_stRunOnceCombaseInit, Status == STATUS_SUCCESS);
+    }
+    if (!NT_SUCCESS(g_lCombaseInitStatus))
+    {
+        return HRESULT_FROM_NT(g_lCombaseInitStatus);
+    }
+
+    /* Create COM Object and set VTable Hooks */
+    hr = g_pfnCoCreateInstanceEx(rCLSID, NULL, CLSCTX_ALL, NULL, 1, &MQI);
+    if (FAILED(hr))
+    {
+        goto _Create_Fail_0;
+    }
+    _Analysis_assume_(MQI.pItf != NULL);
+    Status = detour_hook_table_funcs(bEnable, (PVOID*)MQI.pItf->lpVtbl, ulCount, pHooks);
+    if (!NT_SUCCESS(Status))
+    {
+        hr = HRESULT_FROM_NT(Status);
+        goto _Create_Fail_1;
+    }
+    hr = S_OK;
+
+    /* Cleanup */
+_Create_Fail_1:
+    MQI.pItf->lpVtbl->Release(MQI.pItf);
+_Create_Fail_0:
+    return hr;
 }