dll/appcompat/apphelp/shimeng.c

   1 /*
   2  * PROJECT:     ReactOS Application compatibility module
   3  * LICENSE:     GPL-2.0+ (https://spdx.org/licenses/GPL-2.0+)
   4  * PURPOSE:     Shim engine core
   5  * COPYRIGHT:   Copyright 2015-2018 Mark Jansen (mark.jansen@reactos.org)
   6  */
   7
   8 #define WIN32_NO_STATUS
   9 #include "ntndk.h"
  10 #define IN_APPHELP
  11 #include "shimlib.h"
  12 #include <strsafe.h>
  13 /* Make sure we don't include apphelp logging */
  14 #define APPHELP_NOSDBPAPI
  15 #include "apphelp.h"
  16 #include "shimeng.h"
  17
  18
  19
  20 FARPROC WINAPI StubGetProcAddress(HINSTANCE hModule, LPCSTR lpProcName);
  21 BOOL WINAPI SE_IsShimDll(PVOID BaseAddress);
  22
  23
  24 extern HMODULE g_hInstance;
  25 static UNICODE_STRING g_WindowsDirectory;
  26 static UNICODE_STRING g_System32Directory;
  27 static UNICODE_STRING g_SxsDirectory;
  28 ULONG g_ShimEngDebugLevel = 0xffffffff;
  29 BOOL g_bComPlusImage = FALSE;
  30 BOOL g_bShimDuringInit = FALSE;
  31 BOOL g_bInternalHooksUsed = FALSE;
  32 static ARRAY g_pShimInfo;   /* PSHIMMODULE */
  33 static ARRAY g_pHookArray;  /* HOOKMODULEINFO */
  34 static ARRAY g_InExclude;   /* INEXCLUDE */
  35
  36 /* If we have setup a hook for a function, we should also redirect GetProcAddress for this function */
  37 HOOKAPIEX g_IntHookEx[] =
  38 {
  39     {
  40         "kernel32.dll",     /* LibraryName */
  41         "GetProcAddress",   /* FunctionName */
  42         StubGetProcAddress, /* ReplacementFunction*/
  43         NULL,               /* OriginalFunction */
  44         NULL,               /* pShimInfo */
  45         NULL                /* Unused */
  46     },
  47 };
  48
  49 static inline BOOL ARRAY_InitWorker(PARRAY Array, DWORD ItemSize)
  50 {
  51     Array->Data__ = NULL;
  52     Array->Size__ = Array->MaxSize__ = 0;
  53     Array->ItemSize__ = ItemSize;
  54
  55     return TRUE;
  56 }
  57
  58 static inline BOOL ARRAY_EnsureSize(PARRAY Array, DWORD ItemSize, DWORD GrowWith)
  59 {
  60     PVOID pNewData;
  61     DWORD Count;
  62
  63     ASSERT(Array);
  64     ASSERT(ItemSize == Array->ItemSize__);
  65
  66     if (Array->MaxSize__ > Array->Size__)
  67         return TRUE;
  68
  69     Count = Array->Size__ + GrowWith;
  70     pNewData = SeiAlloc(Count * ItemSize);
  71
  72     if (!pNewData)
  73     {
  74         SHIMENG_FAIL("Failed to allocate %d bytes\n", Count * ItemSize);
  75         return FALSE;
  76     }
  77     Array->MaxSize__ = Count;
  78
  79     if (Array->Data__)
  80     {
  81         memcpy(pNewData, Array->Data__, Array->Size__ * ItemSize);
  82         SeiFree(Array->Data__);
  83     }
  84     Array->Data__ = pNewData;
  85
  86     return TRUE;
  87 }
  88
  89 static inline PVOID ARRAY_AppendWorker(PARRAY Array, DWORD ItemSize, DWORD GrowWith)
  90 {
  91     PBYTE pData;
  92
  93     if (!ARRAY_EnsureSize(Array, ItemSize, GrowWith))
  94         return NULL;
  95
  96     pData = Array->Data__;
  97     pData += (Array->Size__ * ItemSize);
  98     Array->Size__++;
  99
 100     return pData;
 101 }
 102
 103 static inline PVOID ARRAY_AtWorker(PARRAY Array, DWORD ItemSize, DWORD n)
 104 {
 105     PBYTE pData;
 106
 107     ASSERT(Array);
 108     ASSERT(ItemSize == Array->ItemSize__);
 109     ASSERT(n < Array->Size__);
 110
 111     pData = Array->Data__;
 112     return pData + (n * ItemSize);
 113 }
 114
 115
 116 #define ARRAY_Init(Array, TypeOfArray)      ARRAY_InitWorker((Array), sizeof(TypeOfArray))
 117 #define ARRAY_Append(Array, TypeOfArray)    (TypeOfArray*)ARRAY_AppendWorker((Array), sizeof(TypeOfArray), 5)
 118 #define ARRAY_At(Array, TypeOfArray, at)    (TypeOfArray*)ARRAY_AtWorker((Array), sizeof(TypeOfArray), at)
 119 #define ARRAY_Size(Array)                   (Array)->Size__
 120
 121
 122 VOID SeiInitDebugSupport(VOID)
 123 {
 124     static const UNICODE_STRING DebugKey = RTL_CONSTANT_STRING(L"SHIMENG_DEBUG_LEVEL");
 125     UNICODE_STRING DebugValue;
 126     NTSTATUS Status;
 127     ULONG NewLevel = SEI_MSG;   /* Show some basic info in the logs, unless configured different */
 128     WCHAR Buffer[40];
 129
 130     RtlInitEmptyUnicodeString(&DebugValue, Buffer, sizeof(Buffer));
 131
 132     Status = RtlQueryEnvironmentVariable_U(NULL, &DebugKey, &DebugValue);
 133
 134     if (NT_SUCCESS(Status))
 135     {
 136         if (!NT_SUCCESS(RtlUnicodeStringToInteger(&DebugValue, 10, &NewLevel)))
 137             NewLevel = 0;
 138     }
 139     g_ShimEngDebugLevel = NewLevel;
 140 }
 141
 142
 143 /**
 144  * Outputs diagnostic info.
 145  *
 146  * @param [in]  Level           The level to log this message with, choose any of [SHIM_ERR,
 147  *                              SHIM_WARN, SHIM_INFO].
 148  * @param [in]  FunctionName    The function this log should be attributed to.
 149  * @param [in]  Format          The format string.
 150  * @param   ...                 Variable arguments providing additional information.
 151  *
 152  * @return  Success: TRUE Failure: FALSE.
 153  */
 154 BOOL WINAPIV SeiDbgPrint(SEI_LOG_LEVEL Level, PCSTR Function, PCSTR Format, ...)
 155 {
 156     char Buffer[512];
 157     char* Current = Buffer;
 158     const char* LevelStr;
 159     size_t Length = sizeof(Buffer);
 160     va_list ArgList;
 161     HRESULT hr;
 162
 163     if (g_ShimEngDebugLevel == 0xffffffff)
 164         SeiInitDebugSupport();
 165
 166     if (Level > g_ShimEngDebugLevel)
 167         return FALSE;
 168
 169     switch (Level)
 170     {
 171     case SEI_MSG:
 172         LevelStr = "MSG ";
 173         break;
 174     case SEI_FAIL:
 175         LevelStr = "FAIL";
 176         break;
 177     case SEI_WARN:
 178         LevelStr = "WARN";
 179         break;
 180     case SEI_INFO:
 181         LevelStr = "INFO";
 182         break;
 183     default:
 184         LevelStr = "USER";
 185         break;
 186     }
 187
 188     if (Function)
 189         hr = StringCchPrintfExA(Current, Length, &Current, &Length, STRSAFE_NULL_ON_FAILURE, "[%s] [%s] ", LevelStr, Function);
 190     else
 191         hr = StringCchPrintfExA(Current, Length, &Current, &Length, STRSAFE_NULL_ON_FAILURE, "[%s] ", LevelStr);
 192
 193     if (!SUCCEEDED(hr))
 194         return FALSE;
 195
 196     va_start(ArgList, Format);
 197     hr = StringCchVPrintfExA(Current, Length, &Current, &Length, STRSAFE_NULL_ON_FAILURE, Format, ArgList);
 198     va_end(ArgList);
 199     if (!SUCCEEDED(hr))
 200         return FALSE;
 201
 202     DbgPrint("%s", Buffer);
 203     return TRUE;
 204 }
 205
 206
 207 PVOID SeiGetModuleFromAddress(PVOID addr)
 208 {
 209     PVOID hModule = NULL;
 210     RtlPcToFileHeader(addr, &hModule);
 211     return hModule;
 212 }
 213
 214
 215
 216 /* TODO: Guard against recursive calling / calling init multiple times! */
 217 VOID NotifyShims(DWORD dwReason, PVOID Info)
 218 {
 219     DWORD n;
 220
 221     for (n = 0; n < ARRAY_Size(&g_pShimInfo); ++n)
 222     {
 223         PSHIMMODULE pShimModule = *ARRAY_At(&g_pShimInfo, PSHIMMODULE, n);
 224         if (!pShimModule->pNotifyShims)
 225             continue;
 226
 227         pShimModule->pNotifyShims(dwReason, Info);
 228     }
 229 }
 230
 231
 232
 233 VOID SeiCheckComPlusImage(PVOID BaseAddress)
 234 {
 235     ULONG ComSectionSize;
 236     g_bComPlusImage = RtlImageDirectoryEntryToData(BaseAddress, TRUE, IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR, &ComSectionSize) != NULL;
 237
 238     SHIMENG_INFO("COM+ executable %s\n", g_bComPlusImage ? "TRUE" : "FALSE");
 239 }
 240
 241
 242 PSHIMMODULE SeiGetShimModuleInfo(PVOID BaseAddress)
 243 {
 244     DWORD n;
 245
 246     for (n = 0; n < ARRAY_Size(&g_pShimInfo); ++n)
 247     {
 248         PSHIMMODULE pShimModule = *ARRAY_At(&g_pShimInfo, PSHIMMODULE, n);
 249
 250         if (pShimModule->BaseAddress == BaseAddress)
 251             return pShimModule;
 252     }
 253     return NULL;
 254 }
 255
 256 PSHIMMODULE SeiCreateShimModuleInfo(PCWSTR DllName, PVOID BaseAddress)
 257 {
 258     static const ANSI_STRING GetHookAPIs = RTL_CONSTANT_STRING("GetHookAPIs");
 259     static const ANSI_STRING NotifyShims = RTL_CONSTANT_STRING("NotifyShims");
 260     PSHIMMODULE* pData, Data;
 261     PVOID pGetHookAPIs, pNotifyShims;
 262
 263     if (!NT_SUCCESS(LdrGetProcedureAddress(BaseAddress, (PANSI_STRING)&GetHookAPIs, 0, &pGetHookAPIs)) ||
 264         !NT_SUCCESS(LdrGetProcedureAddress(BaseAddress, (PANSI_STRING)&NotifyShims, 0, &pNotifyShims)))
 265     {
 266         SHIMENG_WARN("Failed to resolve entry points for %S\n", DllName);
 267         return NULL;
 268     }
 269
 270     pData = ARRAY_Append(&g_pShimInfo, PSHIMMODULE);
 271     if (!pData)
 272         return NULL;
 273
 274     *pData = SeiAlloc(sizeof(SHIMMODULE));
 275
 276     Data = *pData;
 277
 278     RtlCreateUnicodeString(&Data->Name, DllName);
 279     Data->BaseAddress = BaseAddress;
 280
 281     Data->pGetHookAPIs = pGetHookAPIs;
 282     Data->pNotifyShims = pNotifyShims;
 283
 284     ARRAY_Init(&Data->EnabledShims, PSHIMINFO);
 285
 286     return Data;
 287 }
 288
 289 PSHIMINFO SeiAppendHookInfo(PSHIMMODULE pShimModuleInfo, PHOOKAPIEX pHookApi, DWORD dwHookCount, PCWSTR ShimName)
 290 {
 291     PSHIMINFO* pData, Data;
 292
 293     pData = ARRAY_Append(&pShimModuleInfo->EnabledShims, PSHIMINFO);
 294     if (!pData)
 295         return NULL;
 296
 297     *pData = SeiAlloc(sizeof(SHIMINFO));
 298     Data = *pData;
 299
 300     if (!Data)
 301         return NULL;
 302
 303     Data->ShimName = SdbpStrDup(ShimName);
 304     if (!Data->ShimName)
 305         return NULL;
 306
 307     Data->pHookApi = pHookApi;
 308     Data->dwHookCount = dwHookCount;
 309     Data->pShimModule = pShimModuleInfo;
 310     ARRAY_Init(&Data->InExclude, INEXCLUDE);
 311     return Data;
 312 }
 313
 314 PHOOKMODULEINFO SeiFindHookModuleInfo(PUNICODE_STRING ModuleName, PVOID BaseAddress)
 315 {
 316     DWORD n;
 317
 318     for (n = 0; n < ARRAY_Size(&g_pHookArray); ++n)
 319     {
 320         PHOOKMODULEINFO pModuleInfo = ARRAY_At(&g_pHookArray, HOOKMODULEINFO, n);
 321
 322         if (BaseAddress && BaseAddress == pModuleInfo->BaseAddress)
 323             return pModuleInfo;
 324
 325         if (!BaseAddress && RtlEqualUnicodeString(ModuleName, &pModuleInfo->Name, TRUE))
 326             return pModuleInfo;
 327     }
 328
 329     return NULL;
 330 }
 331
 332 PHOOKMODULEINFO SeiFindHookModuleInfoForImportDescriptor(PBYTE DllBase, PIMAGE_IMPORT_DESCRIPTOR ImportDescriptor)
 333 {
 334     UNICODE_STRING DllName;
 335     PVOID DllHandle;
 336     NTSTATUS Success;
 337
 338     if (!RtlCreateUnicodeStringFromAsciiz(&DllName, (PCSZ)(DllBase + ImportDescriptor->Name)))
 339     {
 340         SHIMENG_FAIL("Unable to convert dll name to unicode\n");
 341         return NULL;
 342     }
 343
 344     Success = LdrGetDllHandle(NULL, NULL, &DllName, &DllHandle);
 345     RtlFreeUnicodeString(&DllName);
 346
 347     if (!NT_SUCCESS(Success))
 348     {
 349         SHIMENG_FAIL("Unable to get module handle for %wZ\n", &DllName);
 350         return NULL;
 351     }
 352
 353     return SeiFindHookModuleInfo(NULL, DllHandle);
 354 }
 355
 356 static LPCWSTR SeiGetStringPtr(PDB pdb, TAGID tag, TAG type)
 357 {
 358     TAGID tagEntry = SdbFindFirstTag(pdb, tag, type);
 359     if (tagEntry == TAGID_NULL)
 360         return NULL;
 361
 362     return SdbGetStringTagPtr(pdb, tagEntry);
 363 }
 364
 365 static DWORD SeiGetDWORD(PDB pdb, TAGID tag, TAG type)
 366 {
 367     TAGID tagEntry = SdbFindFirstTag(pdb, tag, type);
 368     if (tagEntry == TAGID_NULL)
 369         return 0;
 370
 371     return SdbReadDWORDTag(pdb, tagEntry, 0);
 372 }
 373
 374 static QWORD SeiGetQWORD(PDB pdb, TAGID tag, TAG type)
 375 {
 376     TAGID tagEntry = SdbFindFirstTag(pdb, tag, type);
 377     if (tagEntry == TAGID_NULL)
 378         return 0;
 379
 380     return SdbReadQWORDTag(pdb, tagEntry, 0);
 381 }
 382
 383 static VOID SeiAddShim(TAGREF trShimRef, PARRAY pShimRef)
 384 {
 385     TAGREF* Data;
 386
 387     Data = ARRAY_Append(pShimRef, TAGREF);
 388     if (!Data)
 389         return;
 390
 391     *Data = trShimRef;
 392 }
 393
 394 static VOID SeiAddFlag(PDB pdb, TAGID tiFlagRef, PFLAGINFO pFlagInfo)
 395 {
 396     ULARGE_INTEGER Flag;
 397
 398     /* Resolve the FLAG_REF to the real FLAG node */
 399     TAGID FlagTag = SeiGetDWORD(pdb, tiFlagRef, TAG_FLAG_TAGID);
 400
 401     if (FlagTag == TAGID_NULL)
 402         return;
 403
 404     pFlagInfo->AppCompatFlags.QuadPart |= SeiGetQWORD(pdb, FlagTag, TAG_FLAG_MASK_KERNEL);
 405     pFlagInfo->AppCompatFlagsUser.QuadPart |= SeiGetQWORD(pdb, FlagTag, TAG_FLAG_MASK_USER);
 406     Flag.QuadPart = SeiGetQWORD(pdb, FlagTag, TAG_FLAG_PROCESSPARAM);
 407     pFlagInfo->ProcessParameters_Flags |= Flag.LowPart;
 408 }
 409
 410 /* Propagate layers to child processes */
 411 static VOID SeiSetLayerEnvVar(LPCWSTR wszLayer)
 412 {
 413     NTSTATUS Status;
 414     UNICODE_STRING VarName = RTL_CONSTANT_STRING(L"__COMPAT_LAYER");
 415     UNICODE_STRING Value;
 416
 417     RtlInitUnicodeString(&Value, wszLayer);
 418
 419     Status = RtlSetEnvironmentVariable(NULL, &VarName, &Value);
 420     if (NT_SUCCESS(Status))
 421         SHIMENG_INFO("Set env var %wZ=%wZ\n", &VarName, &Value);
 422     else
 423         SHIMENG_FAIL("Failed to set %wZ: 0x%x\n", &VarName, Status);
 424 }
 425
 426 #define MAX_LAYER_LENGTH            256
 427
 428 /* Translate all Exe and Layer entries to Shims, and propagate all layers */
 429 static VOID SeiBuildShimRefArray(HSDB hsdb, SDBQUERYRESULT* pQuery, PARRAY pShimRef, PFLAGINFO pFlagInfo)
 430 {
 431     WCHAR wszLayerEnvVar[MAX_LAYER_LENGTH] = { 0 };
 432     DWORD n;
 433
 434     for (n = 0; n < pQuery->dwExeCount; ++n)
 435     {
 436         PDB pdb;
 437         TAGID tag;
 438         if (SdbTagRefToTagID(hsdb, pQuery->atrExes[n], &pdb, &tag))
 439         {
 440             LPCWSTR ExeName = SeiGetStringPtr(pdb, tag, TAG_NAME);
 441             TAGID ShimRef = SdbFindFirstTag(pdb, tag, TAG_SHIM_REF);
 442             TAGID FlagRef = SdbFindFirstTag(pdb, tag, TAG_FLAG_REF);
 443
 444             if (ExeName)
 445                 SeiDbgPrint(SEI_MSG, NULL, "ShimInfo(Exe(%S))\n", ExeName);
 446
 447             while (ShimRef != TAGID_NULL)
 448             {
 449                 TAGREF trShimRef;
 450                 if (SdbTagIDToTagRef(hsdb, pdb, ShimRef, &trShimRef))
 451                     SeiAddShim(trShimRef, pShimRef);
 452
 453                 ShimRef = SdbFindNextTag(pdb, tag, ShimRef);
 454             }
 455
 456             while (FlagRef != TAGID_NULL)
 457             {
 458                 SeiAddFlag(pdb, FlagRef, pFlagInfo);
 459
 460                 FlagRef = SdbFindNextTag(pdb, tag, FlagRef);
 461             }
 462         }
 463     }
 464
 465
 466     for (n = 0; n < pQuery->dwLayerCount; ++n)
 467     {
 468         PDB pdb;
 469         TAGID tag;
 470         if (SdbTagRefToTagID(hsdb, pQuery->atrLayers[n], &pdb, &tag))
 471         {
 472             LPCWSTR LayerName = SeiGetStringPtr(pdb, tag, TAG_NAME);
 473             TAGID ShimRef = SdbFindFirstTag(pdb, tag, TAG_SHIM_REF);
 474             TAGID FlagRef = SdbFindFirstTag(pdb, tag, TAG_FLAG_REF);
 475
 476             if (LayerName)
 477             {
 478                 HRESULT hr;
 479                 SeiDbgPrint(SEI_MSG, NULL, "ShimInfo(Layer(%S))\n", LayerName);
 480                 if (wszLayerEnvVar[0])
 481                     StringCchCatW(wszLayerEnvVar, ARRAYSIZE(wszLayerEnvVar), L" ");
 482                 hr = StringCchCatW(wszLayerEnvVar, ARRAYSIZE(wszLayerEnvVar), LayerName);
 483                 if (!SUCCEEDED(hr))
 484                 {
 485                     SHIMENG_FAIL("Unable to append %S\n", LayerName);
 486                 }
 487             }
 488
 489             while (ShimRef != TAGID_NULL)
 490             {
 491                 TAGREF trShimRef;
 492                 if (SdbTagIDToTagRef(hsdb, pdb, ShimRef, &trShimRef))
 493                     SeiAddShim(trShimRef, pShimRef);
 494
 495                 ShimRef = SdbFindNextTag(pdb, tag, ShimRef);
 496             }
 497
 498             while (FlagRef != TAGID_NULL)
 499             {
 500                 SeiAddFlag(pdb, FlagRef, pFlagInfo);
 501
 502                 FlagRef = SdbFindNextTag(pdb, tag, FlagRef);
 503             }
 504         }
 505     }
 506     if (wszLayerEnvVar[0])
 507         SeiSetLayerEnvVar(wszLayerEnvVar);
 508 }
 509
 510 /* Given the hooks from one shim, find the relevant modules and store the combination of module + hook */
 511 VOID SeiAddHooks(PHOOKAPIEX hooks, DWORD dwHookCount, PSHIMINFO pShim)
 512 {
 513     DWORD n, j;
 514     UNICODE_STRING UnicodeModName;
 515     WCHAR Buf[512];
 516
 517     RtlInitEmptyUnicodeString(&UnicodeModName, Buf, sizeof(Buf));
 518
 519     for (n = 0; n < dwHookCount; ++n)
 520     {
 521         ANSI_STRING AnsiString;
 522         PVOID DllHandle;
 523         PHOOKAPIEX hook = hooks + n;
 524         PHOOKAPIEX* pHookApi;
 525         PHOOKMODULEINFO HookModuleInfo;
 526
 527         RtlInitAnsiString(&AnsiString, hook->LibraryName);
 528         if (!NT_SUCCESS(RtlAnsiStringToUnicodeString(&UnicodeModName, &AnsiString, FALSE)))
 529         {
 530             SHIMENG_FAIL("Unable to convert %s to Unicode\n", hook->LibraryName);
 531             continue;
 532         }
 533
 534         RtlInitAnsiString(&AnsiString, hook->FunctionName);
 535         if (NT_SUCCESS(LdrGetDllHandle(NULL, 0, &UnicodeModName, &DllHandle)))
 536         {
 537             PVOID ProcAddress;
 538
 539
 540             if (!NT_SUCCESS(LdrGetProcedureAddress(DllHandle, &AnsiString, 0, &ProcAddress)))
 541             {
 542                 SHIMENG_FAIL("Unable to retrieve %s!%s\n", hook->LibraryName, hook->FunctionName);
 543                 continue;
 544             }
 545
 546             HookModuleInfo = SeiFindHookModuleInfo(NULL, DllHandle);
 547             hook->OriginalFunction = ProcAddress;
 548         }
 549         else
 550         {
 551             HookModuleInfo = SeiFindHookModuleInfo(&UnicodeModName, NULL);
 552             DllHandle = NULL;
 553         }
 554
 555         if (!HookModuleInfo)
 556         {
 557             HookModuleInfo = ARRAY_Append(&g_pHookArray, HOOKMODULEINFO);
 558             if (!HookModuleInfo)
 559                 continue;
 560
 561             HookModuleInfo->BaseAddress = DllHandle;
 562             ARRAY_Init(&HookModuleInfo->HookApis, PHOOKAPIEX);
 563             RtlCreateUnicodeString(&HookModuleInfo->Name, UnicodeModName.Buffer);
 564         }
 565
 566         hook->pShimInfo = pShim;
 567
 568         for (j = 0; j < ARRAY_Size(&HookModuleInfo->HookApis); ++j)
 569         {
 570             PHOOKAPIEX HookApi = *ARRAY_At(&HookModuleInfo->HookApis, PHOOKAPIEX, j);
 571             int CmpResult = strcmp(hook->FunctionName, HookApi->FunctionName);
 572             if (CmpResult == 0)
 573             {
 574                 /* Multiple hooks on one function? --> use ApiLink */
 575                 SHIMENG_FAIL("Multiple hooks on one API is not yet supported!\n");
 576                 ASSERT(0);
 577             }
 578         }
 579         pHookApi = ARRAY_Append(&HookModuleInfo->HookApis, PHOOKAPIEX);
 580         *pHookApi = hook;
 581     }
 582 }
 583
 584 typedef FARPROC(WINAPI* GETPROCADDRESSPROC)(HINSTANCE, LPCSTR);
 585
 586 /* Check if we should fake the return from GetProcAddress (because we also redirected the iat for this module) */
 587 FARPROC WINAPI StubGetProcAddress(HINSTANCE hModule, LPCSTR lpProcName)
 588 {
 589     char szOrdProcName[10] = "";
 590     LPCSTR lpPrintName = lpProcName;
 591     PVOID Addr = _ReturnAddress();
 592     PHOOKMODULEINFO HookModuleInfo;
 593     FARPROC proc = ((GETPROCADDRESSPROC)g_IntHookEx[0].OriginalFunction)(hModule, lpProcName);
 594
 595     if (!HIWORD(lpProcName))
 596     {
 597         sprintf(szOrdProcName, "#%lu", (DWORD)lpProcName);
 598         lpPrintName = szOrdProcName;
 599     }
 600
 601     Addr = SeiGetModuleFromAddress(Addr);
 602     if (SE_IsShimDll(Addr))
 603     {
 604         SHIMENG_MSG("Not touching GetProcAddress for shim dll (%p!%s)", hModule, lpPrintName);
 605         return proc;
 606     }
 607
 608     SHIMENG_INFO("(GetProcAddress(%p!%s) => %p\n", hModule, lpPrintName, proc);
 609
 610     HookModuleInfo = SeiFindHookModuleInfo(NULL, hModule);
 611
 612     /* FIXME: Ordinal not yet supported */
 613     if (HookModuleInfo && HIWORD(lpProcName))
 614     {
 615         DWORD n;
 616         for (n = 0; n < ARRAY_Size(&HookModuleInfo->HookApis); ++n)
 617         {
 618             PHOOKAPIEX HookApi = *ARRAY_At(&HookModuleInfo->HookApis, PHOOKAPIEX, n);
 619             int CmpResult = strcmp(lpProcName, HookApi->FunctionName);
 620             if (CmpResult == 0)
 621             {
 622                 SHIMENG_MSG("Redirecting %p to %p\n", proc, HookApi->ReplacementFunction);
 623                 proc = HookApi->ReplacementFunction;
 624                 break;
 625             }
 626         }
 627     }
 628
 629     return proc;
 630 }
 631
 632 /* Walk all shim modules / enabled shims, and add their hooks */
 633 VOID SeiResolveAPIs(VOID)
 634 {
 635     DWORD mod, n;
 636
 637     /* Enumerate all Shim modules */
 638     for (mod = 0; mod < ARRAY_Size(&g_pShimInfo); ++mod)
 639     {
 640         PSHIMMODULE pShimModule = *ARRAY_At(&g_pShimInfo, PSHIMMODULE, mod);
 641         DWORD dwShimCount = ARRAY_Size(&pShimModule->EnabledShims);
 642
 643         /* Enumerate all Shims */
 644         for (n = 0; n < dwShimCount; ++n)
 645         {
 646             PSHIMINFO pShim = *ARRAY_At(&pShimModule->EnabledShims, PSHIMINFO, n);
 647
 648             PHOOKAPIEX hooks = pShim->pHookApi;
 649             DWORD dwHookCount = pShim->dwHookCount;
 650
 651             SeiAddHooks(hooks, dwHookCount, pShim);
 652         }
 653     }
 654 }
 655
 656 /* If we hooked something, we should also redirect GetProcAddress */
 657 VOID SeiAddInternalHooks(DWORD dwNumHooks)
 658 {
 659     if (dwNumHooks == 0)
 660     {
 661         g_bInternalHooksUsed = FALSE;
 662         return;
 663     }
 664
 665     SeiAddHooks(g_IntHookEx, ARRAYSIZE(g_IntHookEx), NULL);
 666     g_bInternalHooksUsed = TRUE;
 667 }
 668
 669 /* Patch one function in the iat */
 670 VOID SeiPatchNewImport(PIMAGE_THUNK_DATA FirstThunk, PHOOKAPIEX HookApi, PLDR_DATA_TABLE_ENTRY LdrEntry)
 671 {
 672     ULONG OldProtection = 0;
 673     PVOID Ptr;
 674     ULONG Size;
 675     NTSTATUS Status;
 676
 677     SHIMENG_INFO("Hooking API \"%s!%s\" for DLL \"%wZ\"\n", HookApi->LibraryName, HookApi->FunctionName, &LdrEntry->BaseDllName);
 678
 679     Ptr = &FirstThunk->u1.Function;
 680     Size = sizeof(FirstThunk->u1.Function);
 681     Status = NtProtectVirtualMemory(NtCurrentProcess(), &Ptr, &Size, PAGE_EXECUTE_READWRITE, &OldProtection);
 682
 683     if (!NT_SUCCESS(Status))
 684     {
 685         SHIMENG_FAIL("Unable to unprotect 0x%p\n", &FirstThunk->u1.Function);
 686         return;
 687     }
 688
 689     SHIMENG_INFO("changing 0x%p to 0x%p\n", FirstThunk->u1.Function, HookApi->ReplacementFunction);
 690 #ifdef _WIN64
 691     FirstThunk->u1.Function = (ULONGLONG)HookApi->ReplacementFunction;
 692 #else
 693     FirstThunk->u1.Function = (DWORD)HookApi->ReplacementFunction;
 694 #endif
 695
 696     Size = sizeof(FirstThunk->u1.Function);
 697     Status = NtProtectVirtualMemory(NtCurrentProcess(), &Ptr, &Size, OldProtection, &OldProtection);
 698
 699     if (!NT_SUCCESS(Status))
 700     {
 701         SHIMENG_WARN("Unable to reprotect 0x%p\n", &FirstThunk->u1.Function);
 702     }
 703 }
 704
 705
 706 PINEXCLUDE SeiFindInExclude(PARRAY InExclude, PCUNICODE_STRING DllName)
 707 {
 708     DWORD n;
 709
 710     for (n = 0; n < ARRAY_Size(InExclude); ++n)
 711     {
 712         PINEXCLUDE InEx = ARRAY_At(InExclude, INEXCLUDE, n);
 713
 714         if (RtlEqualUnicodeString(&InEx->Module, DllName, TRUE))
 715             return InEx;
 716     }
 717
 718     return NULL;
 719 }
 720
 721 BOOL SeiIsExcluded(PLDR_DATA_TABLE_ENTRY LdrEntry, PHOOKAPIEX HookApi)
 722 {
 723     PSHIMINFO pShimInfo = HookApi->pShimInfo;
 724     PINEXCLUDE InExclude;
 725     BOOL IsExcluded = FALSE;
 726
 727     if (!pShimInfo)
 728     {
 729         /* Internal hook, do not exclude it */
 730         return FALSE;
 731     }
 732
 733     /* By default, everything from System32 or WinSxs is excluded */
 734     if (RtlPrefixUnicodeString(&g_System32Directory, &LdrEntry->FullDllName, TRUE) ||
 735         RtlPrefixUnicodeString(&g_SxsDirectory, &LdrEntry->FullDllName, TRUE))
 736         IsExcluded = TRUE;
 737
 738     InExclude = SeiFindInExclude(&pShimInfo->InExclude, &LdrEntry->BaseDllName);
 739     if (InExclude)
 740     {
 741         /* If it is on the 'exclude' list, bail out */
 742         if (!InExclude->Include)
 743         {
 744             SHIMENG_INFO("Module '%wZ' excluded for shim %S, API '%s!%s', because it on in the exclude list.\n",
 745                          &LdrEntry->BaseDllName, pShimInfo->ShimName, HookApi->LibraryName, HookApi->FunctionName);
 746
 747             return TRUE;
 748         }
 749         /* If it is on the 'include' list, override System32 / Winsxs check. */
 750         if (IsExcluded)
 751         {
 752             SHIMENG_INFO("Module '%wZ' included for shim %S, API '%s!%s', because it is on the include list.\n",
 753                          &LdrEntry->BaseDllName, pShimInfo->ShimName, HookApi->LibraryName, HookApi->FunctionName);
 754
 755         }
 756         IsExcluded = FALSE;
 757     }
 758
 759     if (IsExcluded)
 760     {
 761         SHIMENG_INFO("Module '%wZ' excluded for shim %S, API '%s!%s', because it is in System32/WinSXS.\n",
 762                      &LdrEntry->BaseDllName, pShimInfo->ShimName, HookApi->LibraryName, HookApi->FunctionName);
 763     }
 764
 765     return IsExcluded;
 766 }
 767
 768 VOID SeiAppendInExclude(PARRAY dest, PCWSTR ModuleName, BOOL IsInclude)
 769 {
 770     PINEXCLUDE InExclude;
 771     UNICODE_STRING ModuleNameU;
 772     RtlInitUnicodeString(&ModuleNameU, ModuleName);
 773
 774     InExclude = SeiFindInExclude(dest, &ModuleNameU);
 775     if (InExclude)
 776     {
 777         InExclude->Include = IsInclude;
 778         return;
 779     }
 780
 781     InExclude = ARRAY_Append(dest, INEXCLUDE);
 782     if (InExclude)
 783     {
 784         PCWSTR ModuleNameCopy = SdbpStrDup(ModuleName);
 785         RtlInitUnicodeString(&InExclude->Module, ModuleNameCopy);
 786         InExclude->Include = IsInclude;
 787     }
 788 }
 789
 790 /* Read the INEXCLUD tags from a given parent tag */
 791 VOID SeiReadInExclude(PDB pdb, TAGID parent, PARRAY dest)
 792 {
 793     TAGID InExcludeTag;
 794
 795     InExcludeTag = SdbFindFirstTag(pdb, parent, TAG_INEXCLUD);
 796
 797     while (InExcludeTag != TAGID_NULL)
 798     {
 799         PCWSTR ModuleName;
 800         TAGID ModuleTag = SdbFindFirstTag(pdb, InExcludeTag, TAG_MODULE);
 801         TAGID IncludeTag = SdbFindFirstTag(pdb, InExcludeTag, TAG_INCLUDE);
 802
 803         ModuleName = SdbGetStringTagPtr(pdb, ModuleTag);
 804         if (ModuleName)
 805         {
 806             SeiAppendInExclude(dest, ModuleName, IncludeTag != TAGID_NULL);
 807         }
 808         else
 809         {
 810             SHIMENG_WARN("INEXCLUDE without Module: 0x%x\n", InExcludeTag);
 811         }
 812
 813         InExcludeTag = SdbFindNextTag(pdb, parent, InExcludeTag);
 814     }
 815 }
 816
 817 VOID SeiBuildGlobalInclExclList(HSDB hsdb)
 818 {
 819     PDB pdb;
 820     TAGREF tr = TAGREF_ROOT;
 821     TAGID root, db, library;
 822
 823     if (!SdbTagRefToTagID(hsdb, tr, &pdb, &root))
 824     {
 825         SHIMENG_WARN("Unable to resolve database root\n");
 826         return;
 827     }
 828     db = SdbFindFirstTag(pdb, root, TAG_DATABASE);
 829     if (db == TAGID_NULL)
 830     {
 831         SHIMENG_WARN("Unable to resolve database\n");
 832         return;
 833     }
 834     library = SdbFindFirstTag(pdb, db, TAG_LIBRARY);
 835     if (library == TAGID_NULL)
 836     {
 837         SHIMENG_WARN("Unable to resolve library\n");
 838         return;
 839     }
 840
 841     SeiReadInExclude(pdb, library, &g_InExclude);
 842 }
 843
 844 VOID SeiBuildInclExclList(PDB pdb, TAGID ShimTag, PSHIMINFO pShimInfo)
 845 {
 846     DWORD n;
 847
 848     /* First duplicate the global in/excludes */
 849     for (n = 0; n < ARRAY_Size(&g_InExclude); ++n)
 850     {
 851         PINEXCLUDE InEx = ARRAY_At(&g_InExclude, INEXCLUDE, n);
 852         SeiAppendInExclude(&pShimInfo->InExclude, InEx->Module.Buffer, InEx->Include);
 853     }
 854
 855     /* Now read this shim's in/excludes (possibly overriding the global ones) */
 856     SeiReadInExclude(pdb, ShimTag, &pShimInfo->InExclude);
 857 }
 858
 859 /* Given one loaded module, redirect (hook) all functions from the iat that are registered by shims */
 860 VOID SeiHookImports(PLDR_DATA_TABLE_ENTRY LdrEntry)
 861 {
 862     ULONG Size;
 863     PIMAGE_IMPORT_DESCRIPTOR ImportDescriptor;
 864     PBYTE DllBase = LdrEntry->DllBase;
 865
 866     if (SE_IsShimDll(DllBase) || g_hInstance == LdrEntry->DllBase)
 867     {
 868         SHIMENG_INFO("Skipping shim module 0x%p \"%wZ\"\n", LdrEntry->DllBase, &LdrEntry->BaseDllName);
 869         return;
 870     }
 871
 872     if (LdrEntry->Flags & LDRP_COMPAT_DATABASE_PROCESSED)
 873     {
 874         SHIMENG_INFO("Skipping module 0x%p \"%wZ\" because it was already processed\n", LdrEntry->DllBase, &LdrEntry->BaseDllName);
 875         return;
 876     }
 877
 878     ImportDescriptor = RtlImageDirectoryEntryToData(DllBase, TRUE, IMAGE_DIRECTORY_ENTRY_IMPORT, &Size);
 879     if (!ImportDescriptor)
 880     {
 881         SHIMENG_INFO("Skipping module 0x%p \"%wZ\" due to no iat found\n", LdrEntry->DllBase, &LdrEntry->BaseDllName);
 882         return;
 883     }
 884
 885     SHIMENG_INFO("Hooking module 0x%p \"%wZ\"\n", LdrEntry->DllBase, &LdrEntry->BaseDllName);
 886
 887     for ( ;ImportDescriptor->Name && ImportDescriptor->OriginalFirstThunk; ImportDescriptor++)
 888     {
 889         PHOOKMODULEINFO HookModuleInfo;
 890
 891         /* Do we have hooks for this module? */
 892         HookModuleInfo = SeiFindHookModuleInfoForImportDescriptor(DllBase, ImportDescriptor);
 893
 894         if (HookModuleInfo)
 895         {
 896             PIMAGE_THUNK_DATA OriginalThunk, FirstThunk;
 897             DWORD n;
 898
 899             for (n = 0; n < ARRAY_Size(&HookModuleInfo->HookApis); ++n)
 900             {
 901                 DWORD dwFound = 0;
 902                 PHOOKAPIEX HookApi = *ARRAY_At(&HookModuleInfo->HookApis, PHOOKAPIEX, n);
 903
 904                 /* Check if this module should be excluded from being hooked (system32/winsxs, global or shim exclude) */
 905                 if (SeiIsExcluded(LdrEntry, HookApi))
 906                 {
 907                     continue;
 908                 }
 909
 910                 OriginalThunk = (PIMAGE_THUNK_DATA)(DllBase + ImportDescriptor->OriginalFirstThunk);
 911                 FirstThunk = (PIMAGE_THUNK_DATA)(DllBase + ImportDescriptor->FirstThunk);
 912
 913                 /* Walk all imports */
 914                 for (;OriginalThunk->u1.AddressOfData && FirstThunk->u1.Function; OriginalThunk++, FirstThunk++)
 915                 {
 916                     if (!IMAGE_SNAP_BY_ORDINAL32(OriginalThunk->u1.AddressOfData))
 917                     {
 918                         PIMAGE_IMPORT_BY_NAME ImportName;
 919
 920                         ImportName = (PIMAGE_IMPORT_BY_NAME)(DllBase + OriginalThunk->u1.AddressOfData);
 921                         if (!strcmp((PCSTR)ImportName->Name, HookApi->FunctionName))
 922                         {
 923                             SeiPatchNewImport(FirstThunk, HookApi, LdrEntry);
 924
 925                             /* Sadly, iat does not have to be sorted, and can even contain duplicate entries. */
 926                             dwFound++;
 927                         }
 928                     }
 929                     else
 930                     {
 931                         SHIMENG_FAIL("Ordinals not yet supported\n");
 932                         ASSERT(0);
 933                     }
 934                 }
 935
 936                 if (dwFound != 1)
 937                 {
 938                     /* One entry not found. */
 939                     if (!dwFound)
 940                         SHIMENG_INFO("Entry \"%s!%s\" not found for \"%wZ\"\n", HookApi->LibraryName, HookApi->FunctionName, &LdrEntry->BaseDllName);
 941                     else
 942                         SHIMENG_INFO("Entry \"%s!%s\" found %d times for \"%wZ\"\n", HookApi->LibraryName, HookApi->FunctionName, dwFound, &LdrEntry->BaseDllName);
 943                 }
 944             }
 945         }
 946     }
 947
 948     /* Mark this module as processed. */
 949     LdrEntry->Flags |= LDRP_COMPAT_DATABASE_PROCESSED;
 950 }
 951
 952
 953 VOID PatchNewModules(PPEB Peb)
 954 {
 955     PLIST_ENTRY ListHead, ListEntry;
 956     PLDR_DATA_TABLE_ENTRY LdrEntry;
 957
 958     ListHead = &NtCurrentPeb()->Ldr->InLoadOrderModuleList;
 959     ListEntry = ListHead->Flink;
 960
 961     while (ListHead != ListEntry)
 962     {
 963         LdrEntry = CONTAINING_RECORD(ListEntry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
 964         SeiHookImports(LdrEntry);
 965
 966         ListEntry = ListEntry->Flink;
 967     }
 968 }
 969
 970
 971 VOID SeiInitPaths(VOID)
 972 {
 973 #define SYSTEM32  L"\\system32"
 974 #define WINSXS  L"\\winsxs"
 975
 976     PWSTR WindowsDirectory = SdbpStrDup(SharedUserData->NtSystemRoot);
 977     RtlInitUnicodeString(&g_WindowsDirectory, WindowsDirectory);
 978
 979     g_System32Directory.MaximumLength = g_WindowsDirectory.Length + SdbpStrsize(SYSTEM32);
 980     g_System32Directory.Buffer = SdbpAlloc(g_System32Directory.MaximumLength);
 981     RtlCopyUnicodeString(&g_System32Directory, &g_WindowsDirectory);
 982     RtlAppendUnicodeToString(&g_System32Directory, SYSTEM32);
 983
 984     g_SxsDirectory.MaximumLength = g_WindowsDirectory.Length + SdbpStrsize(WINSXS);
 985     g_SxsDirectory.Buffer = SdbpAlloc(g_SxsDirectory.MaximumLength);
 986     RtlCopyUnicodeString(&g_SxsDirectory, &g_WindowsDirectory);
 987     RtlAppendUnicodeToString(&g_SxsDirectory, WINSXS);
 988
 989 #undef SYSTEM32
 990 #undef WINSXS
 991 }
 992
 993 VOID SeiInit(PUNICODE_STRING ProcessImage, HSDB hsdb, SDBQUERYRESULT* pQuery)
 994 {
 995     DWORD n;
 996     ARRAY ShimRefArray;
 997     DWORD dwTotalHooks = 0;
 998     FLAGINFO ShimFlags;
 999
1000     PPEB Peb = NtCurrentPeb();
1001
1002     /* We should only be called once! */
1003     ASSERT(g_pShimInfo.ItemSize__ == 0);
1004
1005     ARRAY_Init(&ShimRefArray, TAGREF);
1006     ARRAY_Init(&g_pShimInfo, PSHIMMODULE);
1007     ARRAY_Init(&g_pHookArray, HOOKMODULEINFO);
1008     ARRAY_Init(&g_InExclude, INEXCLUDE);
1009     RtlZeroMemory(&ShimFlags, sizeof(ShimFlags));
1010
1011     SeiInitPaths();
1012
1013     SeiCheckComPlusImage(Peb->ImageBaseAddress);
1014
1015     /* TODO:
1016     if (pQuery->trApphelp)
1017         SeiDisplayAppHelp(?pQuery->trApphelp?);
1018     */
1019
1020     SeiDbgPrint(SEI_MSG, NULL, "ShimInfo(ExePath(%wZ))\n", ProcessImage);
1021     SeiBuildShimRefArray(hsdb, pQuery, &ShimRefArray, &ShimFlags);
1022     if (ShimFlags.AppCompatFlags.QuadPart)
1023     {
1024         SeiDbgPrint(SEI_MSG, NULL, "Using KERNEL apphack flags 0x%I64x\n", ShimFlags.AppCompatFlags.QuadPart);
1025         Peb->AppCompatFlags.QuadPart |= ShimFlags.AppCompatFlags.QuadPart;
1026     }
1027     if (ShimFlags.AppCompatFlagsUser.QuadPart)
1028     {
1029         SeiDbgPrint(SEI_MSG, NULL, "Using USER apphack flags 0x%I64x\n", ShimFlags.AppCompatFlagsUser.QuadPart);
1030         Peb->AppCompatFlagsUser.QuadPart |= ShimFlags.AppCompatFlagsUser.QuadPart;
1031     }
1032     if (ShimFlags.ProcessParameters_Flags)
1033     {
1034         SeiDbgPrint(SEI_MSG, NULL, "Using ProcessParameters flags 0x%x\n", ShimFlags.ProcessParameters_Flags);
1035         Peb->ProcessParameters->Flags |= ShimFlags.ProcessParameters_Flags;
1036     }
1037     SeiDbgPrint(SEI_MSG, NULL, "ShimInfo(Complete)\n");
1038
1039     SHIMENG_INFO("Got %d shims\n", ARRAY_Size(&ShimRefArray));
1040     SeiBuildGlobalInclExclList(hsdb);
1041
1042     /* Walk all shims referenced (in layers + exes), and load their modules */
1043     for (n = 0; n < ARRAY_Size(&ShimRefArray); ++n)
1044     {
1045         PDB pdb;
1046         TAGID ShimRef;
1047
1048         TAGREF tr = *ARRAY_At(&ShimRefArray, TAGREF, n);
1049
1050         if (SdbTagRefToTagID(hsdb, tr, &pdb, &ShimRef))
1051         {
1052             LPCWSTR ShimName, DllName, CommandLine = NULL;
1053             TAGID ShimTag;
1054             WCHAR FullNameBuffer[MAX_PATH];
1055             UNICODE_STRING UnicodeDllName;
1056             PVOID BaseAddress;
1057             PSHIMMODULE pShimModuleInfo = NULL;
1058             ANSI_STRING AnsiCommandLine = RTL_CONSTANT_STRING("");
1059             PSHIMINFO pShimInfo = NULL;
1060             PHOOKAPIEX pHookApi;
1061             DWORD dwHookCount;
1062
1063             ShimName = SeiGetStringPtr(pdb, ShimRef, TAG_NAME);
1064             if (!ShimName)
1065             {
1066                 SHIMENG_FAIL("Failed to retrieve the name for 0x%x\n", tr);
1067                 continue;
1068             }
1069
1070             CommandLine = SeiGetStringPtr(pdb, ShimRef, TAG_COMMAND_LINE);
1071             if (CommandLine && *CommandLine)
1072             {
1073                 RtlInitUnicodeString(&UnicodeDllName, CommandLine);
1074                 if (NT_SUCCESS(RtlUnicodeStringToAnsiString(&AnsiCommandLine, &UnicodeDllName, TRUE)))
1075                 {
1076                     SHIMENG_INFO("COMMAND LINE %s for %S", AnsiCommandLine.Buffer, ShimName);
1077                 }
1078                 else
1079                 {
1080                     AnsiCommandLine.Buffer = "";
1081                     CommandLine = NULL;
1082                 }
1083             }
1084
1085             ShimTag = SeiGetDWORD(pdb, ShimRef, TAG_SHIM_TAGID);
1086             if (!ShimTag)
1087             {
1088                 SHIMENG_FAIL("Failed to resolve %S to a shim\n", ShimName);
1089                 continue;
1090             }
1091
1092             if (!SUCCEEDED(SdbGetAppPatchDir(NULL, FullNameBuffer, ARRAYSIZE(FullNameBuffer))))
1093             {
1094                 SHIMENG_WARN("Failed to get the AppPatch dir\n");
1095                 continue;
1096             }
1097
1098             DllName = SeiGetStringPtr(pdb, ShimTag, TAG_DLLFILE);
1099             if (DllName == NULL ||
1100                 !SUCCEEDED(StringCchCatW(FullNameBuffer, ARRAYSIZE(FullNameBuffer), L"\\")) ||
1101                 !SUCCEEDED(StringCchCatW(FullNameBuffer, ARRAYSIZE(FullNameBuffer), DllName)))
1102             {
1103                 SHIMENG_WARN("Failed to build a full path for %S\n", ShimName);
1104                 continue;
1105             }
1106
1107             RtlInitUnicodeString(&UnicodeDllName, FullNameBuffer);
1108             if (NT_SUCCESS(LdrGetDllHandle(NULL, NULL, &UnicodeDllName, &BaseAddress)))
1109             {
1110                 /* This shim dll was already loaded, let's find it */
1111                 pShimModuleInfo = SeiGetShimModuleInfo(BaseAddress);
1112             }
1113             else if (!NT_SUCCESS(LdrLoadDll(NULL, NULL, &UnicodeDllName, &BaseAddress)))
1114             {
1115                 SHIMENG_WARN("Failed to load %wZ for %S\n", &UnicodeDllName, ShimName);
1116                 continue;
1117             }
1118             /* No shim module found (or we just loaded it) */
1119             if (!pShimModuleInfo)
1120             {
1121                 pShimModuleInfo = SeiCreateShimModuleInfo(DllName, BaseAddress);
1122                 if (!pShimModuleInfo)
1123                 {
1124                     SHIMENG_FAIL("Failed to allocate ShimInfo for %S\n", DllName);
1125                     continue;
1126                 }
1127             }
1128
1129             SHIMENG_INFO("Shim DLL 0x%p \"%wZ\" loaded\n", BaseAddress, &UnicodeDllName);
1130             SHIMENG_INFO("Using SHIM \"%S!%S\"\n", DllName, ShimName);
1131
1132             /* Ask this shim what hooks it needs (and pass along the commandline) */
1133             pHookApi = pShimModuleInfo->pGetHookAPIs(AnsiCommandLine.Buffer, ShimName, &dwHookCount);
1134             SHIMENG_INFO("GetHookAPIs returns %d hooks for DLL \"%wZ\" SHIM \"%S\"\n", dwHookCount, &UnicodeDllName, ShimName);
1135             if (dwHookCount)
1136                 pShimInfo = SeiAppendHookInfo(pShimModuleInfo, pHookApi, dwHookCount, ShimName);
1137
1138             /* If this shim has hooks, create the include / exclude lists */
1139             if (pShimInfo)
1140                 SeiBuildInclExclList(pdb, ShimTag, pShimInfo);
1141
1142             if (CommandLine && *CommandLine)
1143                 RtlFreeAnsiString(&AnsiCommandLine);
1144
1145             dwTotalHooks += dwHookCount;
1146         }
1147     }
1148
1149     SeiAddInternalHooks(dwTotalHooks);
1150     SeiResolveAPIs();
1151     PatchNewModules(Peb);
1152 }
1153
1154
1155 /* Load the database + unpack the shim data (if this process is allowed) */
1156 BOOL SeiGetShimData(PUNICODE_STRING ProcessImage, PVOID pShimData, HSDB* pHsdb, SDBQUERYRESULT* pQuery)
1157 {
1158     static const UNICODE_STRING ForbiddenShimmingApps[] = {
1159         RTL_CONSTANT_STRING(L"ntsd.exe"),
1160         RTL_CONSTANT_STRING(L"windbg.exe"),
1161 #if WINVER >= 0x600
1162         RTL_CONSTANT_STRING(L"slsvc.exe"),
1163 #endif
1164     };
1165     static const UNICODE_STRING BackSlash = RTL_CONSTANT_STRING(L"\\");
1166     static const UNICODE_STRING ForwdSlash = RTL_CONSTANT_STRING(L"/");
1167     UNICODE_STRING ProcessName;
1168     USHORT Back, Forward;
1169     HSDB hsdb;
1170     DWORD n;
1171
1172     if (!NT_SUCCESS(RtlFindCharInUnicodeString(RTL_FIND_CHAR_IN_UNICODE_STRING_START_AT_END, ProcessImage, &BackSlash, &Back)))
1173         Back = 0;
1174
1175     if (!NT_SUCCESS(RtlFindCharInUnicodeString(RTL_FIND_CHAR_IN_UNICODE_STRING_START_AT_END, ProcessImage, &ForwdSlash, &Forward)))
1176         Forward = 0;
1177
1178     if (Back < Forward)
1179         Back = Forward;
1180
1181     if (Back)
1182         Back += sizeof(WCHAR);
1183
1184     ProcessName.Buffer = ProcessImage->Buffer + Back / sizeof(WCHAR);
1185     ProcessName.Length = ProcessImage->Length - Back;
1186     ProcessName.MaximumLength = ProcessImage->MaximumLength - Back;
1187
1188     for (n = 0; n < ARRAYSIZE(ForbiddenShimmingApps); ++n)
1189     {
1190         if (RtlEqualUnicodeString(&ProcessName, ForbiddenShimmingApps + n, TRUE))
1191         {
1192             SHIMENG_MSG("Not shimming %wZ\n", ForbiddenShimmingApps + n);
1193             return FALSE;
1194         }
1195     }
1196
1197     /* We should probably load all db's here, but since we do not support that yet... */
1198     hsdb = SdbInitDatabase(HID_DOS_PATHS | SDB_DATABASE_MAIN_SHIM, NULL);
1199     if (hsdb)
1200     {
1201         if (SdbUnpackAppCompatData(hsdb, ProcessImage->Buffer, pShimData, pQuery))
1202         {
1203             *pHsdb = hsdb;
1204             return TRUE;
1205         }
1206         SdbReleaseDatabase(hsdb);
1207     }
1208     return FALSE;
1209 }
1210
1211
1212
1213 VOID NTAPI SE_InstallBeforeInit(PUNICODE_STRING ProcessImage, PVOID pShimData)
1214 {
1215     HSDB hsdb = NULL;
1216     SDBQUERYRESULT QueryResult = { { 0 } };
1217     SHIMENG_INFO("(%wZ, %p)\n", ProcessImage, pShimData);
1218
1219     if (!SeiGetShimData(ProcessImage, pShimData, &hsdb, &QueryResult))
1220     {
1221         SHIMENG_FAIL("Failed to get shim data\n");
1222         return;
1223     }
1224
1225     g_bShimDuringInit = TRUE;
1226     SeiInit(ProcessImage, hsdb, &QueryResult);
1227     g_bShimDuringInit = FALSE;
1228
1229     SdbReleaseDatabase(hsdb);
1230 }
1231
1232 VOID NTAPI SE_InstallAfterInit(PUNICODE_STRING ProcessImage, PVOID pShimData)
1233 {
1234     NotifyShims(SHIM_NOTIFY_ATTACH, NULL);
1235 }
1236
1237 VOID NTAPI SE_ProcessDying(VOID)
1238 {
1239     SHIMENG_MSG("()\n");
1240     NotifyShims(SHIM_NOTIFY_DETACH, NULL);
1241 }
1242
1243 VOID WINAPI SE_DllLoaded(PLDR_DATA_TABLE_ENTRY LdrEntry)
1244 {
1245     SHIMENG_INFO("%sINIT. loading DLL \"%wZ\"\n", g_bShimDuringInit ? "" : "AFTER ", &LdrEntry->BaseDllName);
1246
1247     SeiHookImports(LdrEntry);
1248
1249     NotifyShims(SHIM_REASON_DLL_LOAD, LdrEntry);
1250 }
1251
1252 VOID WINAPI SE_DllUnloaded(PLDR_DATA_TABLE_ENTRY LdrEntry)
1253 {
1254     SHIMENG_INFO("(%p)\n", LdrEntry);
1255
1256     /* Should we unhook here? */
1257
1258     NotifyShims(SHIM_REASON_DLL_UNLOAD, LdrEntry);
1259 }
1260
1261 BOOL WINAPI SE_IsShimDll(PVOID BaseAddress)
1262 {
1263     SHIMENG_INFO("(%p)\n", BaseAddress);
1264
1265     return SeiGetShimModuleInfo(BaseAddress) != NULL;
1266 }
1267