[FREELDR][NTOS:MM] Add security cookie generation to FreeLoader (#6270)

* [NTOS:MM] Misc improvements for cookie generation code

- Improve support for 64 bit images
- Improve LdrpFetchAddressOfSecurityCookie code

* [FREELDR] Add security cookie generation to FreeLoader

CORE-17808

boot/freeldr/freeldr/include/peloader.h

@@ -56,3 +56,7 @@ PeLdrCheckForLoadedDll(
     IN OUT PLIST_ENTRY ModuleListHead,
     IN PCH DllName,
     OUT PLDR_DATA_TABLE_ENTRY *LoadedEntry);
+
+PVOID
+PeLdrInitSecurityCookie(
+    _In_ PLDR_DATA_TABLE_ENTRY LdrEntry);

boot/freeldr/freeldr/lib/peloader.c

@@ -27,9 +27,51 @@ DBG_DEFAULT_CHANNEL(PELOADER);
 
 PELDR_IMPORTDLL_LOAD_CALLBACK PeLdrImportDllLoadCallback = NULL;
 
+#ifdef _WIN64
+#define COOKIE_MAX 0x0000FFFFFFFFFFFFll
+#define DEFAULT_SECURITY_COOKIE 0x00002B992DDFA232ll
+#else
+#define DEFAULT_SECURITY_COOKIE 0xBB40E64E
+#endif
+
 
 /* PRIVATE FUNCTIONS *********************************************************/
 
+static PVOID
+PeLdrpFetchAddressOfSecurityCookie(PVOID BaseAddress, ULONG SizeOfImage)
+{
+    PIMAGE_LOAD_CONFIG_DIRECTORY ConfigDir;
+    ULONG DirSize;
+    PULONG_PTR Cookie = NULL;
+
+    /* Get the pointer to the config directory */
+    ConfigDir = RtlImageDirectoryEntryToData(BaseAddress,
+                                             TRUE,
+                                             IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG,
+                                             &DirSize);
+
+    /* Check for sanity */
+    if (!ConfigDir ||
+        DirSize < RTL_SIZEOF_THROUGH_FIELD(IMAGE_LOAD_CONFIG_DIRECTORY, SecurityCookie))
+    {
+        /* Invalid directory*/
+        return NULL;
+    }
+
+    /* Now get the cookie */
+    Cookie = VaToPa((PULONG_PTR)ConfigDir->SecurityCookie);
+
+    /* Check this cookie */
+    if ((PCHAR)Cookie <= (PCHAR)BaseAddress ||
+        (PCHAR)Cookie >= (PCHAR)BaseAddress + SizeOfImage - sizeof(*Cookie))
+    {
+        Cookie = NULL;
+    }
+
+    /* Return validated security cookie */
+    return Cookie;
+}
+
 /* DllName - physical, UnicodeString->Buffer - virtual */
 static BOOLEAN
 PeLdrpCompareDllName(
@@ -386,6 +428,9 @@ PeLdrpLoadAndScanReferencedDll(
         return Success;
     }
 
+    /* Init security cookie */
+    PeLdrInitSecurityCookie(*DataTableEntry);
+
     (*DataTableEntry)->Flags |= LDRP_DRIVER_DEPENDENT_DLL;
 
     /* Scan its dependencies too */
@@ -473,6 +518,44 @@ PeLdrpScanImportAddressTable(
 
 /* FUNCTIONS *****************************************************************/
 
+PVOID
+PeLdrInitSecurityCookie(PLDR_DATA_TABLE_ENTRY LdrEntry)
+{
+    PULONG_PTR Cookie;
+    ULONG_PTR NewCookie;
+
+    /* Fetch address of the cookie */
+    Cookie = PeLdrpFetchAddressOfSecurityCookie(VaToPa(LdrEntry->DllBase), LdrEntry->SizeOfImage);
+
+    if (!Cookie)
+        return NULL;
+
+    /* Check if it's a default one */
+    if ((*Cookie == DEFAULT_SECURITY_COOKIE) ||
+        (*Cookie == 0))
+    {
+        /* Generate new cookie using cookie address and time as seed */
+        NewCookie = (ULONG_PTR)Cookie ^ (ULONG_PTR)ArcGetRelativeTime();
+#ifdef _WIN64
+        /* Some images expect first 16 bits to be kept clean (like in default cookie) */
+        if (NewCookie > COOKIE_MAX)
+        {
+            NewCookie >>= 16;
+        }
+#endif
+        /* If the result is 0 or the same as we got, just add one to the default value */
+        if ((NewCookie == 0) || (NewCookie == *Cookie))
+        {
+            NewCookie = DEFAULT_SECURITY_COOKIE + 1;
+        }
+
+        /* Set the new cookie value */
+        *Cookie = NewCookie;
+    }
+
+    return Cookie;
+}
+
 /* Returns TRUE if DLL has already been loaded - looks in LoadOrderList in LPB */
 BOOLEAN
 PeLdrCheckForLoadedDll(

boot/freeldr/freeldr/ntldr/winldr.c

@@ -342,6 +342,9 @@ WinLdrLoadDeviceDriver(PLIST_ENTRY LoadOrderListHead,
         return FALSE;
     }
 
+    /* Init security cookie */
+    PeLdrInitSecurityCookie(*DriverDTE);
+
     // Modify any flags, if needed
     (*DriverDTE)->Flags |= Flags;
 
@@ -537,9 +540,12 @@ LoadModule(
         /* Cleanup and bail out */
         ERR("PeLdrAllocateDataTableEntry('%s') failed\n", FullFileName);
         MmFreeMemory(BaseAddress);
-        BaseAddress = NULL;
+        return NULL;
     }
 
+    /* Init security cookie */
+    PeLdrInitSecurityCookie(*Dte);
+
     return BaseAddress;
 }
 

ntoskrnl/mm/ARM3/sysldr.c

@@ -38,6 +38,7 @@ ULONG_PTR ExPoolCodeStart, ExPoolCodeEnd, MmPoolCodeStart, MmPoolCodeEnd;
 ULONG_PTR MmPteCodeStart, MmPteCodeEnd;
 
 #ifdef _WIN64
+#define COOKIE_MAX 0x0000FFFFFFFFFFFFll
 #define DEFAULT_SECURITY_COOKIE 0x00002B992DDFA232ll
 #else
 #define DEFAULT_SECURITY_COOKIE 0xBB40E64E
@@ -2859,10 +2860,7 @@ LdrpFetchAddressOfSecurityCookie(PVOID BaseAddress, ULONG SizeOfImage)
 {
     PIMAGE_LOAD_CONFIG_DIRECTORY ConfigDir;
     ULONG DirSize;
-    PVOID Cookie = NULL;
-
-    /* Check NT header first */
-    if (!RtlImageNtHeader(BaseAddress)) return NULL;
+    PULONG_PTR Cookie = NULL;
 
     /* Get the pointer to the config directory */
     ConfigDir = RtlImageDirectoryEntryToData(BaseAddress,
@@ -2872,19 +2870,18 @@ LdrpFetchAddressOfSecurityCookie(PVOID BaseAddress, ULONG SizeOfImage)
 
     /* Check for sanity */
     if (!ConfigDir ||
-        DirSize < FIELD_OFFSET(IMAGE_LOAD_CONFIG_DIRECTORY, SEHandlerTable) ||  /* SEHandlerTable is after SecurityCookie */
-        (ConfigDir->Size != DirSize))
+        DirSize < RTL_SIZEOF_THROUGH_FIELD(IMAGE_LOAD_CONFIG_DIRECTORY, SecurityCookie))
     {
         /* Invalid directory*/
         return NULL;
     }
 
     /* Now get the cookie */
-    Cookie = (PVOID)ConfigDir->SecurityCookie;
+    Cookie = (PULONG_PTR)ConfigDir->SecurityCookie;
 
     /* Check this cookie */
     if ((PCHAR)Cookie <= (PCHAR)BaseAddress ||
-        (PCHAR)Cookie >= (PCHAR)BaseAddress + SizeOfImage)
+        (PCHAR)Cookie >= (PCHAR)BaseAddress + SizeOfImage - sizeof(*Cookie))
     {
         Cookie = NULL;
     }
@@ -2903,28 +2900,34 @@ LdrpInitSecurityCookie(PLDR_DATA_TABLE_ENTRY LdrEntry)
     /* Fetch address of the cookie */
     Cookie = LdrpFetchAddressOfSecurityCookie(LdrEntry->DllBase, LdrEntry->SizeOfImage);
 
-    if (Cookie)
+    if (!Cookie)
+        return NULL;
+    
+    /* Check if it's a default one */
+    if ((*Cookie == DEFAULT_SECURITY_COOKIE) ||
+        (*Cookie == 0))
     {
-        /* Check if it's a default one */
-        if ((*Cookie == DEFAULT_SECURITY_COOKIE) ||
-            (*Cookie == 0))
+        LARGE_INTEGER Counter = KeQueryPerformanceCounter(NULL);
+        /* The address should be unique */
+        NewCookie = (ULONG_PTR)Cookie;
+
+        /* We just need a simple tick, don't care about precision and whatnot */
+        NewCookie ^= (ULONG_PTR)Counter.LowPart;
+#ifdef _WIN64
+        /* Some images expect first 16 bits to be kept clean (like in default cookie) */
+        if (NewCookie > COOKIE_MAX)
         {
-            LARGE_INTEGER Counter = KeQueryPerformanceCounter(NULL);
-            /* The address should be unique */
-            NewCookie = (ULONG_PTR)Cookie;
-
-            /* We just need a simple tick, don't care about precision and whatnot */
-            NewCookie ^= (ULONG_PTR)Counter.LowPart;
-
-            /* If the result is 0 or the same as we got, just add one to the default value */
-            if ((NewCookie == 0) || (NewCookie == *Cookie))
-            {
-                NewCookie = DEFAULT_SECURITY_COOKIE + 1;
-            }
-
-            /* Set the new cookie value */
-            *Cookie = NewCookie;
+            NewCookie >>= 16;
         }
+#endif
+        /* If the result is 0 or the same as we got, just add one to the default value */
+        if ((NewCookie == 0) || (NewCookie == *Cookie))
+        {
+            NewCookie = DEFAULT_SECURITY_COOKIE + 1;
+        }
+
+        /* Set the new cookie value */
+        *Cookie = NewCookie; 
     }
 
     return Cookie;