From 966064c9494e416f5503fc0cc9fb2c2ba6bc2097 Mon Sep 17 00:00:00 2001
From: David Majnemer <david.majnemer@gmail.com>
Date: Fri, 14 Nov 2014 19:35:59 +0000
Subject: [PATCH] yaml2obj, COFF: Correctly calculate SizeOfImage and
 SizeOfHeaders

SizeOfHeaders must be aligned to the FileAlignment.
SizeOfImage must be at least the SizeOfHeaders aligned to the
SectionAlignment.

llvm-svn: 222030
---
 llvm/tools/yaml2obj/yaml2coff.cpp | 25 ++++++++++++++++---------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/llvm/tools/yaml2obj/yaml2coff.cpp b/llvm/tools/yaml2obj/yaml2coff.cpp
index 2c71988..8ed241f9 100644
--- a/llvm/tools/yaml2obj/yaml2coff.cpp
+++ b/llvm/tools/yaml2obj/yaml2coff.cpp
@@ -144,10 +144,11 @@ struct COFFParser {
 static bool layoutOptionalHeader(COFFParser &CP) {
   if (!CP.isPE())
     return true;
+  unsigned PEHeaderSize = CP.is64Bit() ? sizeof(object::pe32plus_header)
+                                       : sizeof(object::pe32_header);
   CP.Obj.Header.SizeOfOptionalHeader =
-      (CP.is64Bit() ? sizeof(object::pe32plus_header)
-                    : sizeof(object::pe32_header)) +
-      (sizeof(object::data_directory) * (COFF::NUM_DATA_DIRECTORIES + 1));
+      PEHeaderSize +
+      sizeof(object::data_directory) * (COFF::NUM_DATA_DIRECTORIES + 1);
   return true;
 }
 
@@ -281,15 +282,18 @@ num_zeros_impl num_zeros(size_t N) {
 }
 
 template <typename T>
-static void initializeOptionalHeader(COFFParser &CP, uint16_t Magic, T Header) {
+static uint32_t initializeOptionalHeader(COFFParser &CP, uint16_t Magic, T Header) {
   memset(Header, 0, sizeof(*Header));
   Header->Magic = Magic;
   Header->SectionAlignment = CP.Obj.OptionalHeader->Header.SectionAlignment;
+  Header->FileAlignment = CP.Obj.OptionalHeader->Header.FileAlignment;
   uint32_t SizeOfCode = 0, SizeOfInitializedData = 0,
            SizeOfUninitializedData = 0;
   uint32_t SizeOfHeaders = RoundUpToAlignment(
-      CP.SectionTableStart + CP.SectionTableSize, Header->SectionAlignment);
-  uint32_t SizeOfImage = SizeOfHeaders;
+      CP.SectionTableStart + CP.SectionTableSize, Header->FileAlignment);
+  uint32_t SizeOfImage =
+      RoundUpToAlignment(SizeOfHeaders, Header->SectionAlignment);
+  uint32_t BaseOfData = 0;
   for (const COFFYAML::Section &S : CP.Obj.Sections) {
     if (S.Header.Characteristics & COFF::IMAGE_SCN_CNT_CODE)
       SizeOfCode += S.Header.SizeOfRawData;
@@ -298,7 +302,9 @@ static void initializeOptionalHeader(COFFParser &CP, uint16_t Magic, T Header) {
     if (S.Header.Characteristics & COFF::IMAGE_SCN_CNT_UNINITIALIZED_DATA)
       SizeOfUninitializedData += S.Header.SizeOfRawData;
     if (S.Name.equals(".text"))
-      Header->BaseOfCode = S.Header.VirtualAddress;          // RVA
+      Header->BaseOfCode = S.Header.VirtualAddress; // RVA
+    else if (S.Name.equals(".data"))
+      BaseOfData = S.Header.VirtualAddress; // RVA
     if (S.Header.VirtualAddress)
       SizeOfImage +=
           RoundUpToAlignment(S.Header.VirtualSize, Header->SectionAlignment);
@@ -309,7 +315,6 @@ static void initializeOptionalHeader(COFFParser &CP, uint16_t Magic, T Header) {
   Header->AddressOfEntryPoint =
       CP.Obj.OptionalHeader->Header.AddressOfEntryPoint; // RVA
   Header->ImageBase = CP.Obj.OptionalHeader->Header.ImageBase;
-  Header->FileAlignment = CP.Obj.OptionalHeader->Header.FileAlignment;
   Header->MajorOperatingSystemVersion =
       CP.Obj.OptionalHeader->Header.MajorOperatingSystemVersion;
   Header->MinorOperatingSystemVersion =
@@ -331,6 +336,7 @@ static void initializeOptionalHeader(COFFParser &CP, uint16_t Magic, T Header) {
   Header->SizeOfHeapReserve = CP.Obj.OptionalHeader->Header.SizeOfHeapReserve;
   Header->SizeOfHeapCommit = CP.Obj.OptionalHeader->Header.SizeOfHeapCommit;
   Header->NumberOfRvaAndSize = COFF::NUM_DATA_DIRECTORIES + 1;
+  return BaseOfData;
 }
 
 static bool writeCOFF(COFFParser &CP, raw_ostream &OS) {
@@ -387,7 +393,8 @@ static bool writeCOFF(COFFParser &CP, raw_ostream &OS) {
       OS.write(reinterpret_cast<char *>(&PEH), sizeof(PEH));
     } else {
       object::pe32_header PEH;
-      initializeOptionalHeader(CP, COFF::PE32Header::PE32, &PEH);
+      uint32_t BaseOfData = initializeOptionalHeader(CP, COFF::PE32Header::PE32, &PEH);
+      PEH.BaseOfData = BaseOfData;
       OS.write(reinterpret_cast<char *>(&PEH), sizeof(PEH));
     }
     for (const Optional<COFF::DataDirectory> &DD :
-- 
2.7.4