rostests/apitests/include/apitest_iathook.h

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