Minor improvements

KNSoft · Feb 18, 2024 · bfd5da9 · bfd5da9
1 parent 24f39db
commit bfd5da9
Show file tree

Hide file tree

Showing 8 changed files with 32 additions and 47 deletions.
diff --git a/Source/SlimDetours/Disasm.c b/Source/SlimDetours/Disasm.c
@@ -1579,21 +1579,6 @@ pp is like VEX but only instructions with 0 are defined
     }
 }
 
-PVOID NTAPI SlimDetoursCopyInstruction(
-    _In_opt_ PVOID pDst,
-    _Inout_opt_ PVOID* ppDstPool,
-    _In_ PVOID pSrc,
-    _Out_opt_ PVOID* ppTarget,
-    _Out_opt_ LONG* plExtra)
-{
-    UNREFERENCED_PARAMETER(ppDstPool);  // x86 & x64 don't use a constant pool.
-
-    DETOUR_DISASM Disasm;
-
-    detour_disasm_init(&Disasm, (PBYTE*)ppTarget, plExtra);
-    return CopyInstruction(&Disasm, (PBYTE)pDst, (PBYTE)pSrc);
-}
-
 #endif // defined(_M_X64) || defined(_M_IX86)
 
 #if defined(_M_ARM64)
@@ -2324,23 +2309,27 @@ static PBYTE CopyInstruction(_In_ PDETOUR_DISASM pDisasm, _In_opt_ PBYTE pDst, _
     return pSrc + 4;
 }
 
+#endif // defined(_M_ARM64)
+
 PVOID NTAPI SlimDetoursCopyInstruction(
     _In_opt_ PVOID pDst,
-    _Inout_opt_ PVOID* ppDstPool,
     _In_ PVOID pSrc,
     _Out_opt_ PVOID* ppTarget,
     _Out_opt_ LONG* plExtra)
 {
-    UNREFERENCED_PARAMETER(ppDstPool);
-
     DETOUR_DISASM Disasm;
 
+#if defined(_M_X64) || defined(_M_IX86)
+    detour_disasm_init(&Disasm, (PBYTE*)ppTarget, plExtra);
+    return (PVOID)CopyInstruction(&Disasm, (PBYTE)pDst, (PBYTE)pSrc);
+#elif defined(_M_ARM64)
     detour_disasm_init(&Disasm);
     return (PVOID)CopyInstruction(&Disasm,
                                   (PBYTE)pDst,
                                   (PBYTE)pSrc,
                                   (PBYTE*)ppTarget,
                                   plExtra);
+#else
+    return NULL;
+#endif
 }
-
-#endif // defined(_M_ARM64)
diff --git a/Source/SlimDetours/Memory.c b/Source/SlimDetours/Memory.c
@@ -79,13 +79,16 @@ VOID detour_memory_init()
     }
 #endif
 
-    g_hHeap = RtlCreateHeap(HEAP_NO_SERIALIZE | HEAP_GROWABLE, NULL, 0, 0, NULL, NULL);
+    if (g_hHeap == NULL)
+    {
+        g_hHeap = RtlCreateHeap(HEAP_NO_SERIALIZE | HEAP_GROWABLE, NULL, 0, 0, NULL, NULL);
+    }
 }
 
 _Must_inspect_result_
 _Ret_maybenull_
 _Post_writable_byte_size_(Size)
-PVOID detour_memory_alloc(_In_range_(>, 0) SIZE_T Size)
+PVOID detour_memory_alloc(_In_ SIZE_T Size)
 {
     if (g_hHeap == NULL)
     {

diff --git a/Source/SlimDetours/SlimDetours.h b/Source/SlimDetours/SlimDetours.h
@@ -52,7 +52,6 @@ NTSTATUS NTAPI SlimDetoursDelayAttach(
 PVOID NTAPI SlimDetoursCodeFromPointer(_In_ PVOID pPointer);
 PVOID NTAPI SlimDetoursCopyInstruction(
     _In_opt_ PVOID pDst,
-    _Inout_opt_ PVOID* ppDstPool,
     _In_ PVOID pSrc,
     _Out_opt_ PVOID* ppTarget,
     _Out_opt_ LONG* plExtra);

diff --git a/Source/SlimDetours/SlimDetours.inl b/Source/SlimDetours/SlimDetours.inl
@@ -124,7 +124,7 @@ VOID detour_memory_init();
 _Must_inspect_result_
 _Ret_maybenull_
 _Post_writable_byte_size_(Size)
-PVOID detour_memory_alloc(_In_range_(> , 0) SIZE_T Size);
+PVOID detour_memory_alloc(_In_ SIZE_T Size);
 
 BOOL detour_memory_free(_Frees_ptr_ PVOID BaseAddress);
 

diff --git a/Source/SlimDetours/Thread.c b/Source/SlimDetours/Thread.c
@@ -139,7 +139,6 @@ NTSTATUS detour_thread_update(_In_ HANDLE ThreadHandle, _In_ PDETOUR_OPERATION P
             if (cxt.CONTEXT_PC >= (ULONG_PTR)o->pTrampoline &&
                 cxt.CONTEXT_PC < ((ULONG_PTR)o->pTrampoline + sizeof(o->pTrampoline)))
             {
-
                 cxt.CONTEXT_PC = (ULONG_PTR)o->pbTarget +
                     detour_align_from_trampoline(o->pTrampoline, (BYTE)(cxt.CONTEXT_PC - (ULONG_PTR)o->pTrampoline));
                 bUpdateContext = TRUE;
@@ -149,7 +148,6 @@ NTSTATUS detour_thread_update(_In_ HANDLE ThreadHandle, _In_ PDETOUR_OPERATION P
             if (cxt.CONTEXT_PC >= (ULONG_PTR)o->pbTarget &&
                 cxt.CONTEXT_PC < ((ULONG_PTR)o->pbTarget + o->pTrampoline->cbRestore))
             {
-
                 cxt.CONTEXT_PC = (ULONG_PTR)o->pTrampoline +
                     detour_align_from_target(o->pTrampoline, (BYTE)(cxt.CONTEXT_PC - (ULONG_PTR)o->pbTarget));
                 bUpdateContext = TRUE;

diff --git a/Source/SlimDetours/Transaction.c b/Source/SlimDetours/Transaction.c
@@ -323,7 +323,7 @@ NTSTATUS NTAPI SlimDetoursAttach(_Inout_ PVOID* ppPointer, _In_ PVOID pDetour)
         LONG lExtra = 0;
 
         DETOUR_TRACE(" SlimDetoursCopyInstruction(%p,%p)\n", pbTrampoline, pbSrc);
-        pbSrc = (PBYTE)SlimDetoursCopyInstruction(pbTrampoline, (PVOID*)&pbPool, pbSrc, NULL, &lExtra);
+        pbSrc = (PBYTE)SlimDetoursCopyInstruction(pbTrampoline, pbSrc, NULL, &lExtra);
         DETOUR_TRACE(" SlimDetoursCopyInstruction() = %p (%d bytes)\n", pbSrc, (int)(pbSrc - pbOp));
         pbTrampoline += (pbSrc - pbOp) + lExtra;
         cbTarget = PtrOffset(pbTarget, pbSrc);
@@ -588,7 +588,7 @@ static VOID CALLBACK detour_dll_notify_proc(
     while (pAttach != NULL)
     {
         /* Match Dll name */
-        if (RtlCompareUnicodeString(&pAttach->usDllName, NotificationData->Loaded.BaseDllName, FALSE) != 0)
+        if (!RtlEqualUnicodeString(&pAttach->usDllName, NotificationData->Loaded.BaseDllName, FALSE))
         {
             pPrevAttach = pAttach;
             pAttach = pAttach->pNext;
@@ -620,7 +620,6 @@ static VOID CALLBACK detour_dll_notify_proc(
             g_DelayedAttaches = pNextAttach;
         }
         pAttach = pNextAttach;
-        continue;
     }
 
     RtlReleaseSRWLockExclusive(&g_DelayedAttachesLock);

diff --git a/Source/Test/Main.cpp b/Source/Test/Main.cpp
@@ -49,26 +49,23 @@ static VOID CALLBACK Delay_attach_proc(
     }
 }
 
-int WINAPI wWinMain(
-    _In_     HINSTANCE hInstance,
-    _In_opt_ HINSTANCE hPrevInstance,
-    _In_     LPWSTR    lpCmdLine,
-    _In_     int       nShowCmd)
+int wmain()
 {
     NTSTATUS Status;
     PVOID User32Base;
 
 #if 1 // Test delay attach
     FN_MessageBoxW* pfnMessageBoxW;
 
-    if (!NT_SUCCESS(SlimDetoursDelayAttach((PVOID*)&g_pfnMessageBoxW,
-                                           Hooked_MessageBoxW,
-                                           g_usUser32.Buffer,
-                                           g_asMessageBoxW.Buffer,
-                                           Delay_attach_proc,
-                                           NULL)))
+    Status = SlimDetoursDelayAttach((PVOID*)&g_pfnMessageBoxW,
+                                    Hooked_MessageBoxW,
+                                    g_usUser32.Buffer,
+                                    g_asMessageBoxW.Buffer,
+                                    Delay_attach_proc,
+                                    NULL);
+    if (!NT_SUCCESS(Status))
     {
-        return STATUS_UNSUCCESSFUL;
+        return Status;
     }
 
     if (!NT_SUCCESS(LdrLoadDll(NULL, NULL, &g_usUser32, &User32Base)) ||

diff --git a/Source/Test/Test.vcxproj b/Source/Test/Test.vcxproj
@@ -119,7 +119,7 @@
       <OmitDefaultLibName>true</OmitDefaultLibName>
     </ClCompile>
     <Link>
-      <SubSystem>Windows</SubSystem>
+      <SubSystem>Console</SubSystem>
     </Link>
     <ProjectReference />
   </ItemDefinitionGroup>
@@ -130,7 +130,7 @@
       <OmitDefaultLibName>true</OmitDefaultLibName>
     </ClCompile>
     <Link>
-      <SubSystem>Windows</SubSystem>
+      <SubSystem>Console</SubSystem>
     </Link>
     <ProjectReference />
   </ItemDefinitionGroup>
@@ -143,7 +143,7 @@
       <OmitDefaultLibName>true</OmitDefaultLibName>
     </ClCompile>
     <Link>
-      <SubSystem>Windows</SubSystem>
+      <SubSystem>Console</SubSystem>
       <EnableCOMDATFolding>true</EnableCOMDATFolding>
       <OptimizeReferences>true</OptimizeReferences>
     </Link>
@@ -158,7 +158,7 @@
       <OmitDefaultLibName>true</OmitDefaultLibName>
     </ClCompile>
     <Link>
-      <SubSystem>Windows</SubSystem>
+      <SubSystem>Console</SubSystem>
       <EnableCOMDATFolding>true</EnableCOMDATFolding>
       <OptimizeReferences>true</OptimizeReferences>
     </Link>
@@ -171,7 +171,7 @@
       <OmitDefaultLibName>true</OmitDefaultLibName>
     </ClCompile>
     <Link>
-      <SubSystem>Windows</SubSystem>
+      <SubSystem>Console</SubSystem>
     </Link>
     <ProjectReference />
   </ItemDefinitionGroup>
@@ -184,7 +184,7 @@
       <OmitDefaultLibName>true</OmitDefaultLibName>
     </ClCompile>
     <Link>
-      <SubSystem>Windows</SubSystem>
+      <SubSystem>Console</SubSystem>
       <EnableCOMDATFolding>true</EnableCOMDATFolding>
       <OptimizeReferences>true</OptimizeReferences>
     </Link>