From 6602def1108ea647eeaf74e7dde497f48d1a6b85 Mon Sep 17 00:00:00 2001
From: g00356712 <guochenyang@huawei.com>
Date: Sun, 20 Nov 2022 16:31:53 +0800
Subject: [PATCH] fix context setting errors in DetourEnumerateImports API

---
 src/modules.cpp           | 19 ++++++++++++++++---
 tests/test_module_api.cpp | 18 +++++++++++++++++-
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/src/modules.cpp b/src/modules.cpp
index 52fc2cfb..a781e4ad 100644
--- a/src/modules.cpp
+++ b/src/modules.cpp
@@ -645,6 +645,7 @@ BOOL WINAPI DetourEnumerateImportsEx(_In_opt_ HMODULE hModule,
 struct _DETOUR_ENUMERATE_IMPORTS_THUNK_CONTEXT
 {
     PVOID pContext;
+    PF_DETOUR_IMPORT_FILE_CALLBACK pfImportFile;
     PF_DETOUR_IMPORT_FUNC_CALLBACK pfImportFunc;
 };
 
@@ -664,6 +665,19 @@ DetourEnumerateImportsThunk(_In_ PVOID VoidContext,
     return pContext->pfImportFunc(pContext->pContext, nOrdinal, pszFunc, ppvFunc ? *ppvFunc : NULL);
 }
 
+static
+BOOL
+CALLBACK
+DetourEnumerateImportsFile(_In_opt_ PVOID VoidContext,
+                            _In_opt_ HMODULE hModule,
+                            _In_opt_ LPCSTR pszFile)
+{
+    _DETOUR_ENUMERATE_IMPORTS_THUNK_CONTEXT const * const
+        pContext = (_DETOUR_ENUMERATE_IMPORTS_THUNK_CONTEXT*)VoidContext;
+    return pContext->pfImportFile(pContext->pContext, hModule, pszFile);
+}
+
+
 BOOL WINAPI DetourEnumerateImports(_In_opt_ HMODULE hModule,
                                    _In_opt_ PVOID pContext,
                                    _In_opt_ PF_DETOUR_IMPORT_FILE_CALLBACK pfImportFile,
@@ -674,11 +688,10 @@ BOOL WINAPI DetourEnumerateImports(_In_opt_ HMODULE hModule,
         return FALSE;
     }
 
-    _DETOUR_ENUMERATE_IMPORTS_THUNK_CONTEXT const context = { pContext, pfImportFunc };
-
+    _DETOUR_ENUMERATE_IMPORTS_THUNK_CONTEXT const context = { pContext, pfImportFile, pfImportFunc };
     return DetourEnumerateImportsEx(hModule,
                                     (PVOID)&context,
-                                    pfImportFile,
+                                    &DetourEnumerateImportsFile,
                                     &DetourEnumerateImportsThunk);
 }
 
diff --git a/tests/test_module_api.cpp b/tests/test_module_api.cpp
index 3e2928cd..fee84bd6 100644
--- a/tests/test_module_api.cpp
+++ b/tests/test_module_api.cpp
@@ -409,7 +409,9 @@ BOOL WINAPI ImportFileCallback(PVOID pContext, HMODULE, PCSTR pszFile)
         reinterpret_cast<EnumerateImportsTestContext*>(pContext);
 
     context->ImportCount++;
-    context->ImportModuleFound |= Catch::contains(pszFile, "ntdll");
+    if (pszFile != NULL) {
+        context->ImportModuleFound |= Catch::contains(pszFile, "ntdll");
+    }
 
     return TRUE;
 }
@@ -471,6 +473,20 @@ TEST_CASE("DetourEnumerateImports", "[module]")
         REQUIRE( context.ImportFuncCount == 0 );
         REQUIRE_FALSE( context.ImportFuncFound );
     }
+
+    SECTION("The context transferred during the input parameter is the same as the context parsed in the callback.")
+    {
+        SetLastError(ERROR_INVALID_HANDLE);
+
+        EnumerateImportsTestContext context {};
+        auto success = DetourEnumerateImports(NULL, &context, ImportFileCallback, ImportFuncCallback);
+
+        REQUIRE( GetLastError() == 0 );
+        REQUIRE( success  == true );
+
+        REQUIRE( context.ImportFuncCount != 0 );
+        REQUIRE( context.ImportCount != 0 );
+    }
 }
 
 TEST_CASE("DetourGetSizeOfPayloads", "[module]")