Fix buffer overflow when the raw size of a section and the virtual size of a section...

[mirror_edk2.git] / Tools / Source / TianoTools / PeiRebase / PeiRebaseExe.c

diff --git a/Tools/Source/TianoTools/PeiRebase/PeiRebaseExe.c b/Tools/Source/TianoTools/PeiRebase/PeiRebaseExe.c
index cc07fa06bb98aa2b2f0fc217af961aea1f6d9a2a..35329c74d69cbf08c53b5118f3dc341244fee4f2 100644 (file)
--- a/Tools/Source/TianoTools/PeiRebase/PeiRebaseExe.c
+++ b/Tools/Source/TianoTools/PeiRebase/PeiRebaseExe.c
@@ -624,12 +624,12 @@ Returns:
     // Allocate a buffer for the image to be loaded into.\r
     //\r
     Pe32ImageSize       = GetLength (CurrentPe32Section.Pe32Section->CommonHeader.Size) - sizeof (EFI_PE32_SECTION);\r
-    MemoryImagePointer  = (UINTN) (malloc (Pe32ImageSize + 0x1000));\r
+    MemoryImagePointer  = (UINTN) (malloc (Pe32ImageSize + 0x1000 + ImageContext.SectionAlignment));\r
     if (MemoryImagePointer == 0) {\r
       Error (NULL, 0, 0, "memory allocation failure", NULL);\r
       return EFI_OUT_OF_RESOURCES;\r
     }\r
-    memset ((void *) MemoryImagePointer, 0, Pe32ImageSize + 0x1000);\r
+    memset ((void *) MemoryImagePointer, 0, Pe32ImageSize + 0x1000 + ImageContext.SectionAlignment);\r
     MemoryImagePointerAligned = (MemoryImagePointer + 0x0FFF) & (-1 << 12);\r
     \r
 \r
@@ -675,14 +675,14 @@ Returns:
     //\r
     PeHdr = (VOID *) ((UINTN) ImageAddress + ImageContext.PeCoffHeaderOffset);\r
     if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_IA32) {\r
-      PeHdrSizeOfImage  = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER32 *) &PeHdr->OptionalHeader).SizeOfImage);\r
-      PeHdrChecksum     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER32 *) &PeHdr->OptionalHeader).CheckSum);\r
+      PeHdrSizeOfImage     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER32 *) &PeHdr->OptionalHeader).SizeOfImage);\r
+      PeHdrChecksum        = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER32 *) &PeHdr->OptionalHeader).CheckSum);\r
     } else if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_IA64) {\r
-      PeHdrSizeOfImage  = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).SizeOfImage);\r
-      PeHdrChecksum     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).CheckSum);\r
+      PeHdrSizeOfImage     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).SizeOfImage);\r
+      PeHdrChecksum        = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).CheckSum);\r
     } else if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_X64) {\r
-      PeHdrSizeOfImage  = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).SizeOfImage);\r
-      PeHdrChecksum     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).CheckSum);\r
+      PeHdrSizeOfImage     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).SizeOfImage);\r
+      PeHdrChecksum        = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).CheckSum);\r
     } else {\r
       Error (\r
         NULL,\r
@@ -705,6 +705,24 @@ Returns:
     }\r
 \r
     memcpy (CurrentPe32Section.Pe32Section + 1, (VOID *) MemoryImagePointerAligned, (UINT32) ImageSize);\r
+    \r
+    //\r
+    // Get EntryPoint in Flash Region.\r
+    //\r
+    EntryPoint = NewPe32BaseAddress + EntryPoint - ImageAddress;\r
+\r
+    //\r
+    // If a map file was selected output mapping information for any file that\r
+    // was rebased.\r
+    //\r
+    if (MapFile != NULL) {\r
+      fprintf (MapFile, "PE32 File: %s Base:%08lx", FileGuidString, BaseAddress);\r
+      fprintf (MapFile, " EntryPoint:%08lx", EntryPoint);\r
+      if (ImageContext.PdbPointer != NULL) {\r
+        fprintf (MapFile, " FileName: %s", ImageContext.PdbPointer);\r
+      }\r
+      fprintf (MapFile, "\n");\r
+    }\r
 \r
     free ((VOID *) MemoryImagePointer);\r
 \r
@@ -798,6 +816,8 @@ Returns:
     //\r
     PeHdr->FileHeader.SizeOfOptionalHeader = (UINT16) (TEImageHeader->StrippedSize - 0x40 - sizeof (UINT32) - sizeof (EFI_IMAGE_FILE_HEADER));\r
     PeHdr->OptionalHeader.ImageBase = (UINTN) (TEImageHeader->ImageBase - TEImageHeader->StrippedSize + sizeof (EFI_TE_IMAGE_HEADER));\r
+    PeHdr->OptionalHeader.AddressOfEntryPoint = TEImageHeader->AddressOfEntryPoint;\r
+    PeHdr->OptionalHeader.BaseOfCode  = TEImageHeader->BaseOfCode;\r
     PeHdr->OptionalHeader.SizeOfImage = Pe32ImageSize;\r
     PeHdr->OptionalHeader.Subsystem   = TEImageHeader->Subsystem;\r
     PeHdr->OptionalHeader.SizeOfImage = Pe32ImageSize;\r
@@ -856,13 +876,13 @@ Returns:
     //\r
     // Allocate a buffer for the image to be loaded into.\r
     //\r
-    MemoryImagePointer = (UINTN) (malloc (Pe32ImageSize + 0x1000));\r
+    MemoryImagePointer = (UINTN) (malloc (Pe32ImageSize + 0x1000 + ImageContext.SectionAlignment));\r
     if (MemoryImagePointer == 0) {\r
       Error (NULL, 0, 0, "memory allocation error on rebase of TE image", FileGuidString);\r
       free (TEBuffer);\r
       return EFI_OUT_OF_RESOURCES;\r
     }\r
-    memset ((void *) MemoryImagePointer, 0, Pe32ImageSize + 0x1000);\r
+    memset ((void *) MemoryImagePointer, 0, Pe32ImageSize + 0x1000 + ImageContext.SectionAlignment);\r
     MemoryImagePointerAligned = (MemoryImagePointer + 0x0FFF) & (-1 << 12);\r
     \r
 \r
@@ -895,11 +915,11 @@ Returns:
     //\r
     PeHdr = (VOID *) ((UINTN) ImageAddress + ImageContext.PeCoffHeaderOffset);\r
     if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_IA32) {\r
-      PeHdrSizeOfImage  = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER32 *) &PeHdr->OptionalHeader).SizeOfImage);\r
-      PeHdrChecksum     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER32 *) &PeHdr->OptionalHeader).CheckSum);\r
+      PeHdrSizeOfImage     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER32 *) &PeHdr->OptionalHeader).SizeOfImage);\r
+      PeHdrChecksum        = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER32 *) &PeHdr->OptionalHeader).CheckSum);\r
     } else if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_IA64) {\r
-      PeHdrSizeOfImage  = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).SizeOfImage);\r
-      PeHdrChecksum     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).CheckSum);\r
+      PeHdrSizeOfImage     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).SizeOfImage);\r
+      PeHdrChecksum        = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).CheckSum);\r
     } else {\r
       Error (\r
         NULL,\r
@@ -929,6 +949,25 @@ Returns:
       GetLength (CurrentPe32Section.Pe32Section->CommonHeader.Size) - sizeof (EFI_PE32_SECTION) -\r
       sizeof (EFI_TE_IMAGE_HEADER)\r
       );\r
+    \r
+    //\r
+    // Get EntryPoint in Flash Region.\r
+    //\r
+    EntryPoint = NewPe32BaseAddress + EntryPoint - ImageAddress;\r
+\r
+    //\r
+    // If a map file was selected output mapping information for any file that\r
+    // was rebased.\r
+    //\r
+    if (MapFile != NULL) {\r
+      fprintf (MapFile, "TE   File: %s Base:%08lx", FileGuidString, BaseAddress);\r
+      fprintf (MapFile, " EntryPoint:%08lx", EntryPoint);\r
+      if (ImageContext.PdbPointer != NULL) {\r
+        fprintf (MapFile, " FileName: %s", ImageContext.PdbPointer);\r
+      }\r
+      fprintf (MapFile, "\n");\r
+    }\r
+\r
     free ((VOID *) MemoryImagePointer);\r
     free (TEBuffer);\r
     if (FfsFile->Attributes & FFS_ATTRIB_TAIL_PRESENT) {\r
@@ -972,18 +1011,6 @@ Returns:
       return EFI_NOT_FOUND;\r
     }\r
   }\r
-\r
-  //\r
-  // If a map file was selected output mapping information for any file that\r
-  // was rebased.\r
-  //\r
-  if (MapFile != NULL) {\r
-    fprintf (MapFile, "File: %s Base:%08lx", FileGuidString, BaseAddress);\r
-    if (ImageContext.PdbPointer != NULL) {\r
-      fprintf (MapFile, " FileName: %s", ImageContext.PdbPointer);\r
-    }\r
-    fprintf (MapFile, "\n");\r
-  }\r
   \r
   return EFI_SUCCESS;\r
 }\r