obj2yaml, yaml2obj: Add support for COFF executables

In support of serializing executables, obj2yaml now records the virtual address
and size of sections.  It also serializes whatever we strictly need from
the PE header, it expects that it can reconstitute everything else via
inference.

yaml2obj can reconstitute a fully linked executable.

In order to get executables correctly serialized/deserialized, other
bugs were fixed as a circumstance.  We now properly respect file and
section alignments.  We also avoid writing out string tables unless they
are strictly necessary.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@221975 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/Object/COFF.h

@@ -110,12 +110,14 @@ struct pe32_header {
   support::ulittle32_t SizeOfHeaders;
   support::ulittle32_t CheckSum;
   support::ulittle16_t Subsystem;
+  // FIXME: This should be DllCharacteristics.
   support::ulittle16_t DLLCharacteristics;
   support::ulittle32_t SizeOfStackReserve;
   support::ulittle32_t SizeOfStackCommit;
   support::ulittle32_t SizeOfHeapReserve;
   support::ulittle32_t SizeOfHeapCommit;
   support::ulittle32_t LoaderFlags;
+  // FIXME: This should be NumberOfRvaAndSizes.
   support::ulittle32_t NumberOfRvaAndSize;
 };
 

include/llvm/Object/COFFYAML.h

@@ -31,6 +31,12 @@ inline SectionCharacteristics operator|(SectionCharacteristics a,
   uint32_t Ret = static_cast<uint32_t>(a) | static_cast<uint32_t>(b);
   return static_cast<SectionCharacteristics>(Ret);
 }
+
+inline DLLCharacteristics operator|(DLLCharacteristics a,
+                                    DLLCharacteristics b) {
+  uint16_t Ret = static_cast<uint16_t>(a) | static_cast<uint16_t>(b);
+  return static_cast<DLLCharacteristics>(Ret);
+}
 }
 
 // The structure of the yaml files is not an exact 1:1 match to COFF. In order
@@ -69,7 +75,13 @@ namespace COFFYAML {
     Symbol();
   };
 
+  struct PEHeader {
+    COFF::PE32Header Header;
+    Optional<COFF::DataDirectory> DataDirectories[COFF::NUM_DATA_DIRECTORIES];
+  };
+
   struct Object {
+    Optional<PEHeader> OptionalHeader;
     COFF::header Header;
     std::vector<Section> Sections;
     std::vector<Symbol> Symbols;
@@ -130,6 +142,11 @@ struct ScalarEnumerationTraits<COFF::RelocationTypeAMD64> {
   static void enumeration(IO &IO, COFF::RelocationTypeAMD64 &Value);
 };
 
+template <>
+struct ScalarEnumerationTraits<COFF::WindowsSubsystem> {
+  static void enumeration(IO &IO, COFF::WindowsSubsystem &Value);
+};
+
 template <>
 struct ScalarBitSetTraits<COFF::Characteristics> {
   static void bitset(IO &IO, COFF::Characteristics &Value);
@@ -140,11 +157,26 @@ struct ScalarBitSetTraits<COFF::SectionCharacteristics> {
   static void bitset(IO &IO, COFF::SectionCharacteristics &Value);
 };
 
+template <>
+struct ScalarBitSetTraits<COFF::DLLCharacteristics> {
+  static void bitset(IO &IO, COFF::DLLCharacteristics &Value);
+};
+
 template <>
 struct MappingTraits<COFFYAML::Relocation> {
   static void mapping(IO &IO, COFFYAML::Relocation &Rel);
 };
 
+template <>
+struct MappingTraits<COFFYAML::PEHeader> {
+  static void mapping(IO &IO, COFFYAML::PEHeader &PH);
+};
+
+template <>
+struct MappingTraits<COFF::DataDirectory> {
+  static void mapping(IO &IO, COFF::DataDirectory &DD);
+};
+
 template <>
 struct MappingTraits<COFF::header> {
   static void mapping(IO &IO, COFF::header &H);

include/llvm/Support/COFF.h

@@ -515,12 +515,14 @@ namespace COFF {
     uint32_t SizeOfHeaders;
     uint32_t CheckSum;
     uint16_t Subsystem;
+    // FIXME: This should be DllCharacteristics to match the COFF spec.
     uint16_t DLLCharacteristics;
     uint32_t SizeOfStackReserve;
     uint32_t SizeOfStackCommit;
     uint32_t SizeOfHeapReserve;
     uint32_t SizeOfHeapCommit;
     uint32_t LoaderFlags;
+    // FIXME: This should be NumberOfRvaAndSizes to match the COFF spec.
     uint32_t NumberOfRvaAndSize;
   };
 
@@ -544,10 +546,12 @@ namespace COFF {
     BOUND_IMPORT,
     IAT,
     DELAY_IMPORT_DESCRIPTOR,
-    CLR_RUNTIME_HEADER
+    CLR_RUNTIME_HEADER,
+
+    NUM_DATA_DIRECTORIES
   };
 
-  enum WindowsSubsystem {
+  enum WindowsSubsystem : uint16_t {
     IMAGE_SUBSYSTEM_UNKNOWN = 0, ///< An unknown subsystem.
     IMAGE_SUBSYSTEM_NATIVE = 1, ///< Device drivers and native Windows processes
     IMAGE_SUBSYSTEM_WINDOWS_GUI = 2, ///< The Windows GUI subsystem.
@@ -566,7 +570,7 @@ namespace COFF {
     IMAGE_SUBSYSTEM_WINDOWS_BOOT_APPLICATION = 16 ///< A BCD application.
   };
 
-  enum DLLCharacteristics {
+  enum DLLCharacteristics : uint16_t {
     /// ASLR with 64 bit address space.
     IMAGE_DLL_CHARACTERISTICS_HIGH_ENTROPY_VA = 0x0020,
     /// DLL can be relocated at load time.

lib/Object/COFFObjectFile.cpp

@@ -460,9 +460,9 @@ std::error_code COFFObjectFile::initSymbolTablePtr() {
   // Find string table. The first four byte of the string table contains the
   // total size of the string table, including the size field itself. If the
   // string table is empty, the value of the first four byte would be 4.
-  const uint8_t *StringTableAddr =
+  uint32_t StringTableOffset = getPointerToSymbolTable() +
-      base() + getPointerToSymbolTable() +
+                               getNumberOfSymbols() * getSymbolTableEntrySize();
-      getNumberOfSymbols() * getSymbolTableEntrySize();
+  const uint8_t *StringTableAddr = base() + StringTableOffset;
   const ulittle32_t *StringTableSizePtr;
   if (std::error_code EC = getObject(StringTableSizePtr, Data, StringTableAddr))
     return EC;
@@ -826,13 +826,17 @@ std::error_code
 COFFObjectFile::getDataDirectory(uint32_t Index,
                                  const data_directory *&Res) const {
   // Error if if there's no data directory or the index is out of range.
-  if (!DataDirectory)
+  if (!DataDirectory) {
+    Res = nullptr;
     return object_error::parse_failed;
+  }
   assert(PE32Header || PE32PlusHeader);
   uint32_t NumEnt = PE32Header ? PE32Header->NumberOfRvaAndSize
                                : PE32PlusHeader->NumberOfRvaAndSize;
-  if (Index > NumEnt)
+  if (Index >= NumEnt) {
+    Res = nullptr;
     return object_error::parse_failed;
+  }
   Res = &DataDirectory[Index];
   return object_error::success;
 }

lib/Object/COFFYAML.cpp

@@ -168,6 +168,24 @@ void ScalarEnumerationTraits<COFF::RelocationTypeAMD64>::enumeration(
   ECase(IMAGE_REL_AMD64_PAIR);
   ECase(IMAGE_REL_AMD64_SSPAN32);
 }
+
+void ScalarEnumerationTraits<COFF::WindowsSubsystem>::enumeration(
+    IO &IO, COFF::WindowsSubsystem &Value) {
+    ECase(IMAGE_SUBSYSTEM_UNKNOWN);
+    ECase(IMAGE_SUBSYSTEM_NATIVE);
+    ECase(IMAGE_SUBSYSTEM_WINDOWS_GUI);
+    ECase(IMAGE_SUBSYSTEM_WINDOWS_CUI);
+    ECase(IMAGE_SUBSYSTEM_OS2_CUI);
+    ECase(IMAGE_SUBSYSTEM_POSIX_CUI);
+    ECase(IMAGE_SUBSYSTEM_NATIVE_WINDOWS);
+    ECase(IMAGE_SUBSYSTEM_WINDOWS_CE_GUI);
+    ECase(IMAGE_SUBSYSTEM_EFI_APPLICATION);
+    ECase(IMAGE_SUBSYSTEM_EFI_BOOT_SERVICE_DRIVER);
+    ECase(IMAGE_SUBSYSTEM_EFI_RUNTIME_DRIVER);
+    ECase(IMAGE_SUBSYSTEM_EFI_ROM);
+    ECase(IMAGE_SUBSYSTEM_XBOX);
+    ECase(IMAGE_SUBSYSTEM_WINDOWS_BOOT_APPLICATION);
+}
 #undef ECase
 
 #define BCase(X) IO.bitSetCase(Value, #X, COFF::X);
@@ -214,6 +232,21 @@ void ScalarBitSetTraits<COFF::SectionCharacteristics>::bitset(
   BCase(IMAGE_SCN_MEM_READ);
   BCase(IMAGE_SCN_MEM_WRITE);
 }
+
+void ScalarBitSetTraits<COFF::DLLCharacteristics>::bitset(
+    IO &IO, COFF::DLLCharacteristics &Value) {
+  BCase(IMAGE_DLL_CHARACTERISTICS_HIGH_ENTROPY_VA);
+  BCase(IMAGE_DLL_CHARACTERISTICS_DYNAMIC_BASE);
+  BCase(IMAGE_DLL_CHARACTERISTICS_FORCE_INTEGRITY);
+  BCase(IMAGE_DLL_CHARACTERISTICS_NX_COMPAT);
+  BCase(IMAGE_DLL_CHARACTERISTICS_NO_ISOLATION);
+  BCase(IMAGE_DLL_CHARACTERISTICS_NO_SEH);
+  BCase(IMAGE_DLL_CHARACTERISTICS_NO_BIND);
+  BCase(IMAGE_DLL_CHARACTERISTICS_APPCONTAINER);
+  BCase(IMAGE_DLL_CHARACTERISTICS_WDM_DRIVER);
+  BCase(IMAGE_DLL_CHARACTERISTICS_GUARD_CF);
+  BCase(IMAGE_DLL_CHARACTERISTICS_TERMINAL_SERVER_AWARE);
+}
 #undef BCase
 
 namespace {
@@ -285,6 +318,23 @@ struct NType {
   RelocType Type;
 };
 
+struct NWindowsSubsystem {
+  NWindowsSubsystem(IO &) : Subsystem(COFF::WindowsSubsystem(0)) {}
+  NWindowsSubsystem(IO &, uint16_t C) : Subsystem(COFF::WindowsSubsystem(C)) {}
+  uint16_t denormalize(IO &) { return Subsystem; }
+
+  COFF::WindowsSubsystem Subsystem;
+};
+
+struct NDLLCharacteristics {
+  NDLLCharacteristics(IO &) : Characteristics(COFF::DLLCharacteristics(0)) {}
+  NDLLCharacteristics(IO &, uint16_t C)
+      : Characteristics(COFF::DLLCharacteristics(C)) {}
+  uint16_t denormalize(IO &) { return Characteristics; }
+
+  COFF::DLLCharacteristics Characteristics;
+};
+
 }
 
 void MappingTraits<COFFYAML::Relocation>::mapping(IO &IO,
@@ -306,6 +356,59 @@ void MappingTraits<COFFYAML::Relocation>::mapping(IO &IO,
   }
 }
 
+void MappingTraits<COFF::DataDirectory>::mapping(IO &IO,
+                                                 COFF::DataDirectory &DD) {
+  IO.mapRequired("RelativeVirtualAddress", DD.RelativeVirtualAddress);
+  IO.mapRequired("Size", DD.Size);
+}
+
+void MappingTraits<COFFYAML::PEHeader>::mapping(IO &IO,
+                                                COFFYAML::PEHeader &PH) {
+  MappingNormalization<NWindowsSubsystem, uint16_t> NWS(IO,
+                                                        PH.Header.Subsystem);
+  MappingNormalization<NDLLCharacteristics, uint16_t> NDC(
+      IO, PH.Header.DLLCharacteristics);
+
+  IO.mapRequired("AddressOfEntryPoint", PH.Header.AddressOfEntryPoint);
+  IO.mapRequired("ImageBase", PH.Header.ImageBase);
+  IO.mapRequired("SectionAlignment", PH.Header.SectionAlignment);
+  IO.mapRequired("FileAlignment", PH.Header.FileAlignment);
+  IO.mapRequired("MajorOperatingSystemVersion",
+                 PH.Header.MajorOperatingSystemVersion);
+  IO.mapRequired("MinorOperatingSystemVersion",
+                 PH.Header.MinorOperatingSystemVersion);
+  IO.mapRequired("MajorImageVersion", PH.Header.MajorImageVersion);
+  IO.mapRequired("MinorImageVersion", PH.Header.MinorImageVersion);
+  IO.mapRequired("MajorSubsystemVersion", PH.Header.MajorSubsystemVersion);
+  IO.mapRequired("MinorSubsystemVersion", PH.Header.MinorSubsystemVersion);
+  IO.mapRequired("Subsystem", NWS->Subsystem);
+  IO.mapRequired("DLLCharacteristics", NDC->Characteristics);
+  IO.mapRequired("SizeOfStackReserve", PH.Header.SizeOfStackReserve);
+  IO.mapRequired("SizeOfStackCommit", PH.Header.SizeOfStackCommit);
+  IO.mapRequired("SizeOfHeapReserve", PH.Header.SizeOfHeapReserve);
+  IO.mapRequired("SizeOfHeapCommit", PH.Header.SizeOfHeapCommit);
+
+  IO.mapOptional("ExportTable", PH.DataDirectories[COFF::EXPORT_TABLE]);
+  IO.mapOptional("ImportTable", PH.DataDirectories[COFF::IMPORT_TABLE]);
+  IO.mapOptional("ResourceTable", PH.DataDirectories[COFF::RESOURCE_TABLE]);
+  IO.mapOptional("ExceptionTable", PH.DataDirectories[COFF::EXCEPTION_TABLE]);
+  IO.mapOptional("CertificateTable", PH.DataDirectories[COFF::CERTIFICATE_TABLE]);
+  IO.mapOptional("BaseRelocationTable",
+                 PH.DataDirectories[COFF::BASE_RELOCATION_TABLE]);
+  IO.mapOptional("Debug", PH.DataDirectories[COFF::DEBUG]);
+  IO.mapOptional("Architecture", PH.DataDirectories[COFF::ARCHITECTURE]);
+  IO.mapOptional("GlobalPtr", PH.DataDirectories[COFF::GLOBAL_PTR]);
+  IO.mapOptional("TlsTable", PH.DataDirectories[COFF::TLS_TABLE]);
+  IO.mapOptional("LoadConfigTable",
+                 PH.DataDirectories[COFF::LOAD_CONFIG_TABLE]);
+  IO.mapOptional("BoundImport", PH.DataDirectories[COFF::BOUND_IMPORT]);
+  IO.mapOptional("IAT", PH.DataDirectories[COFF::IAT]);
+  IO.mapOptional("DelayImportDescriptor",
+                 PH.DataDirectories[COFF::DELAY_IMPORT_DESCRIPTOR]);
+  IO.mapOptional("ClrRuntimeHeader",
+                 PH.DataDirectories[COFF::CLR_RUNTIME_HEADER]);
+}
+
 void MappingTraits<COFF::header>::mapping(IO &IO, COFF::header &H) {
   MappingNormalization<NMachine, uint16_t> NM(IO, H.Machine);
   MappingNormalization<NHeaderCharacteristics, uint16_t> NC(IO,
@@ -380,12 +483,15 @@ void MappingTraits<COFFYAML::Section>::mapping(IO &IO, COFFYAML::Section &Sec) {
       IO, Sec.Header.Characteristics);
   IO.mapRequired("Name", Sec.Name);
   IO.mapRequired("Characteristics", NC->Characteristics);
+  IO.mapOptional("VirtualAddress", Sec.Header.VirtualAddress, 0U);
+  IO.mapOptional("VirtualSize", Sec.Header.VirtualSize, 0U);
   IO.mapOptional("Alignment", Sec.Alignment);
   IO.mapRequired("SectionData", Sec.SectionData);
   IO.mapOptional("Relocations", Sec.Relocations);
 }
 
 void MappingTraits<COFFYAML::Object>::mapping(IO &IO, COFFYAML::Object &Obj) {
+  IO.mapOptional("OptionalHeader", Obj.OptionalHeader);
   IO.mapRequired("header", Obj.Header);
   IO.mapRequired("sections", Obj.Sections);
   IO.mapRequired("symbols", Obj.Symbols);

tools/obj2yaml/coff2yaml.cpp

@@ -20,6 +20,8 @@ namespace {
 class COFFDumper {
   const object::COFFObjectFile &Obj;
   COFFYAML::Object YAMLObj;
+  template <typename T>
+  void dumpOptionalHeader(T OptionalHeader);
   void dumpHeader();
   void dumpSections(unsigned numSections);
   void dumpSymbols(unsigned numSymbols);
@@ -32,32 +34,89 @@ public:
 }
 
 COFFDumper::COFFDumper(const object::COFFObjectFile &Obj) : Obj(Obj) {
+  const object::pe32_header *PE32Header = nullptr;
+  Obj.getPE32Header(PE32Header);
+  if (PE32Header) {
+    dumpOptionalHeader(PE32Header);
+  } else {
+    const object::pe32plus_header *PE32PlusHeader = nullptr;
+    Obj.getPE32PlusHeader(PE32PlusHeader);
+    if (PE32PlusHeader) {
+      dumpOptionalHeader(PE32PlusHeader);
+    }
+  }
   dumpHeader();
   dumpSections(Obj.getNumberOfSections());
   dumpSymbols(Obj.getNumberOfSymbols());
 }
 
+template <typename T> void COFFDumper::dumpOptionalHeader(T OptionalHeader) {
+  YAMLObj.OptionalHeader = COFFYAML::PEHeader();
+  YAMLObj.OptionalHeader->Header.AddressOfEntryPoint =
+      OptionalHeader->AddressOfEntryPoint;
+  YAMLObj.OptionalHeader->Header.AddressOfEntryPoint =
+      OptionalHeader->AddressOfEntryPoint;
+  YAMLObj.OptionalHeader->Header.ImageBase = OptionalHeader->ImageBase;
+  YAMLObj.OptionalHeader->Header.SectionAlignment =
+      OptionalHeader->SectionAlignment;
+  YAMLObj.OptionalHeader->Header.FileAlignment = OptionalHeader->FileAlignment;
+  YAMLObj.OptionalHeader->Header.MajorOperatingSystemVersion =
+      OptionalHeader->MajorOperatingSystemVersion;
+  YAMLObj.OptionalHeader->Header.MinorOperatingSystemVersion =
+      OptionalHeader->MinorOperatingSystemVersion;
+  YAMLObj.OptionalHeader->Header.MajorImageVersion =
+      OptionalHeader->MajorImageVersion;
+  YAMLObj.OptionalHeader->Header.MinorImageVersion =
+      OptionalHeader->MinorImageVersion;
+  YAMLObj.OptionalHeader->Header.MajorSubsystemVersion =
+      OptionalHeader->MajorSubsystemVersion;
+  YAMLObj.OptionalHeader->Header.MinorSubsystemVersion =
+      OptionalHeader->MinorSubsystemVersion;
+  YAMLObj.OptionalHeader->Header.Subsystem = OptionalHeader->Subsystem;
+  YAMLObj.OptionalHeader->Header.DLLCharacteristics =
+      OptionalHeader->DLLCharacteristics;
+  YAMLObj.OptionalHeader->Header.SizeOfStackReserve =
+      OptionalHeader->SizeOfStackReserve;
+  YAMLObj.OptionalHeader->Header.SizeOfStackCommit =
+      OptionalHeader->SizeOfStackCommit;
+  YAMLObj.OptionalHeader->Header.SizeOfHeapReserve =
+      OptionalHeader->SizeOfHeapReserve;
+  YAMLObj.OptionalHeader->Header.SizeOfHeapCommit =
+      OptionalHeader->SizeOfHeapCommit;
+  unsigned I = 0;
+  for (auto &DestDD : YAMLObj.OptionalHeader->DataDirectories) {
+    const object::data_directory *DD;
+    if (Obj.getDataDirectory(I++, DD))
+      continue;
+    DestDD = COFF::DataDirectory();
+    DestDD->RelativeVirtualAddress = DD->RelativeVirtualAddress;
+    DestDD->Size = DD->Size;
+  }
+}
+
 void COFFDumper::dumpHeader() {
   YAMLObj.Header.Machine = Obj.getMachine();
   YAMLObj.Header.Characteristics = Obj.getCharacteristics();
 }
 
 void COFFDumper::dumpSections(unsigned NumSections) {
-  std::vector<COFFYAML::Section> &Sections = YAMLObj.Sections;
+  std::vector<COFFYAML::Section> &YAMLSections = YAMLObj.Sections;
-  for (const auto &Section : Obj.sections()) {
+  for (const auto &ObjSection : Obj.sections()) {
-    const object::coff_section *Sect = Obj.getCOFFSection(Section);
+    const object::coff_section *COFFSection = Obj.getCOFFSection(ObjSection);
-    COFFYAML::Section Sec;
+    COFFYAML::Section NewYAMLSection;
-    Section.getName(Sec.Name);
+    ObjSection.getName(NewYAMLSection.Name);
-    Sec.Header.Characteristics = Sect->Characteristics;
+    NewYAMLSection.Header.Characteristics = COFFSection->Characteristics;
-    Sec.Alignment = Section.getAlignment();
+    NewYAMLSection.Header.VirtualAddress = ObjSection.getAddress();
+    NewYAMLSection.Header.VirtualSize = COFFSection->VirtualSize;
+    NewYAMLSection.Alignment = ObjSection.getAlignment();
 
     ArrayRef<uint8_t> sectionData;
-    if (!Section.isBSS())
+    if (!ObjSection.isBSS())
-      Obj.getSectionContents(Sect, sectionData);
+      Obj.getSectionContents(COFFSection, sectionData);
-    Sec.SectionData = yaml::BinaryRef(sectionData);
+    NewYAMLSection.SectionData = yaml::BinaryRef(sectionData);
 
     std::vector<COFFYAML::Relocation> Relocations;
-    for (const auto &Reloc : Section.relocations()) {
+    for (const auto &Reloc : ObjSection.relocations()) {
       const object::coff_relocation *reloc = Obj.getCOFFRelocation(Reloc);
       COFFYAML::Relocation Rel;
       object::symbol_iterator Sym = Reloc.getSymbol();
@@ -66,8 +125,8 @@ void COFFDumper::dumpSections(unsigned NumSections) {
       Rel.Type = reloc->Type;
       Relocations.push_back(Rel);
     }
-    Sec.Relocations = Relocations;
+    NewYAMLSection.Relocations = Relocations;
-    Sections.push_back(Sec);
+    YAMLSections.push_back(NewYAMLSection);
   }
 }
 

tools/yaml2obj/yaml2coff.cpp

@@ -19,6 +19,7 @@
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Object/COFFYAML.h"
+#include "llvm/Object/COFF.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
@@ -30,7 +31,8 @@ using namespace llvm;
 /// This parses a yaml stream that represents a COFF object file.
 /// See docs/yaml2obj for the yaml scheema.
 struct COFFParser {
-  COFFParser(COFFYAML::Object &Obj) : Obj(Obj) {
+  COFFParser(COFFYAML::Object &Obj)
+      : Obj(Obj), SectionTableStart(0), SectionTableSize(0) {
     // A COFF string table always starts with a 4 byte size field. Offsets into
     // it include this size, so allocate it now.
     StringTable.append(4, char(0));
@@ -41,6 +43,15 @@ struct COFFParser {
            COFF::MaxNumberOfSections16;
   }
 
+  bool isPE() const { return Obj.OptionalHeader.hasValue(); }
+  bool is64Bit() const {
+    return Obj.Header.Machine == COFF::IMAGE_FILE_MACHINE_AMD64;
+  }
+
+  uint32_t getFileAlignment() const {
+    return Obj.OptionalHeader->Header.FileAlignment;
+  }
+
   unsigned getHeaderSize() const {
     return useBigObj() ? COFF::Header32Size : COFF::Header16Size;
   }
@@ -124,39 +135,54 @@ struct COFFParser {
 
   StringMap<unsigned> StringTableMap;
   std::string StringTable;
+  uint32_t SectionTableStart;
+  uint32_t SectionTableSize;
 };
 
 // Take a CP and assign addresses and sizes to everything. Returns false if the
 // layout is not valid to do.
-static bool layoutCOFF(COFFParser &CP) {
+static bool layoutOptionalHeader(COFFParser &CP) {
-  uint32_t SectionTableStart = 0;
+  if (!CP.isPE())
-  uint32_t SectionTableSize  = 0;
+    return true;
+  CP.Obj.Header.SizeOfOptionalHeader =
+      (CP.is64Bit() ? sizeof(object::pe32plus_header)
+                    : sizeof(object::pe32_header)) +
+      (sizeof(object::data_directory) * (COFF::NUM_DATA_DIRECTORIES + 1));
+  return true;
+}
 
+// Take a CP and assign addresses and sizes to everything. Returns false if the
+// layout is not valid to do.
+static bool layoutCOFF(COFFParser &CP) {
   // The section table starts immediately after the header, including the
   // optional header.
-  SectionTableStart = CP.getHeaderSize() + CP.Obj.Header.SizeOfOptionalHeader;
+  CP.SectionTableStart =
-  SectionTableSize = COFF::SectionSize * CP.Obj.Sections.size();
+      CP.getHeaderSize() + CP.Obj.Header.SizeOfOptionalHeader;
+  CP.SectionTableSize = COFF::SectionSize * CP.Obj.Sections.size();
 
-  uint32_t CurrentSectionDataOffset = SectionTableStart + SectionTableSize;
+  uint32_t CurrentSectionDataOffset =
+      CP.SectionTableStart + CP.SectionTableSize;
 
   // Assign each section data address consecutively.
-  for (std::vector<COFFYAML::Section>::iterator i = CP.Obj.Sections.begin(),
+  for (COFFYAML::Section &S : CP.Obj.Sections) {
-                                                e = CP.Obj.Sections.end();
+    if (S.SectionData.binary_size() > 0) {
-                                                i != e; ++i) {
+      CurrentSectionDataOffset = RoundUpToAlignment(
-    if (i->SectionData.binary_size() > 0) {
+          CurrentSectionDataOffset, CP.isPE() ? CP.getFileAlignment() : 4);
-      i->Header.SizeOfRawData = i->SectionData.binary_size();
+      S.Header.SizeOfRawData = S.SectionData.binary_size();
-      i->Header.PointerToRawData = CurrentSectionDataOffset;
+      if (CP.isPE())
-      CurrentSectionDataOffset += i->Header.SizeOfRawData;
+        S.Header.SizeOfRawData =
-      if (!i->Relocations.empty()) {
+            RoundUpToAlignment(S.Header.SizeOfRawData, CP.getFileAlignment());
-        i->Header.PointerToRelocations = CurrentSectionDataOffset;
+      S.Header.PointerToRawData = CurrentSectionDataOffset;
-        i->Header.NumberOfRelocations = i->Relocations.size();
+      CurrentSectionDataOffset += S.Header.SizeOfRawData;
-        CurrentSectionDataOffset += i->Header.NumberOfRelocations *
+      if (!S.Relocations.empty()) {
-          COFF::RelocationSize;
+        S.Header.PointerToRelocations = CurrentSectionDataOffset;
+        S.Header.NumberOfRelocations = S.Relocations.size();
+        CurrentSectionDataOffset +=
+            S.Header.NumberOfRelocations * COFF::RelocationSize;
       }
-      // TODO: Handle alignment.
     } else {
-      i->Header.SizeOfRawData = 0;
+      S.Header.SizeOfRawData = 0;
-      i->Header.PointerToRawData = 0;
+      S.Header.PointerToRawData = 0;
     }
   }
 
@@ -188,7 +214,10 @@ static bool layoutCOFF(COFFParser &CP) {
   // Store all the allocated start addresses in the header.
   CP.Obj.Header.NumberOfSections = CP.Obj.Sections.size();
   CP.Obj.Header.NumberOfSymbols = NumberOfSymbols;
-  CP.Obj.Header.PointerToSymbolTable = SymbolTableStart;
+  if (NumberOfSymbols > 0 || CP.StringTable.size() > 4)
+    CP.Obj.Header.PointerToSymbolTable = SymbolTableStart;
+  else
+    CP.Obj.Header.PointerToSymbolTable = 0;
 
   *reinterpret_cast<support::ulittle32_t *>(&CP.StringTable[0])
     = CP.StringTable.size();
@@ -251,7 +280,83 @@ num_zeros_impl num_zeros(size_t N) {
   return NZI;
 }
 
-bool writeCOFF(COFFParser &CP, raw_ostream &OS) {
+template <typename T>
+static void initializeOptionalHeader(COFFParser &CP, uint16_t Magic, T Header) {
+  memset(Header, 0, sizeof(*Header));
+  Header->Magic = Magic;
+  Header->SectionAlignment = CP.Obj.OptionalHeader->Header.SectionAlignment;
+  uint32_t SizeOfCode = 0, SizeOfInitializedData = 0,
+           SizeOfUninitializedData = 0;
+  uint32_t SizeOfHeaders = RoundUpToAlignment(
+      CP.SectionTableStart + CP.SectionTableSize, Header->SectionAlignment);
+  uint32_t SizeOfImage = SizeOfHeaders;
+  for (const COFFYAML::Section &S : CP.Obj.Sections) {
+    if (S.Header.Characteristics & COFF::IMAGE_SCN_CNT_CODE)
+      SizeOfCode += S.Header.SizeOfRawData;
+    if (S.Header.Characteristics & COFF::IMAGE_SCN_CNT_INITIALIZED_DATA)
+      SizeOfInitializedData += S.Header.SizeOfRawData;
+    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
+    if (S.Header.VirtualAddress)
+      SizeOfImage +=
+          RoundUpToAlignment(S.Header.VirtualSize, Header->SectionAlignment);
+  }
+  Header->SizeOfCode = SizeOfCode;
+  Header->SizeOfInitializedData = SizeOfInitializedData;
+  Header->SizeOfUninitializedData = SizeOfUninitializedData;
+  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 =
+      CP.Obj.OptionalHeader->Header.MinorOperatingSystemVersion;
+  Header->MajorImageVersion =
+      CP.Obj.OptionalHeader->Header.MajorImageVersion;
+  Header->MinorImageVersion =
+      CP.Obj.OptionalHeader->Header.MinorImageVersion;
+  Header->MajorSubsystemVersion =
+      CP.Obj.OptionalHeader->Header.MajorSubsystemVersion;
+  Header->MinorSubsystemVersion =
+      CP.Obj.OptionalHeader->Header.MinorSubsystemVersion;
+  Header->SizeOfImage = SizeOfImage;
+  Header->SizeOfHeaders = SizeOfHeaders;
+  Header->Subsystem = CP.Obj.OptionalHeader->Header.Subsystem;
+  Header->DLLCharacteristics = CP.Obj.OptionalHeader->Header.DLLCharacteristics;
+  Header->SizeOfStackReserve = CP.Obj.OptionalHeader->Header.SizeOfStackReserve;
+  Header->SizeOfStackCommit = CP.Obj.OptionalHeader->Header.SizeOfStackCommit;
+  Header->SizeOfHeapReserve = CP.Obj.OptionalHeader->Header.SizeOfHeapReserve;
+  Header->SizeOfHeapCommit = CP.Obj.OptionalHeader->Header.SizeOfHeapCommit;
+  Header->NumberOfRvaAndSize = COFF::NUM_DATA_DIRECTORIES + 1;
+}
+
+static bool writeCOFF(COFFParser &CP, raw_ostream &OS) {
+  if (CP.isPE()) {
+    // PE files start with a DOS stub.
+    object::dos_header DH;
+    memset(&DH, 0, sizeof(DH));
+
+    // DOS EXEs start with "MZ" magic.
+    DH.Magic[0] = 'M';
+    DH.Magic[1] = 'Z';
+    // Initializing the AddressOfRelocationTable is strictly optional but
+    // mollifies certain tools which expect it to have a value greater than
+    // 0x40.
+    DH.AddressOfRelocationTable = sizeof(DH);
+    // This is the address of the PE signature.
+    DH.AddressOfNewExeHeader = 128;
+
+    // Write out our DOS stub.
+    OS.write(reinterpret_cast<char *>(&DH), sizeof(DH));
+    // Write padding until we reach the position of where our PE signature
+    // should live.
+    OS << num_zeros(DH.AddressOfNewExeHeader - sizeof(DH));
+    // Write out the PE signature.
+    OS.write(COFF::PEMagic, sizeof(COFF::PEMagic));
+  }
   if (CP.useBigObj()) {
     OS << binary_le(static_cast<uint16_t>(COFF::IMAGE_FILE_MACHINE_UNKNOWN))
        << binary_le(static_cast<uint16_t>(0xffff))
@@ -275,6 +380,29 @@ bool writeCOFF(COFFParser &CP, raw_ostream &OS) {
        << binary_le(CP.Obj.Header.SizeOfOptionalHeader)
        << binary_le(CP.Obj.Header.Characteristics);
   }
+  if (CP.isPE()) {
+    if (CP.is64Bit()) {
+      object::pe32plus_header PEH;
+      initializeOptionalHeader(CP, COFF::PE32Header::PE32_PLUS, &PEH);
+      OS.write(reinterpret_cast<char *>(&PEH), sizeof(PEH));
+    } else {
+      object::pe32_header PEH;
+      initializeOptionalHeader(CP, COFF::PE32Header::PE32, &PEH);
+      OS.write(reinterpret_cast<char *>(&PEH), sizeof(PEH));
+    }
+    for (const Optional<COFF::DataDirectory> &DD :
+         CP.Obj.OptionalHeader->DataDirectories) {
+      if (!DD.hasValue()) {
+        OS << zeros(uint32_t(0));
+        OS << zeros(uint32_t(0));
+      } else {
+        OS << binary_le(DD->RelativeVirtualAddress);
+        OS << binary_le(DD->Size);
+      }
+    }
+    OS << zeros(uint32_t(0));
+    OS << zeros(uint32_t(0));
+  }
 
   // Output section table.
   for (std::vector<COFFYAML::Section>::iterator i = CP.Obj.Sections.begin(),
@@ -302,12 +430,13 @@ bool writeCOFF(COFFParser &CP, raw_ostream &OS) {
   }
 
   // Output section data.
-  for (std::vector<COFFYAML::Section>::iterator i = CP.Obj.Sections.begin(),
+  for (const COFFYAML::Section &S : CP.Obj.Sections) {
-                                                e = CP.Obj.Sections.end();
+    if (!S.Header.SizeOfRawData)
-                                                i != e; ++i) {
+      continue;
-    i->SectionData.writeAsBinary(OS);
+    OS << num_zeros(S.Header.PointerToRawData - OS.tell());
-    for (unsigned I2 = 0, E2 = i->Relocations.size(); I2 != E2; ++I2) {
+    S.SectionData.writeAsBinary(OS);
-      const COFFYAML::Relocation &R = i->Relocations[I2];
+    OS << num_zeros(S.Header.SizeOfRawData - S.SectionData.binary_size());
+    for (const COFFYAML::Relocation &R : S.Relocations) {
       uint32_t SymbolTableIndex = SymbolTableIndexMap[R.SymbolName];
       OS << binary_le(R.VirtualAddress)
          << binary_le(SymbolTableIndex)
@@ -377,7 +506,8 @@ bool writeCOFF(COFFParser &CP, raw_ostream &OS) {
   }
 
   // Output string table.
-  OS.write(&CP.StringTable[0], CP.StringTable.size());
+  if (CP.Obj.Header.PointerToSymbolTable)
+    OS.write(&CP.StringTable[0], CP.StringTable.size());
   return true;
 }
 
@@ -395,6 +525,10 @@ int yaml2coff(yaml::Input &YIn, raw_ostream &Out) {
     return 1;
   }
 
+  if (!layoutOptionalHeader(CP)) {
+    errs() << "yaml2obj: Failed to layout optional header for COFF file!\n";
+    return 1;
+  }
   if (!layoutCOFF(CP)) {
     errs() << "yaml2obj: Failed to layout COFF file!\n";
     return 1;