obj2yaml, yaml2obj: Add support for COFF executables
authorDavid Majnemer <david.majnemer@gmail.com>
Fri, 14 Nov 2014 08:15:42 +0000 (08:15 +0000)
committerDavid Majnemer <david.majnemer@gmail.com>
Fri, 14 Nov 2014 08:15:42 +0000 (08:15 +0000)
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
include/llvm/Object/COFFYAML.h
include/llvm/Support/COFF.h
lib/Object/COFFObjectFile.cpp
lib/Object/COFFYAML.cpp
tools/obj2yaml/coff2yaml.cpp
tools/yaml2obj/yaml2coff.cpp

index e04954577f80eb850467a22cdffddce546c5b40e..100f4199173092fad526ce638f60df8c1aa70a5d 100644 (file)
@@ -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;
 };
 
index 4aba08f75ddc90e14591e3d4f9b0100f84104299..12a25223bd37e71016de9f7c899068d558f87edd 100644 (file)
@@ -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);
index 86698e3c5a2716249e41ed9c0393bfeb733dc37d..074d4f9c9d139c230b74d09fb77c484e604a1e34 100644 (file)
@@ -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.
index 68bf7fd7f1ca719682bbf7cfbfb9c2b235ddb289..a1fc51ad4ca5784af337f11ed00ded5a9b5b2a32 100644 (file)
@@ -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 =
-      base() + getPointerToSymbolTable() +
-      getNumberOfSymbols() * getSymbolTableEntrySize();
+  uint32_t StringTableOffset = getPointerToSymbolTable() +
+                               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;
 }
index 49c5ddaee0a78121bfb692d47f09d2aa8a84ce41..9a24b531da9e5abe839570e54ab2fd304cf0d879 100644 (file)
@@ -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);
index f11f818e5a2787297846a3c6e762765838af485b..5baa644557cdbb002d57190b0e0b4f5b70b1c5b2 100644 (file)
@@ -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;
-  for (const auto &Section : Obj.sections()) {
-    const object::coff_section *Sect = Obj.getCOFFSection(Section);
-    COFFYAML::Section Sec;
-    Section.getName(Sec.Name);
-    Sec.Header.Characteristics = Sect->Characteristics;
-    Sec.Alignment = Section.getAlignment();
+  std::vector<COFFYAML::Section> &YAMLSections = YAMLObj.Sections;
+  for (const auto &ObjSection : Obj.sections()) {
+    const object::coff_section *COFFSection = Obj.getCOFFSection(ObjSection);
+    COFFYAML::Section NewYAMLSection;
+    ObjSection.getName(NewYAMLSection.Name);
+    NewYAMLSection.Header.Characteristics = COFFSection->Characteristics;
+    NewYAMLSection.Header.VirtualAddress = ObjSection.getAddress();
+    NewYAMLSection.Header.VirtualSize = COFFSection->VirtualSize;
+    NewYAMLSection.Alignment = ObjSection.getAlignment();
 
     ArrayRef<uint8_t> sectionData;
-    if (!Section.isBSS())
-      Obj.getSectionContents(Sect, sectionData);
-    Sec.SectionData = yaml::BinaryRef(sectionData);
+    if (!ObjSection.isBSS())
+      Obj.getSectionContents(COFFSection, 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;
-    Sections.push_back(Sec);
+    NewYAMLSection.Relocations = Relocations;
+    YAMLSections.push_back(NewYAMLSection);
   }
 }
 
index 5edce17f81aca8f622027f95355a564e53f94b46..2c71988bad905b332ce04ac19158058397b11af9 100644 (file)
@@ -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) {
-  uint32_t SectionTableStart = 0;
-  uint32_t SectionTableSize  = 0;
+static bool layoutOptionalHeader(COFFParser &CP) {
+  if (!CP.isPE())
+    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;
-  SectionTableSize = COFF::SectionSize * CP.Obj.Sections.size();
+  CP.SectionTableStart =
+      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(),
-                                                e = CP.Obj.Sections.end();
-                                                i != e; ++i) {
-    if (i->SectionData.binary_size() > 0) {
-      i->Header.SizeOfRawData = i->SectionData.binary_size();
-      i->Header.PointerToRawData = CurrentSectionDataOffset;
-      CurrentSectionDataOffset += i->Header.SizeOfRawData;
-      if (!i->Relocations.empty()) {
-        i->Header.PointerToRelocations = CurrentSectionDataOffset;
-        i->Header.NumberOfRelocations = i->Relocations.size();
-        CurrentSectionDataOffset += i->Header.NumberOfRelocations *
-          COFF::RelocationSize;
+  for (COFFYAML::Section &S : CP.Obj.Sections) {
+    if (S.SectionData.binary_size() > 0) {
+      CurrentSectionDataOffset = RoundUpToAlignment(
+          CurrentSectionDataOffset, CP.isPE() ? CP.getFileAlignment() : 4);
+      S.Header.SizeOfRawData = S.SectionData.binary_size();
+      if (CP.isPE())
+        S.Header.SizeOfRawData =
+            RoundUpToAlignment(S.Header.SizeOfRawData, CP.getFileAlignment());
+      S.Header.PointerToRawData = CurrentSectionDataOffset;
+      CurrentSectionDataOffset += S.Header.SizeOfRawData;
+      if (!S.Relocations.empty()) {
+        S.Header.PointerToRelocations = CurrentSectionDataOffset;
+        S.Header.NumberOfRelocations = S.Relocations.size();
+        CurrentSectionDataOffset +=
+            S.Header.NumberOfRelocations * COFF::RelocationSize;
       }
-      // TODO: Handle alignment.
     } else {
-      i->Header.SizeOfRawData = 0;
-      i->Header.PointerToRawData = 0;
+      S.Header.SizeOfRawData = 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(),
-                                                e = CP.Obj.Sections.end();
-                                                i != e; ++i) {
-    i->SectionData.writeAsBinary(OS);
-    for (unsigned I2 = 0, E2 = i->Relocations.size(); I2 != E2; ++I2) {
-      const COFFYAML::Relocation &R = i->Relocations[I2];
+  for (const COFFYAML::Section &S : CP.Obj.Sections) {
+    if (!S.Header.SizeOfRawData)
+      continue;
+    OS << num_zeros(S.Header.PointerToRawData - OS.tell());
+    S.SectionData.writeAsBinary(OS);
+    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;