dll/appcompat/apphelp/shimeng.c

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