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 27c646e486fab1b86699270498f41ec2e0fe001b..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
     // 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
     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
     MemoryImagePointerAligned = (MemoryImagePointer + 0x0FFF) & (-1 << 12);\r
     \r
 \r


@@ -876,13 +876,13 @@ Returns:
     //\r
     // Allocate a buffer for the image to be loaded into.\r
     //\r
     //\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
     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
     MemoryImagePointerAligned = (MemoryImagePointer + 0x0FFF) & (-1 << 12);\r
     \r
 \r