reactos/ntoskrnl/ps/apphelp.c

   1 /*
   2  * PROJECT:         ReactOS Kernel
   3  * LICENSE:         BSD - See COPYING.ARM in the top level directory
   4  * FILE:            ntoskrnl/ps/apphelp.c
   5  * PURPOSE:         SHIM engine caching.
   6  *                  This caching speeds up checks for the apphelp compatibility layer.
   7  * PROGRAMMERS:     Mark Jansen
   8  */
   9
  10 /*
  11 Useful references:
  12 https://github.com/mandiant/ShimCacheParser/blob/master/ShimCacheParser.py
  13 http://technet.microsoft.com/en-us/library/dd837644(v=ws.10).aspx
  14 http://msdn.microsoft.com/en-us/library/bb432182(v=vs.85).aspx
  15 http://www.alex-ionescu.com/?p=43
  16 http://recxltd.blogspot.nl/2012/04/windows-appcompat-research-notes-part-1.html
  17 http://journeyintoir.blogspot.ch/2013/12/revealing-recentfilecachebcf-file.html
  18 https://dl.mandiant.com/EE/library/Whitepaper_ShimCacheParser.pdf
  19 */
  20
  21 /* INCLUDES ******************************************************************/
  22
  23 #include <ntoskrnl.h>
  24 #define NDEBUG
  25 #include <debug.h>
  26
  27 /* GLOBALS *******************************************************************/
  28
  29 static BOOLEAN ApphelpCacheEnabled = FALSE;
  30 static ERESOURCE ApphelpCacheLock;
  31 static RTL_AVL_TABLE ApphelpShimCache;
  32 static LIST_ENTRY ApphelpShimCacheAge;
  33
  34 extern ULONG InitSafeBootMode;
  35
  36 static UNICODE_STRING AppCompatCacheKey = RTL_CONSTANT_STRING(L"\\Registry\\MACHINE\\System\\CurrentControlSet\\Control\\Session Manager\\AppCompatCache");
  37 static OBJECT_ATTRIBUTES AppCompatKeyAttributes = RTL_CONSTANT_OBJECT_ATTRIBUTES(&AppCompatCacheKey, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE);
  38 static UNICODE_STRING AppCompatCacheValue = RTL_CONSTANT_STRING(L"AppCompatCache");
  39
  40 #define EMPTY_SHIM_ENTRY    { { 0 }, { { 0 } }, 0 }
  41 #define MAX_SHIM_ENTRIES    0x200
  42 #define TAG_SHIM            'MIHS'
  43
  44 #ifndef INVALID_HANDLE_VALUE
  45 #define INVALID_HANDLE_VALUE (HANDLE)(-1)
  46 #endif
  47
  48 #include <pshpack1.h>
  49
  50 typedef struct SHIM_PERSISTENT_CACHE_HEADER_52
  51 {
  52     ULONG Magic;
  53     ULONG NumEntries;
  54 } SHIM_PERSISTENT_CACHE_HEADER_52, *PSHIM_PERSISTENT_CACHE_HEADER_52;
  55
  56 /* The data that is present in the registry (Win2k3 version) */
  57 typedef struct SHIM_PERSISTENT_CACHE_ENTRY_52
  58 {
  59     UNICODE_STRING ImageName;
  60     LARGE_INTEGER DateTime;
  61     LARGE_INTEGER FileSize;
  62 } SHIM_PERSISTENT_CACHE_ENTRY_52, *PSHIM_PERSISTENT_CACHE_ENTRY_52;
  63
  64 #include <poppack.h>
  65
  66 #define CACHE_MAGIC_NT_52 0xbadc0ffe
  67 #define CACHE_HEADER_SIZE_NT_52 0x8
  68 #define NT52_PERSISTENT_ENTRY_SIZE32 0x18
  69 #define NT52_PERSISTENT_ENTRY_SIZE64 0x20
  70
  71 //#define CACHE_MAGIC_NT_61 0xbadc0fee
  72 //#define CACHE_HEADER_SIZE_NT_61 0x80
  73 //#define NT61_PERSISTENT_ENTRY_SIZE32 0x20
  74 //#define NT61_PERSISTENT_ENTRY_SIZE64 0x30
  75
  76 #define SHIM_CACHE_MAGIC                    CACHE_MAGIC_NT_52
  77 #define SHIM_CACHE_HEADER_SIZE              CACHE_HEADER_SIZE_NT_52
  78 #ifdef _WIN64
  79 #define SHIM_PERSISTENT_CACHE_ENTRY_SIZE    NT52_PERSISTENT_ENTRY_SIZE64
  80 #else
  81 #define SHIM_PERSISTENT_CACHE_ENTRY_SIZE    NT52_PERSISTENT_ENTRY_SIZE32
  82 #endif
  83 #define SHIM_PERSISTENT_CACHE_HEADER        SHIM_PERSISTENT_CACHE_HEADER_52
  84 #define PSHIM_PERSISTENT_CACHE_HEADER      PSHIM_PERSISTENT_CACHE_HEADER_52
  85 #define SHIM_PERSISTENT_CACHE_ENTRY         SHIM_PERSISTENT_CACHE_ENTRY_52
  86 #define PSHIM_PERSISTENT_CACHE_ENTRY       PSHIM_PERSISTENT_CACHE_ENTRY_52
  87
  88 C_ASSERT(sizeof(SHIM_PERSISTENT_CACHE_ENTRY) == SHIM_PERSISTENT_CACHE_ENTRY_SIZE);
  89 C_ASSERT(sizeof(SHIM_PERSISTENT_CACHE_HEADER) == SHIM_CACHE_HEADER_SIZE);
  90
  91 /* The struct we keep in memory */
  92 typedef struct SHIM_CACHE_ENTRY
  93 {
  94     LIST_ENTRY List;
  95     SHIM_PERSISTENT_CACHE_ENTRY Persistent;
  96     ULONG CompatFlags;
  97 } SHIM_CACHE_ENTRY, *PSHIM_CACHE_ENTRY;
  98
  99 /* PRIVATE FUNCTIONS *********************************************************/
 100
 101 PVOID
 102 ApphelpAlloc(
 103     _In_ ULONG ByteSize)
 104 {
 105     return ExAllocatePoolWithTag(PagedPool, ByteSize, TAG_SHIM);
 106 }
 107
 108 VOID
 109 ApphelpFree(
 110     _In_ PVOID Data)
 111 {
 112     ExFreePoolWithTag(Data, TAG_SHIM);
 113 }
 114
 115 VOID
 116 ApphelpCacheAcquireLock(VOID)
 117 {
 118     KeEnterCriticalRegion();
 119     ExAcquireResourceExclusiveLite(&ApphelpCacheLock, TRUE);
 120 }
 121
 122 BOOLEAN
 123 ApphelpCacheTryAcquireLock(VOID)
 124 {
 125     KeEnterCriticalRegion();
 126     if (!ExTryToAcquireResourceExclusiveLite(&ApphelpCacheLock))
 127     {
 128         KeLeaveCriticalRegion();
 129         return FALSE;
 130     }
 131     return TRUE;
 132 }
 133
 134 VOID
 135 ApphelpCacheReleaseLock(VOID)
 136 {
 137     ExReleaseResourceLite(&ApphelpCacheLock);
 138     KeLeaveCriticalRegion();
 139 }
 140
 141 VOID
 142 ApphelpDuplicateUnicodeString(
 143     _Out_ PUNICODE_STRING Destination,
 144     _In_ PCUNICODE_STRING Source)
 145 {
 146     Destination->Length = Source->Length;
 147     if (Destination->Length)
 148     {
 149         Destination->MaximumLength = Destination->Length + sizeof(WCHAR);
 150         Destination->Buffer = ApphelpAlloc(Destination->MaximumLength);
 151         RtlCopyMemory(Destination->Buffer, Source->Buffer, Destination->Length);
 152         Destination->Buffer[Destination->Length / sizeof(WCHAR)] = UNICODE_NULL;
 153     }
 154     else
 155     {
 156         Destination->MaximumLength = 0;
 157         Destination->Buffer = NULL;
 158     }
 159 }
 160
 161 VOID
 162 ApphelpFreeUnicodeString(
 163     _Inout_ PUNICODE_STRING String)
 164 {
 165     if (String->Buffer)
 166     {
 167         ApphelpFree(String->Buffer);
 168     }
 169     String->Length = 0;
 170     String->MaximumLength = 0;
 171     String->Buffer = NULL;
 172 }
 173
 174 /* Query file info from a handle, storing it in Entry */
 175 NTSTATUS
 176 ApphelpCacheQueryInfo(
 177     _In_ HANDLE ImageHandle,
 178     _Out_ PSHIM_CACHE_ENTRY Entry)
 179 {
 180     IO_STATUS_BLOCK IoStatusBlock;
 181     FILE_BASIC_INFORMATION FileBasic;
 182     FILE_STANDARD_INFORMATION FileStandard;
 183     NTSTATUS Status;
 184
 185     Status = ZwQueryInformationFile(ImageHandle,
 186                                     &IoStatusBlock,
 187                                     &FileBasic,
 188                                     sizeof(FileBasic),
 189                                     FileBasicInformation);
 190     if (!NT_SUCCESS(Status))
 191     {
 192         return Status;
 193     }
 194
 195     Status = ZwQueryInformationFile(ImageHandle,
 196                                     &IoStatusBlock,
 197                                     &FileStandard,
 198                                     sizeof(FileStandard),
 199                                     FileStandardInformation);
 200     if (NT_SUCCESS(Status))
 201     {
 202         Entry->Persistent.DateTime = FileBasic.LastWriteTime;
 203         Entry->Persistent.FileSize = FileStandard.EndOfFile;
 204     }
 205     return Status;
 206 }
 207
 208 RTL_GENERIC_COMPARE_RESULTS
 209 NTAPI
 210 ApphelpShimCacheCompareRoutine(
 211     _In_ PRTL_AVL_TABLE Table,
 212     _In_ PVOID FirstStruct,
 213     _In_ PVOID SecondStruct)
 214 {
 215     PSHIM_CACHE_ENTRY FirstEntry = FirstStruct;
 216     PSHIM_CACHE_ENTRY SecondEntry = SecondStruct;
 217     LONG Result;
 218
 219     Result = RtlCompareUnicodeString(&FirstEntry->Persistent.ImageName,
 220                                      &SecondEntry->Persistent.ImageName,
 221                                      TRUE);
 222     if (Result < 0)
 223     {
 224         return GenericLessThan;
 225     }
 226     else if (Result == 0)
 227     {
 228         return GenericEqual;
 229     }
 230     return GenericGreaterThan;
 231 }
 232
 233 PVOID
 234 NTAPI
 235 ApphelpShimCacheAllocateRoutine(
 236     _In_ PRTL_AVL_TABLE Table,
 237     _In_ CLONG ByteSize)
 238 {
 239     return ApphelpAlloc(ByteSize);
 240 }
 241
 242 VOID
 243 NTAPI
 244 ApphelpShimCacheFreeRoutine(
 245     _In_ PRTL_AVL_TABLE Table,
 246     _In_ PVOID Buffer)
 247 {
 248     ApphelpFree(Buffer);
 249 }
 250
 251 NTSTATUS
 252 ApphelpCacheParse(
 253     _In_reads_(DataLength) PUCHAR Data,
 254     _In_ ULONG DataLength)
 255 {
 256     PSHIM_PERSISTENT_CACHE_HEADER Header = (PSHIM_PERSISTENT_CACHE_HEADER)Data;
 257     ULONG Cur;
 258     ULONG NumEntries;
 259     UNICODE_STRING String;
 260     SHIM_CACHE_ENTRY Entry = EMPTY_SHIM_ENTRY;
 261     PSHIM_CACHE_ENTRY Result;
 262     PSHIM_PERSISTENT_CACHE_ENTRY Persistent;
 263
 264     if (DataLength < CACHE_HEADER_SIZE_NT_52)
 265     {
 266         DPRINT1("SHIMS: ApphelpCacheParse not enough data for a minimal header (0x%x)\n", DataLength);
 267         return STATUS_INVALID_PARAMETER;
 268     }
 269
 270     if (Header->Magic != SHIM_CACHE_MAGIC)
 271     {
 272         DPRINT1("SHIMS: ApphelpCacheParse found invalid magic (0x%x)\n", Header->Magic);
 273         return STATUS_INVALID_PARAMETER;
 274     }
 275
 276     NumEntries = Header->NumEntries;
 277     DPRINT("SHIMS: ApphelpCacheParse walking %d entries\n", NumEntries);
 278     for (Cur = 0; Cur < NumEntries; ++Cur)
 279     {
 280         Persistent = (PSHIM_PERSISTENT_CACHE_ENTRY)(Data + SHIM_CACHE_HEADER_SIZE +
 281                                                     (Cur * SHIM_PERSISTENT_CACHE_ENTRY_SIZE));
 282         /* The entry in the Persistent storage is not really a UNICODE_STRING,
 283             so we have to convert the offset into a real pointer before using it. */
 284         String.Length = Persistent->ImageName.Length;
 285         String.MaximumLength = Persistent->ImageName.MaximumLength;
 286         String.Buffer = (PWCHAR)((ULONG_PTR)Persistent->ImageName.Buffer + Data);
 287
 288         /* Now we copy all data to a local buffer, that can be safely duplicated by RtlInsert */
 289         Entry.Persistent = *Persistent;
 290         ApphelpDuplicateUnicodeString(&Entry.Persistent.ImageName, &String);
 291         Result = RtlInsertElementGenericTableAvl(&ApphelpShimCache,
 292                                                  &Entry,
 293                                                  sizeof(Entry),
 294                                                  NULL);
 295         if (!Result)
 296         {
 297             DPRINT1("SHIMS: ApphelpCacheParse insert failed\n");
 298             ApphelpFreeUnicodeString(&Entry.Persistent.ImageName);
 299             return STATUS_INVALID_PARAMETER;
 300         }
 301         InsertTailList(&ApphelpShimCacheAge, &Result->List);
 302     }
 303     return STATUS_SUCCESS;
 304 }
 305
 306 BOOLEAN
 307 ApphelpCacheRead(VOID)
 308 {
 309     HANDLE KeyHandle;
 310     NTSTATUS Status;
 311     KEY_VALUE_PARTIAL_INFORMATION KeyValueObject;
 312     PKEY_VALUE_PARTIAL_INFORMATION KeyValueInformation = &KeyValueObject;
 313     ULONG KeyInfoSize, ResultSize;
 314
 315     Status = ZwOpenKey(&KeyHandle, KEY_QUERY_VALUE, &AppCompatKeyAttributes);
 316     if (!NT_SUCCESS(Status))
 317     {
 318         DPRINT1("SHIMS: ApphelpCacheRead could not even open Session Manager\\AppCompatCache (0x%x)\n", Status);
 319         return FALSE;
 320     }
 321
 322     Status = ZwQueryValueKey(KeyHandle,
 323                              &AppCompatCacheValue,
 324                              KeyValuePartialInformation,
 325                              KeyValueInformation,
 326                              sizeof(KeyValueObject),
 327                              &ResultSize);
 328     if (Status == STATUS_BUFFER_OVERFLOW)
 329     {
 330         KeyInfoSize = sizeof(KEY_VALUE_PARTIAL_INFORMATION) + KeyValueInformation->DataLength;
 331         KeyValueInformation = ApphelpAlloc(KeyInfoSize);
 332         if (KeyValueInformation != NULL)
 333         {
 334             Status = ZwQueryValueKey(KeyHandle,
 335                                      &AppCompatCacheValue,
 336                                      KeyValuePartialInformation,
 337                                      KeyValueInformation,
 338                                      KeyInfoSize,
 339                                      &ResultSize);
 340         }
 341     }
 342
 343     if (NT_SUCCESS(Status) && KeyValueInformation->Type == REG_BINARY)
 344     {
 345         Status = ApphelpCacheParse(KeyValueInformation->Data,
 346                                    KeyValueInformation->DataLength);
 347     }
 348     else
 349     {
 350         DPRINT1("SHIMS: ApphelpCacheRead not loaded from registry (0x%x)\n", Status);
 351     }
 352
 353     if (KeyValueInformation != &KeyValueObject && KeyValueInformation != NULL)
 354     {
 355         ApphelpFree(KeyValueInformation);
 356     }
 357
 358     ZwClose(KeyHandle);
 359     return NT_SUCCESS(Status);
 360 }
 361
 362 BOOLEAN
 363 ApphelpCacheWrite(VOID)
 364 {
 365     ULONG Length = SHIM_CACHE_HEADER_SIZE;
 366     ULONG NumEntries = 0;
 367     PLIST_ENTRY ListEntry;
 368     PUCHAR Buffer, BufferNamePos;
 369     PSHIM_PERSISTENT_CACHE_HEADER Header;
 370     PSHIM_PERSISTENT_CACHE_ENTRY WriteEntry;
 371     HANDLE KeyHandle;
 372     NTSTATUS Status;
 373
 374     /* First we have to calculate the required size. */
 375     ApphelpCacheAcquireLock();
 376     ListEntry = ApphelpShimCacheAge.Flink;
 377     while (ListEntry != &ApphelpShimCacheAge)
 378     {
 379         PSHIM_CACHE_ENTRY Entry = CONTAINING_RECORD(ListEntry, SHIM_CACHE_ENTRY, List);
 380         Length += SHIM_PERSISTENT_CACHE_ENTRY_SIZE;
 381         Length += Entry->Persistent.ImageName.MaximumLength;
 382         ++NumEntries;
 383         ListEntry = ListEntry->Flink;
 384     }
 385     DPRINT("SHIMS: ApphelpCacheWrite, %d Entries, total size: %d\n", NumEntries, Length);
 386     Length = ROUND_UP(Length, sizeof(ULONGLONG));
 387     DPRINT("SHIMS: ApphelpCacheWrite, Rounded to: %d\n", Length);
 388
 389     /* Now we allocate and prepare some helpers */
 390     Buffer = ApphelpAlloc(Length);
 391     BufferNamePos = Buffer + Length;
 392     Header = (PSHIM_PERSISTENT_CACHE_HEADER)Buffer;
 393     WriteEntry = (PSHIM_PERSISTENT_CACHE_ENTRY)(Buffer + SHIM_CACHE_HEADER_SIZE);
 394
 395     Header->Magic = SHIM_CACHE_MAGIC;
 396     Header->NumEntries = NumEntries;
 397
 398     ListEntry = ApphelpShimCacheAge.Flink;
 399     while (ListEntry != &ApphelpShimCacheAge)
 400     {
 401         PSHIM_CACHE_ENTRY Entry = CONTAINING_RECORD(ListEntry, SHIM_CACHE_ENTRY, List);
 402         USHORT ImageNameLen = Entry->Persistent.ImageName.MaximumLength;
 403         /* Copy the Persistent structure over */
 404         *WriteEntry = Entry->Persistent;
 405         BufferNamePos -= ImageNameLen;
 406         /* Copy the image name over */
 407         RtlCopyMemory(BufferNamePos, Entry->Persistent.ImageName.Buffer, ImageNameLen);
 408         /* Fix the Persistent structure, so that Buffer is once again an offset */
 409         WriteEntry->ImageName.Buffer = (PWCH)(BufferNamePos - Buffer);
 410
 411         ++WriteEntry;
 412         ListEntry = ListEntry->Flink;
 413     }
 414     ApphelpCacheReleaseLock();
 415
 416     Status = ZwOpenKey(&KeyHandle, KEY_SET_VALUE, &AppCompatKeyAttributes);
 417     if (NT_SUCCESS(Status))
 418     {
 419         Status = ZwSetValueKey(KeyHandle,
 420                                &AppCompatCacheValue,
 421                                0,
 422                                REG_BINARY,
 423                                Buffer,
 424                                Length);
 425         ZwClose(KeyHandle);
 426     }
 427     else
 428     {
 429         DPRINT1("SHIMS: ApphelpCacheWrite could not even open Session Manager\\AppCompatCache (0x%x)\n", Status);
 430     }
 431
 432     ApphelpFree(Buffer);
 433     return NT_SUCCESS(Status);
 434 }
 435
 436
 437 NTSTATUS
 438 NTAPI
 439 INIT_FUNCTION
 440 ApphelpCacheInitialize(VOID)
 441 {
 442     DPRINT("SHIMS: ApphelpCacheInitialize\n");
 443     /* If we are booting in safemode we do not want to use the apphelp cache */
 444     if (InitSafeBootMode)
 445     {
 446         DPRINT1("SHIMS: Safe mode detected, disabling cache.\n");
 447         ApphelpCacheEnabled = FALSE;
 448     }
 449     else
 450     {
 451         ExInitializeResourceLite(&ApphelpCacheLock);
 452         RtlInitializeGenericTableAvl(&ApphelpShimCache,
 453                                      ApphelpShimCacheCompareRoutine,
 454                                      ApphelpShimCacheAllocateRoutine,
 455                                      ApphelpShimCacheFreeRoutine,
 456                                      NULL);
 457         InitializeListHead(&ApphelpShimCacheAge);
 458         ApphelpCacheEnabled = ApphelpCacheRead();
 459     }
 460     DPRINT("SHIMS: ApphelpCacheInitialize: %d\n", ApphelpCacheEnabled);
 461     return STATUS_SUCCESS;
 462 }
 463
 464 VOID
 465 NTAPI
 466 ApphelpCacheShutdown(VOID)
 467 {
 468     if (ApphelpCacheEnabled)
 469     {
 470         ApphelpCacheWrite();
 471     }
 472 }
 473
 474 NTSTATUS
 475 ApphelpValidateData(
 476     _In_opt_ PAPPHELP_CACHE_SERVICE_LOOKUP ServiceData,
 477     _Out_ PUNICODE_STRING ImageName,
 478     _Out_ PHANDLE ImageHandle)
 479 {
 480     NTSTATUS Status = STATUS_INVALID_PARAMETER;
 481
 482     if (ServiceData)
 483     {
 484         UNICODE_STRING LocalImageName;
 485         _SEH2_TRY
 486         {
 487             ProbeForRead(ServiceData,
 488                          sizeof(APPHELP_CACHE_SERVICE_LOOKUP),
 489                          sizeof(ULONG));
 490             LocalImageName = ServiceData->ImageName;
 491             *ImageHandle = ServiceData->ImageHandle;
 492             if (LocalImageName.Length && LocalImageName.Buffer)
 493             {
 494                 ProbeForRead(LocalImageName.Buffer,
 495                              LocalImageName.Length * sizeof(WCHAR),
 496                              1);
 497                 ApphelpDuplicateUnicodeString(ImageName, &LocalImageName);
 498                 Status = STATUS_SUCCESS;
 499             }
 500         }
 501         _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
 502         {
 503             _SEH2_YIELD(return _SEH2_GetExceptionCode());
 504         }
 505         _SEH2_END;
 506     }
 507     if (!NT_SUCCESS(Status))
 508     {
 509         DPRINT1("SHIMS: ApphelpValidateData: invalid data passed\n");
 510     }
 511     return Status;
 512 }
 513
 514 NTSTATUS
 515 ApphelpCacheRemoveEntryNolock(
 516     _In_ PSHIM_CACHE_ENTRY Entry)
 517 {
 518     if (Entry)
 519     {
 520         PWSTR Buffer = Entry->Persistent.ImageName.Buffer;
 521         RemoveEntryList(&Entry->List);
 522         if (RtlDeleteElementGenericTableAvl(&ApphelpShimCache, Entry))
 523         {
 524             ApphelpFree(Buffer);
 525         }
 526         return STATUS_SUCCESS;
 527     }
 528     return STATUS_NOT_FOUND;
 529 }
 530
 531 NTSTATUS
 532 ApphelpCacheLookupEntry(
 533     _In_ PUNICODE_STRING ImageName,
 534     _In_ HANDLE ImageHandle)
 535 {
 536     NTSTATUS Status = STATUS_NOT_FOUND;
 537     SHIM_CACHE_ENTRY Lookup = EMPTY_SHIM_ENTRY;
 538     PSHIM_CACHE_ENTRY Entry;
 539
 540     if (!ApphelpCacheTryAcquireLock())
 541     {
 542         return Status;
 543     }
 544
 545     Lookup.Persistent.ImageName = *ImageName;
 546     Entry = RtlLookupElementGenericTableAvl(&ApphelpShimCache, &Lookup);
 547     if (Entry == NULL)
 548     {
 549         DPRINT("SHIMS: ApphelpCacheLookupEntry: could not find %wZ\n", ImageName);
 550         goto Cleanup;
 551     }
 552
 553     DPRINT("SHIMS: ApphelpCacheLookupEntry: found %wZ\n", ImageName);
 554     if (ImageHandle == INVALID_HANDLE_VALUE)
 555     {
 556         DPRINT("SHIMS: ApphelpCacheLookupEntry: ok\n");
 557         /* just return if we know it, do not query file info */
 558         Status = STATUS_SUCCESS;
 559     }
 560     else
 561     {
 562         Status = ApphelpCacheQueryInfo(ImageHandle, &Lookup);
 563         if (NT_SUCCESS(Status) &&
 564             Lookup.Persistent.DateTime.QuadPart == Entry->Persistent.DateTime.QuadPart &&
 565             Lookup.Persistent.FileSize.QuadPart == Entry->Persistent.FileSize.QuadPart)
 566         {
 567             DPRINT("SHIMS: ApphelpCacheLookupEntry: found & validated\n");
 568             Status = STATUS_SUCCESS;
 569             /* move it to the front to keep it alive */
 570             RemoveEntryList(&Entry->List);
 571             InsertHeadList(&ApphelpShimCacheAge, &Entry->List);
 572         }
 573         else
 574         {
 575             DPRINT1("SHIMS: ApphelpCacheLookupEntry: file info mismatch (%lx)\n", Status);
 576             Status = STATUS_NOT_FOUND;
 577             /* Could not read file info, or it did not match, drop it from the cache */
 578             ApphelpCacheRemoveEntryNolock(Entry);
 579         }
 580     }
 581
 582 Cleanup:
 583     ApphelpCacheReleaseLock();
 584     return Status;
 585 }
 586
 587 NTSTATUS
 588 ApphelpCacheRemoveEntry(
 589     _In_ PUNICODE_STRING ImageName)
 590 {
 591     PSHIM_CACHE_ENTRY Entry;
 592     NTSTATUS Status;
 593
 594     ApphelpCacheAcquireLock();
 595     Entry = RtlLookupElementGenericTableAvl(&ApphelpShimCache, ImageName);
 596     Status = ApphelpCacheRemoveEntryNolock(Entry);
 597     ApphelpCacheReleaseLock();
 598     return Status;
 599 }
 600
 601 /* Validate that we are either called from r0, or from a service-like context */
 602 NTSTATUS
 603 ApphelpCacheAccessCheck(VOID)
 604 {
 605     if (ExGetPreviousMode() != KernelMode)
 606     {
 607         if (!SeSinglePrivilegeCheck(SeTcbPrivilege, UserMode))
 608         {
 609             DPRINT1("SHIMS: ApphelpCacheAccessCheck failed\n");
 610             return STATUS_ACCESS_DENIED;
 611         }
 612     }
 613     return STATUS_SUCCESS;
 614 }
 615
 616 NTSTATUS
 617 ApphelpCacheUpdateEntry(
 618     _In_ PUNICODE_STRING ImageName,
 619     _In_ HANDLE ImageHandle)
 620 {
 621     NTSTATUS Status = STATUS_SUCCESS;
 622     SHIM_CACHE_ENTRY Entry = EMPTY_SHIM_ENTRY;
 623     PSHIM_CACHE_ENTRY Lookup;
 624     PVOID NodeOrParent;
 625     TABLE_SEARCH_RESULT SearchResult;
 626
 627     ApphelpCacheAcquireLock();
 628
 629     /* If we got a file handle, query it for info */
 630     if (ImageHandle != INVALID_HANDLE_VALUE)
 631     {
 632         Status = ApphelpCacheQueryInfo(ImageHandle, &Entry);
 633         if (!NT_SUCCESS(Status))
 634         {
 635             goto Cleanup;
 636         }
 637     }
 638
 639     /* Use ImageName for the lookup, don't actually duplicate it */
 640     Entry.Persistent.ImageName = *ImageName;
 641     Lookup = RtlLookupElementGenericTableFullAvl(&ApphelpShimCache,
 642                                                  &Entry,
 643                                                  &NodeOrParent, &SearchResult);
 644     if (Lookup)
 645     {
 646         DPRINT("SHIMS: ApphelpCacheUpdateEntry: Entry already exists, reusing it\n");
 647         /* Unlink the found item, so we can put it back at the front,
 648             and copy the earlier obtained file info*/
 649         RemoveEntryList(&Lookup->List);
 650         Lookup->Persistent.DateTime = Entry.Persistent.DateTime;
 651         Lookup->Persistent.FileSize = Entry.Persistent.FileSize;
 652     }
 653     else
 654     {
 655         DPRINT("SHIMS: ApphelpCacheUpdateEntry: Inserting new Entry\n");
 656         /* Insert a new entry, with its own copy of the ImageName */
 657         ApphelpDuplicateUnicodeString(&Entry.Persistent.ImageName, ImageName);
 658         Lookup = RtlInsertElementGenericTableFullAvl(&ApphelpShimCache,
 659                                                      &Entry,
 660                                                      sizeof(Entry),
 661                                                      0,
 662                                                      NodeOrParent,
 663                                                      SearchResult);
 664         if (!Lookup)
 665         {
 666             ApphelpFreeUnicodeString(&Entry.Persistent.ImageName);
 667             Status = STATUS_NO_MEMORY;
 668         }
 669     }
 670     if (Lookup)
 671     {
 672         /* Either we re-used an existing item, or we inserted a new one, keep it alive */
 673         InsertHeadList(&ApphelpShimCacheAge, &Lookup->List);
 674         if (RtlNumberGenericTableElementsAvl(&ApphelpShimCache) > MAX_SHIM_ENTRIES)
 675         {
 676             PSHIM_CACHE_ENTRY Remove;
 677             DPRINT1("SHIMS: ApphelpCacheUpdateEntry: Cache growing too big, dropping oldest item\n");
 678             Remove = CONTAINING_RECORD(ApphelpShimCacheAge.Blink, SHIM_CACHE_ENTRY, List);
 679             Status = ApphelpCacheRemoveEntryNolock(Remove);
 680         }
 681     }
 682
 683 Cleanup:
 684     ApphelpCacheReleaseLock();
 685     return Status;
 686 }
 687
 688 NTSTATUS
 689 ApphelpCacheFlush(VOID)
 690 {
 691     PVOID p;
 692
 693     DPRINT1("SHIMS: ApphelpCacheFlush\n");
 694     ApphelpCacheAcquireLock();
 695     while ((p = RtlEnumerateGenericTableAvl(&ApphelpShimCache, TRUE)))
 696     {
 697         ApphelpCacheRemoveEntryNolock((PSHIM_CACHE_ENTRY)p);
 698     }
 699     ApphelpCacheReleaseLock();
 700     return STATUS_SUCCESS;
 701 }
 702
 703 NTSTATUS
 704 ApphelpCacheDump(VOID)
 705 {
 706     PLIST_ENTRY ListEntry;
 707     PSHIM_CACHE_ENTRY Entry;
 708
 709     DPRINT1("SHIMS: NtApphelpCacheControl( Dumping entries, newest to oldest )\n");
 710     ApphelpCacheAcquireLock();
 711     ListEntry = ApphelpShimCacheAge.Flink;
 712     while (ListEntry != &ApphelpShimCacheAge)
 713     {
 714         Entry = CONTAINING_RECORD(ListEntry, SHIM_CACHE_ENTRY, List);
 715         DPRINT1("Entry: %wZ\n", &Entry->Persistent.ImageName);
 716         DPRINT1("DateTime: 0x%I64x\n", Entry->Persistent.DateTime.QuadPart);
 717         DPRINT1("FileSize: 0x%I64x\n", Entry->Persistent.FileSize.QuadPart);
 718         DPRINT1("Flags: 0x%x\n", Entry->CompatFlags);
 719         ListEntry = ListEntry->Flink;
 720     }
 721     ApphelpCacheReleaseLock();
 722     return STATUS_SUCCESS;
 723 }
 724
 725 /* PUBLIC FUNCTIONS **********************************************************/
 726
 727 NTSTATUS
 728 NTAPI
 729 NtApphelpCacheControl(
 730     _In_ APPHELPCACHESERVICECLASS Service,
 731     _In_opt_ PAPPHELP_CACHE_SERVICE_LOOKUP ServiceData)
 732 {
 733     NTSTATUS Status = STATUS_INVALID_PARAMETER;
 734     UNICODE_STRING ImageName = { 0 };
 735     HANDLE Handle = INVALID_HANDLE_VALUE;
 736
 737     if (!ApphelpCacheEnabled)
 738     {
 739         DPRINT1("NtApphelpCacheControl: ApphelpCacheEnabled == 0\n");
 740         return Status;
 741     }
 742     switch (Service)
 743     {
 744         case ApphelpCacheServiceLookup:
 745             DPRINT("SHIMS: NtApphelpCacheControl( ApphelpCacheServiceLookup )\n");
 746             Status = ApphelpValidateData(ServiceData, &ImageName, &Handle);
 747             if (NT_SUCCESS(Status))
 748                 Status = ApphelpCacheLookupEntry(&ImageName, Handle);
 749             break;
 750         case ApphelpCacheServiceRemove:
 751             DPRINT("SHIMS: NtApphelpCacheControl( ApphelpCacheServiceRemove )\n");
 752             Status = ApphelpValidateData(ServiceData, &ImageName, &Handle);
 753             if (NT_SUCCESS(Status))
 754                 Status = ApphelpCacheRemoveEntry(&ImageName);
 755             break;
 756         case ApphelpCacheServiceUpdate:
 757             DPRINT("SHIMS: NtApphelpCacheControl( ApphelpCacheServiceUpdate )\n");
 758             Status = ApphelpCacheAccessCheck();
 759             if (NT_SUCCESS(Status))
 760             {
 761                 Status = ApphelpValidateData(ServiceData, &ImageName, &Handle);
 762                 if (NT_SUCCESS(Status))
 763                     Status = ApphelpCacheUpdateEntry(&ImageName, Handle);
 764             }
 765             break;
 766         case ApphelpCacheServiceFlush:
 767             Status = ApphelpCacheFlush();
 768             break;
 769         case ApphelpCacheServiceDump:
 770             Status = ApphelpCacheDump();
 771             break;
 772         case ApphelpDBGReadRegistry:
 773             DPRINT1("SHIMS: NtApphelpCacheControl( ApphelpDBGReadRegistry ): flushing cache.\n");
 774             ApphelpCacheFlush();
 775             DPRINT1("SHIMS: NtApphelpCacheControl( ApphelpDBGReadRegistry ): reading cache.\n");
 776             Status = ApphelpCacheRead() ? STATUS_SUCCESS : STATUS_NOT_FOUND;
 777             break;
 778         case ApphelpDBGWriteRegistry:
 779             DPRINT1("SHIMS: NtApphelpCacheControl( ApphelpDBGWriteRegistry ): writing cache.\n");
 780             Status = ApphelpCacheWrite() ? STATUS_SUCCESS : STATUS_NOT_FOUND;
 781             break;
 782         default:
 783             DPRINT1("SHIMS: NtApphelpCacheControl( Invalid service requested )\n");
 784             break;
 785     }
 786     if (ImageName.Buffer)
 787     {
 788         ApphelpFreeUnicodeString(&ImageName);
 789     }
 790     return Status;
 791 }
 792