- Implement enable/disable of protection for PE sections in ntdll

- Finally get rid of duplicated relocation code in ntdll by using the relocate function from RTL and enable/disable protection
- Some improvements added to RTL's relocation function, it's split up into two functions now

svn path=/trunk/; revision=24020

reactos/dll/ntdll/ldr/utils.c

@@ -1198,126 +1198,89 @@ LdrGetExportByName(PVOID BaseAddress,
    return (PVOID)NULL;
 }
 
-
+/*++
-/**********************************************************************
+* @name LdrpSetProtection
- * NAME                                                         LOCAL
+*
- *      LdrPerformRelocations
+*     Local routine, used to set or reset pages protection.
- *
+*     Used during PE relocation process.
- * DESCRIPTION
+*
- *      Relocate a DLL's memory image.
+* @param BaseAddress
- *
+*        Base address of the image.
- * ARGUMENTS
+*
- *
+* @param Protection
- * RETURN VALUE
+*        FALSE to disable protection, TRUE to enable it
- *
+*
- * REVISIONS
+* @return Usual NTSTATUS.
- *
+*
- * NOTE
+* @remarks None.
- *
+*
- */
+*--*/
-static NTSTATUS
+NTSTATUS
-LdrPerformRelocations(PIMAGE_NT_HEADERS NTHeaders,
+LdrpSetProtection(IN PVOID BaseAddress,
-                      PVOID ImageBase)
+                  IN BOOLEAN Protection)
 {
-  PIMAGE_DATA_DIRECTORY RelocationDDir;
+    PIMAGE_NT_HEADERS NtHeaders;
-  PIMAGE_BASE_RELOCATION RelocationDir, RelocationEnd;
+    PIMAGE_SECTION_HEADER SectionHeader;
-  ULONG Count, ProtectSize, OldProtect, OldProtect2;
+    ULONG i;
-  PVOID Page, ProtectPage, ProtectPage2;
+    PVOID Address;
-  PUSHORT TypeOffset;
+    SIZE_T Size;
-  ULONG_PTR Delta;
+    NTSTATUS Status;
-  NTSTATUS Status;
+    ULONG NewAccess, OldAccess;
 
-  if (NTHeaders->FileHeader.Characteristics & IMAGE_FILE_RELOCS_STRIPPED)
+    NtHeaders = RtlImageNtHeader(BaseAddress);
+
+    /* Ensure image is valid */
+    if (NtHeaders == NULL)
+        return STATUS_INVALID_IMAGE_FORMAT;
+
+    /* Get the first section header */
+    SectionHeader = IMAGE_FIRST_SECTION(NtHeaders);
+
+    /* Loop through every section */
+    for (i = 0; i < NtHeaders->FileHeader.NumberOfSections; i++, SectionHeader++)
     {
-      return STATUS_UNSUCCESSFUL;
+        /* If section has no data - just skip to the next section */
-    }
+        if (SectionHeader->SizeOfRawData == 0)
+            continue;
 
-  RelocationDDir =
+        /* We need to act only if section isn't RW */
-    &NTHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
+        if (!(SectionHeader->Characteristics & IMAGE_SCN_MEM_WRITE))
-
-  if (RelocationDDir->VirtualAddress == 0 || RelocationDDir->Size == 0)
-    {
-      return STATUS_SUCCESS;
-    }
-
-  ProtectSize = PAGE_SIZE;
-  Delta = (ULONG_PTR)ImageBase - NTHeaders->OptionalHeader.ImageBase;
-  RelocationDir = (PIMAGE_BASE_RELOCATION)((ULONG_PTR)ImageBase +
-                  RelocationDDir->VirtualAddress);
-  RelocationEnd = (PIMAGE_BASE_RELOCATION)((ULONG_PTR)ImageBase +
-                  RelocationDDir->VirtualAddress + RelocationDDir->Size);
-
-  while (RelocationDir < RelocationEnd &&
-         RelocationDir->SizeOfBlock > 0)
-    {
-      Count = (RelocationDir->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) /
-              sizeof(USHORT);
-      Page = (PVOID)((ULONG_PTR)ImageBase + (ULONG_PTR)RelocationDir->VirtualAddress);
-      TypeOffset = (PUSHORT)(RelocationDir + 1);
-
-      /* Unprotect the page(s) we're about to relocate. */
-      ProtectPage = Page;
-      Status = NtProtectVirtualMemory(NtCurrentProcess(),
-                                      &ProtectPage,
-                                      &ProtectSize,
-                                      PAGE_READWRITE,
-                                      &OldProtect);
-      if (!NT_SUCCESS(Status))
         {
-          DPRINT1("Failed to unprotect relocation target.\n");
+            if (Protection)
-          return Status;
-        }
-
-      if (RelocationDir->VirtualAddress + PAGE_SIZE <
-          NTHeaders->OptionalHeader.SizeOfImage)
-        {
-          ProtectPage2 = (PVOID)((ULONG_PTR)ProtectPage + PAGE_SIZE);
-          Status = NtProtectVirtualMemory(NtCurrentProcess(),
-                                          &ProtectPage2,
-                                          &ProtectSize,
-                                          PAGE_READWRITE,
-                                          &OldProtect2);
-          if (!NT_SUCCESS(Status))
             {
-              DPRINT1("Failed to unprotect relocation target (2).\n");
+                /* Check if the section was also executable, then we should
-              NtProtectVirtualMemory(NtCurrentProcess(),
+                   give it PAGE_EXECUTE, otherwise just PAGE_READONLY */
-                                     &ProtectPage,
+                NewAccess = (SectionHeader->Characteristics & IMAGE_SCN_MEM_EXECUTE) ? PAGE_EXECUTE : PAGE_READONLY;
-                                     &ProtectSize,
+
-                                     OldProtect,
+                /* Also check caching */
-                                     &OldProtect);
+                NewAccess |= (SectionHeader->Characteristics & IMAGE_SCN_MEM_NOT_CACHED) ? PAGE_NOCACHE : 0;
-              return Status;
+            }
+            else
+            {
+                /* Set access to RW */
+                NewAccess = PAGE_READWRITE;
             }
-        }
-      else
-        {
-          ProtectPage2 = NULL;
-        }
 
-      RelocationDir = LdrProcessRelocationBlock((ULONG_PTR)Page,
+            /* Get the address and size of the section, and set protectin accordingly */
-                                                Count,
+            Address = (PVOID)((ULONG_PTR)BaseAddress + SectionHeader->VirtualAddress);
-                                                TypeOffset,
+            Size = SectionHeader->SizeOfRawData;
-                                                Delta);
-      if (RelocationDir == NULL)
-        return STATUS_UNSUCCESSFUL;
 
-      /* Restore old page protection. */
+            Status = NtProtectVirtualMemory(NtCurrentProcess(),
-      NtProtectVirtualMemory(NtCurrentProcess(),
+                &Address,
-                             &ProtectPage,
+                &Size,
-                             &ProtectSize,
+                NewAccess,
-                             OldProtect,
+                &OldAccess);
-                             &OldProtect);
 
-      if (ProtectPage2 != NULL)
+            if (!NT_SUCCESS(Status))
-        {
+                return Status;
-          NtProtectVirtualMemory(NtCurrentProcess(),
-                                 &ProtectPage2,
-                                 &ProtectSize,
-                                 OldProtect2,
-                                 &OldProtect2);
         }
     }
 
-  return STATUS_SUCCESS;
+    /* Instruction cache should be flushed if we enabled protection */
+    if (Protection)
+        NtFlushInstructionCache(NtCurrentProcess(), NULL, 0);
+
+    return STATUS_SUCCESS;
 }
 
 static NTSTATUS
@@ -1898,12 +1861,32 @@ PEPFUNC LdrPEStartup (PVOID  ImageBase,
    if (ImageBase != (PVOID) NTHeaders->OptionalHeader.ImageBase)
      {
        DPRINT("LDR: Performing relocations\n");
-       Status = LdrPerformRelocations(NTHeaders, ImageBase);
+
+       /* Remove protection */
+       Status = LdrpSetProtection(ImageBase, FALSE);
        if (!NT_SUCCESS(Status))
          {
-           DPRINT1("LdrPerformRelocations() failed\n");
+           DPRINT1("LdrpSetProtection() failed with Status=0x%08X\n", Status);
            return NULL;
          }
+
+       /* Relocate */
+       Status = LdrRelocateImageWithBias(ImageBase, 0, "", STATUS_SUCCESS,
+           STATUS_CONFLICTING_ADDRESSES, STATUS_INVALID_IMAGE_FORMAT);
+       if (!NT_SUCCESS(Status))
+         {
+           DPRINT1("LdrRelocateImageWithBias() failed with Status=0x%08X\n", Status);
+           return NULL;
+         }
+
+       /* Set protection back */
+       Status = LdrpSetProtection(ImageBase, TRUE);
+       if (!NT_SUCCESS(Status))
+         {
+           DPRINT1("LdrpSetProtection() failed with Status=0x%08X\n", Status);
+           return NULL;
+         }
+
      }
 
    if (Module != NULL)
@@ -2078,10 +2061,13 @@ LdrpLoadModule(IN PWSTR SearchPath OPTIONAL,
           {
             DPRINT1("Relocating (%lx -> %p) %wZ\n",
               NtHeaders->OptionalHeader.ImageBase, ImageBase, &FullDosName);
-            Status = LdrPerformRelocations(NtHeaders, ImageBase);
+
+            Status = LdrRelocateImageWithBias(ImageBase, 0, "", STATUS_SUCCESS,
+                STATUS_CONFLICTING_ADDRESSES, STATUS_INVALID_IMAGE_FORMAT);
+
             if (!NT_SUCCESS(Status))
               {
-                DPRINT1("LdrPerformRelocations() failed\n");
+                DPRINT1("LdrRelocateImageWithBias() failed\n");
                 NtUnmapViewOfSection (NtCurrentProcess (), ImageBase);
                 NtClose (SectionHandle);
                 RtlFreeUnicodeString(&FullDosName);

reactos/lib/rtl/image.c

@@ -3,8 +3,10 @@
  * FILE:            lib/rtl/image.c
  * PURPOSE:         Image handling functions
  *                  Relocate functions were previously located in
- *                  ntoskrnl/ldr/loader.c
+ *                  ntoskrnl/ldr/loader.c and
- * PROGRAMMER:      Eric Kohl + original authors from loader.c file
+ *                  dll/ntdll/ldr/utils.c files
+ * PROGRAMMER:      Eric Kohl + original authors from loader.c and utils.c file
+ *                  Aleksey Bragin
  */
 
 /* INCLUDES *****************************************************************/
@@ -149,6 +151,70 @@ RtlImageRvaToVa (
 	               (ULONG_PTR)Section->VirtualAddress);
 }
 
+PIMAGE_BASE_RELOCATION
+LdrProcessRelocationBlockLongLong(
+    IN ULONG_PTR Address,
+    IN ULONG Count,
+    IN PUSHORT TypeOffset,
+    IN LONGLONG Delta
+    )
+{
+    SHORT Offset;
+    USHORT Type;
+    USHORT i;
+    PUSHORT ShortPtr;
+    PULONG LongPtr;
+
+    for (i = 0; i < Count; i++)
+    {
+        Offset = *TypeOffset & 0xFFF;
+        Type = *TypeOffset >> 12;
+        ShortPtr = (PUSHORT)(RVA(Address, Offset));
+
+        /*
+        * Don't relocate within the relocation section itself.
+        * GCC/LD generates sometimes relocation records for the relocation section.
+        * This is a bug in GCC/LD.
+        * Fix for it disabled, since it was only in ntoskrnl and not in ntdll
+        */
+        /*
+        if ((ULONG_PTR)ShortPtr < (ULONG_PTR)RelocationDir ||
+        (ULONG_PTR)ShortPtr >= (ULONG_PTR)RelocationEnd)
+        {*/
+        switch (Type)
+        {
+        /* case IMAGE_REL_BASED_SECTION : */
+        /* case IMAGE_REL_BASED_REL32 : */
+        case IMAGE_REL_BASED_ABSOLUTE:
+            break;
+
+        case IMAGE_REL_BASED_HIGH:
+            *ShortPtr += HIWORD(Delta);
+            break;
+
+        case IMAGE_REL_BASED_LOW:
+            *ShortPtr += LOWORD(Delta);
+            break;
+
+        case IMAGE_REL_BASED_HIGHLOW:
+            LongPtr = (PULONG)RVA(Address, Offset);
+            *LongPtr += Delta;
+            break;
+
+        case IMAGE_REL_BASED_HIGHADJ:
+        case IMAGE_REL_BASED_MIPS_JMPADDR:
+        default:
+            DPRINT1("Unknown/unsupported fixup type %hu.\n", Type);
+            DPRINT1("Address %x, Current %d, Count %d, *TypeOffset %x\n", Address, i, Count, *TypeOffset);
+            return (PIMAGE_BASE_RELOCATION)NULL;
+        }
+
+        TypeOffset++;
+    }
+
+    return (PIMAGE_BASE_RELOCATION)TypeOffset;
+}
+
 ULONG
 NTAPI
 LdrRelocateImageWithBias(
@@ -163,14 +229,10 @@ LdrRelocateImageWithBias(
     PIMAGE_NT_HEADERS NtHeaders;
     PIMAGE_DATA_DIRECTORY RelocationDDir;
     PIMAGE_BASE_RELOCATION RelocationDir, RelocationEnd;
-    ULONG Count, i;
+    ULONG Count;
-    PVOID Address;//, MaxAddress;
+    ULONG_PTR Address;
     PUSHORT TypeOffset;
-    ULONG_PTR Delta;
+    LONGLONG Delta;
-    SHORT Offset;
-    USHORT Type;
-    PUSHORT ShortPtr;
-    PULONG LongPtr;
 
     NtHeaders = RtlImageNtHeader(BaseAddress);
 
@@ -192,64 +254,24 @@ LdrRelocateImageWithBias(
     Delta = (ULONG_PTR)BaseAddress - NtHeaders->OptionalHeader.ImageBase + AdditionalBias;
     RelocationDir = (PIMAGE_BASE_RELOCATION)((ULONG_PTR)BaseAddress + RelocationDDir->VirtualAddress);
     RelocationEnd = (PIMAGE_BASE_RELOCATION)((ULONG_PTR)RelocationDir + RelocationDDir->Size);
-    //MaxAddress = RVA(BaseAddress, DriverSize);
 
     while (RelocationDir < RelocationEnd &&
         RelocationDir->SizeOfBlock > 0)
     {
         Count = (RelocationDir->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(USHORT);
-        Address = RVA(BaseAddress, RelocationDir->VirtualAddress);
+        Address = (ULONG_PTR)RVA(BaseAddress, RelocationDir->VirtualAddress);
         TypeOffset = (PUSHORT)(RelocationDir + 1);
 
-        for (i = 0; i < Count; i++)
+        RelocationDir = LdrProcessRelocationBlockLongLong(Address,
+                                                          Count,
+                                                          TypeOffset,
+                                                          Delta);
+
+        if (RelocationDir == NULL)
         {
-            Offset = *TypeOffset & 0xFFF;
+            DPRINT1("Error during call to LdrProcessRelocationBlockLongLong()!\n");
-            Type = *TypeOffset >> 12;
+            return Invalid;
-            ShortPtr = (PUSHORT)(RVA(Address, Offset));
-
-            /* Don't relocate after the end of the loaded driver */
-            /*if ((PVOID)ShortPtr >= MaxAddress)
-            {
-                break;
-            }*/
-
-            /*
-            * Don't relocate within the relocation section itself.
-            * GCC/LD generates sometimes relocation records for the relocation section.
-            * This is a bug in GCC/LD.
-            */
-            if ((ULONG_PTR)ShortPtr < (ULONG_PTR)RelocationDir ||
-                (ULONG_PTR)ShortPtr >= (ULONG_PTR)RelocationEnd)
-            {
-                switch (Type)
-                {
-                case IMAGE_REL_BASED_ABSOLUTE:
-                    break;
-
-                case IMAGE_REL_BASED_HIGH:
-                    *ShortPtr += HIWORD(Delta);
-                    break;
-
-                case IMAGE_REL_BASED_LOW:
-                    *ShortPtr += LOWORD(Delta);
-                    break;
-
-                case IMAGE_REL_BASED_HIGHLOW:
-                    LongPtr = (PULONG)ShortPtr;
-                    *LongPtr += Delta;
-                    break;
-
-                case IMAGE_REL_BASED_HIGHADJ:
-                case IMAGE_REL_BASED_MIPS_JMPADDR:
-                default:
-                    DPRINT1("Unknown/unsupported fixup type %hu.\n", Type);
-                    DPRINT1("Address %x, Current %d, Count %d, *TypeOffset %x\n", Address, i, Count, *TypeOffset);
-                    return Invalid;
-                }
-            }
-            TypeOffset++;
         }
-        RelocationDir = (PIMAGE_BASE_RELOCATION)((ULONG_PTR)RelocationDir + RelocationDir->SizeOfBlock);
     }
 
     return Success;