1*c2c66affSColin Finck #ifndef _APITEST_IATHOOK_H
2*c2c66affSColin Finck #define _APITEST_IATHOOK_H
3*c2c66affSColin Finck 
4*c2c66affSColin Finck static PIMAGE_IMPORT_DESCRIPTOR FindImportDescriptor(PBYTE DllBase, PCSTR DllName)
5*c2c66affSColin Finck {
6*c2c66affSColin Finck     ULONG Size;
7*c2c66affSColin Finck     PIMAGE_IMPORT_DESCRIPTOR ImportDescriptor = RtlImageDirectoryEntryToData((HMODULE)DllBase, TRUE, IMAGE_DIRECTORY_ENTRY_IMPORT, &Size);
8*c2c66affSColin Finck     while (ImportDescriptor->Name && ImportDescriptor->OriginalFirstThunk)
9*c2c66affSColin Finck     {
10*c2c66affSColin Finck         PCHAR Name = (PCHAR)(DllBase + ImportDescriptor->Name);
11*c2c66affSColin Finck         if (!lstrcmpiA(Name, DllName))
12*c2c66affSColin Finck         {
13*c2c66affSColin Finck             return ImportDescriptor;
14*c2c66affSColin Finck         }
15*c2c66affSColin Finck         ImportDescriptor++;
16*c2c66affSColin Finck     }
17*c2c66affSColin Finck     return NULL;
18*c2c66affSColin Finck }
19*c2c66affSColin Finck 
20*c2c66affSColin Finck static BOOL RedirectIat(HMODULE TargetDll, PCSTR DllName, PCSTR FunctionName, ULONG_PTR NewFunction, ULONG_PTR* OriginalFunction)
21*c2c66affSColin Finck {
22*c2c66affSColin Finck     PBYTE DllBase = (PBYTE)TargetDll;
23*c2c66affSColin Finck     PIMAGE_IMPORT_DESCRIPTOR ImportDescriptor = FindImportDescriptor(DllBase, DllName);
24*c2c66affSColin Finck     if (ImportDescriptor)
25*c2c66affSColin Finck     {
26*c2c66affSColin Finck         // On loaded images, OriginalFirstThunk points to the name / ordinal of the function
27*c2c66affSColin Finck         PIMAGE_THUNK_DATA OriginalThunk = (PIMAGE_THUNK_DATA)(DllBase + ImportDescriptor->OriginalFirstThunk);
28*c2c66affSColin Finck         // FirstThunk points to the resolved address.
29*c2c66affSColin Finck         PIMAGE_THUNK_DATA FirstThunk = (PIMAGE_THUNK_DATA)(DllBase + ImportDescriptor->FirstThunk);
30*c2c66affSColin Finck         while (OriginalThunk->u1.AddressOfData && FirstThunk->u1.Function)
31*c2c66affSColin Finck         {
32*c2c66affSColin Finck             if (!IMAGE_SNAP_BY_ORDINAL32(OriginalThunk->u1.AddressOfData))
33*c2c66affSColin Finck             {
34*c2c66affSColin Finck                 PIMAGE_IMPORT_BY_NAME ImportName = (PIMAGE_IMPORT_BY_NAME)(DllBase + OriginalThunk->u1.AddressOfData);
35*c2c66affSColin Finck                 if (!lstrcmpiA((PCSTR)ImportName->Name, FunctionName))
36*c2c66affSColin Finck                 {
37*c2c66affSColin Finck                     DWORD dwOld;
38*c2c66affSColin Finck                     VirtualProtect(&FirstThunk->u1.Function, sizeof(ULONG_PTR), PAGE_EXECUTE_READWRITE, &dwOld);
39*c2c66affSColin Finck                     *OriginalFunction = FirstThunk->u1.Function;
40*c2c66affSColin Finck                     FirstThunk->u1.Function = NewFunction;
41*c2c66affSColin Finck                     VirtualProtect(&FirstThunk->u1.Function, sizeof(ULONG_PTR), dwOld, &dwOld);
42*c2c66affSColin Finck                     return TRUE;
43*c2c66affSColin Finck                 }
44*c2c66affSColin Finck             }
45*c2c66affSColin Finck             OriginalThunk++;
46*c2c66affSColin Finck             FirstThunk++;
47*c2c66affSColin Finck         }
48*c2c66affSColin Finck         skip("Unable to find the Import %s!%s\n", DllName, FunctionName);
49*c2c66affSColin Finck     }
50*c2c66affSColin Finck     else
51*c2c66affSColin Finck     {
52*c2c66affSColin Finck         skip("Unable to find the ImportDescriptor for %s\n", DllName);
53*c2c66affSColin Finck     }
54*c2c66affSColin Finck     return FALSE;
55*c2c66affSColin Finck }
56*c2c66affSColin Finck 
57*c2c66affSColin Finck static BOOL RestoreIat(HMODULE TargetDll, PCSTR DllName, PCSTR FunctionName, ULONG_PTR OriginalFunction)
58*c2c66affSColin Finck {
59*c2c66affSColin Finck     ULONG_PTR old = 0;
60*c2c66affSColin Finck     return RedirectIat(TargetDll, DllName, FunctionName, OriginalFunction, &old);
61*c2c66affSColin Finck }
62*c2c66affSColin Finck 
63*c2c66affSColin Finck  #endif // _APITEST_IATHOOK_H
64*c2c66affSColin Finck