[PECOFF] Emit PE32+ file header.

llvm-svn: 200128

3 files changed, 123 insertions, 24 deletions

diff --git a/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp b/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
index 1ff59f4..83457ac 100644
--- a/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
@@ -123,6 +123,7 @@ private:
 };
 
 /// A PEHeaderChunk represents PE header including COFF header.
+template <class PEHeader>
 class PEHeaderChunk : public HeaderChunk {
 public:
   explicit PEHeaderChunk(const PECOFFLinkingContext &ctx);
@@ -136,7 +137,7 @@ public:
 
   void setSizeOfCode(uint64_t size) { _peHeader.SizeOfCode = size; }
   void setBaseOfCode(uint32_t rva) { _peHeader.BaseOfCode = rva; }
-  void setBaseOfData(uint32_t rva) { _peHeader.BaseOfData = rva; }
+  void setBaseOfData(uint32_t rva);
   void setSizeOfImage(uint32_t size) { _peHeader.SizeOfImage = size; }
 
   void setSizeOfInitializedData(uint64_t size) {
@@ -155,7 +156,7 @@ public:
 
 private:
   llvm::object::coff_file_header _coffHeader;
-  llvm::object::pe32_header _peHeader;
+  PEHeader _peHeader;
 };
 
 /// A SectionHeaderTableChunk represents Section Table Header of PE/COFF
@@ -310,7 +311,8 @@ private:
   std::vector<uint8_t> _contents;
 };
 
-PEHeaderChunk::PEHeaderChunk(const PECOFFLinkingContext &ctx)
+template <class PEHeader>
+PEHeaderChunk<PEHeader>::PEHeaderChunk(const PECOFFLinkingContext &ctx)
     : HeaderChunk() {
   // Set the size of the chunk and initialize the header with null bytes.
   _size = sizeof(llvm::COFF::PEMagic) + sizeof(_coffHeader) + sizeof(_peHeader);
@@ -400,7 +402,18 @@ PEHeaderChunk::PEHeaderChunk(const PECOFFLinkingContext &ctx)
   _peHeader.NumberOfRvaAndSize = 16;
 }
 
-void PEHeaderChunk::write(uint8_t *buffer) {
+template <>
+void PEHeaderChunk<llvm::object::pe32_header>::setBaseOfData(uint32_t rva) {
+  _peHeader.BaseOfData = rva;
+}
+
+template <>
+void PEHeaderChunk<llvm::object::pe32plus_header>::setBaseOfData(uint32_t rva) {
+  // BaseOfData field does not exist in PE32+ header.
+}
+
+template <class PEHeader>
+void PEHeaderChunk<PEHeader>::write(uint8_t *buffer) {
   std::memcpy(buffer, llvm::COFF::PEMagic, sizeof(llvm::COFF::PEMagic));
   buffer += sizeof(llvm::COFF::PEMagic);
   std::memcpy(buffer, &_coffHeader, sizeof(_coffHeader));
@@ -742,7 +755,7 @@ public:
       : _ctx(context), _numSections(0), _imageSizeInMemory(PAGE_SIZE),
         _imageSizeOnDisk(0) {}
 
-  void build(const File &linkedFile);
+  template <class PEHeader> void build(const File &linkedFile);
   virtual error_code writeFile(const File &linkedFile, StringRef path);
 
 private:
@@ -751,7 +764,6 @@ private:
   void addChunk(Chunk *chunk);
   void addSectionChunk(SectionChunk *chunk, SectionHeaderTableChunk *table);
   void setImageSizeOnDisk();
-  void setAddressOfEntryPoint(AtomChunk *text, PEHeaderChunk *peHeader);
   uint64_t
   calcSectionSize(llvm::COFF::SectionCharacteristics sectionType) const;
 
@@ -829,13 +841,14 @@ void groupAtoms(const PECOFFLinkingContext &ctx, const File &file,
 }
 
 // Create all chunks that consist of the output file.
+template <class PEHeader>
 void PECOFFWriter::build(const File &linkedFile) {
   AtomVectorMap atoms;
   groupAtoms(_ctx, linkedFile, atoms);
 
   // Create file chunks and add them to the list.
   auto *dosStub = new DOSStubChunk(_ctx);
-  auto *peHeader = new PEHeaderChunk(_ctx);
+  auto *peHeader = new PEHeaderChunk<PEHeader>(_ctx);
   auto *dataDirectory = new DataDirectoryChunk();
   auto *sectionTable = new SectionHeaderTableChunk();
   addChunk(dosStub);
@@ -871,7 +884,19 @@ void PECOFFWriter::build(const File &linkedFile) {
       continue;
     if (section->getSectionName() == ".text") {
       peHeader->setBaseOfCode(section->getVirtualAddress());
-      setAddressOfEntryPoint(dyn_cast<AtomChunk>(section), peHeader);
+
+      // Find the virtual address of the entry point symbol if any.  PECOFF spec
+      // says that entry point for dll images is optional, in which case it must
+      // be set to 0.
+      if (_ctx.entrySymbolName().empty() && _ctx.isDll()) {
+        peHeader->setAddressOfEntryPoint(0);
+      } else {
+        uint64_t entryPointAddress =
+            dyn_cast<AtomChunk>(section)
+                ->getAtomVirtualAddress(_ctx.entrySymbolName());
+        if (entryPointAddress != 0)
+          peHeader->setAddressOfEntryPoint(entryPointAddress);
+      }
     }
     if (section->getSectionName() == ".data")
       peHeader->setBaseOfData(section->getVirtualAddress());
@@ -897,7 +922,11 @@ void PECOFFWriter::build(const File &linkedFile) {
 }
 
 error_code PECOFFWriter::writeFile(const File &linkedFile, StringRef path) {
-  this->build(linkedFile);
+  if (_ctx.is64Bit()) {
+    this->build<llvm::object::pe32plus_header>(linkedFile);
+  } else {
+    this->build<llvm::object::pe32_header>(linkedFile);
+  }
 
   uint64_t totalSize = _chunks.back()->fileOffset() + _chunks.back()->size();
   OwningPtr<llvm::FileOutputBuffer> buffer;
@@ -977,21 +1006,6 @@ void PECOFFWriter::setImageSizeOnDisk() {
   }
 }
 
-void PECOFFWriter::setAddressOfEntryPoint(AtomChunk *text,
-                                          PEHeaderChunk *peHeader) {
-  // Find the virtual address of the entry point symbol if any.
-  // PECOFF spec says that entry point for dll images is optional, in which
-  // case it must be set to 0.
-  if (_ctx.entrySymbolName().empty() && _ctx.isDll()) {
-    peHeader->setAddressOfEntryPoint(0);
-  } else {
-    uint64_t entryPointAddress =
-        text->getAtomVirtualAddress(_ctx.entrySymbolName());
-    if (entryPointAddress != 0)
-      peHeader->setAddressOfEntryPoint(entryPointAddress);
-  }
-}
-
 uint64_t PECOFFWriter::calcSectionSize(
     llvm::COFF::SectionCharacteristics sectionType) const {
   uint64_t ret = 0;
diff --git a/lld/test/pecoff/Inputs/nop64.obj b/lld/test/pecoff/Inputs/nop64.obj
new file mode 100644
index 0000000..eec5310
--- /dev/null
+++ b/lld/test/pecoff/Inputs/nop64.obj
diff --git a/lld/test/pecoff/peplus.test b/lld/test/pecoff/peplus.test
new file mode 100644
index 0000000..0d7bdc0
--- /dev/null
+++ b/lld/test/pecoff/peplus.test
@@ -0,0 +1,85 @@
+# RUN: lld -flavor link /out:%t.exe /subsystem:console /machine:x64 \
+# RUN:   /entry:start -- %p/Inputs/nop64.obj
+# RUN: llvm-readobj -file-headers %t.exe | FileCheck %s
+
+CHECK:      Format: COFF-x86-64
+CHECK-NEXT: Arch: x86_64
+CHECK-NEXT: AddressSize: 64bit
+CHECK-NEXT: ImageFileHeader {
+CHECK-NEXT:   Machine: IMAGE_FILE_MACHINE_AMD64 (0x8664)
+CHECK-NEXT:   SectionCount: 3
+CHECK-NEXT:   TimeDateStamp:
+CHECK-NEXT:   PointerToSymbolTable: 0x0
+CHECK-NEXT:   SymbolCount: 0
+CHECK-NEXT:   OptionalHeaderSize: 224
+CHECK-NEXT:   Characteristics [ (0x102)
+CHECK-NEXT:     IMAGE_FILE_32BIT_MACHINE (0x100)
+CHECK-NEXT:     IMAGE_FILE_EXECUTABLE_IMAGE (0x2)
+CHECK-NEXT:   ]
+CHECK-NEXT: }
+CHECK-NEXT: ImageOptionalHeader {
+CHECK-NEXT:   MajorLinkerVersion: 0
+CHECK-NEXT:   MinorLinkerVersion: 0
+CHECK-NEXT:   SizeOfCode: 1
+CHECK-NEXT:   SizeOfInitializedData: 108
+CHECK-NEXT:   SizeOfUninitializedData: 0
+CHECK-NEXT:   AddressOfEntryPoint: 0x2000
+CHECK-NEXT:   BaseOfCode: 0x2000
+CHECK-NEXT:   ImageBase: 0x400000
+CHECK-NEXT:   SectionAlignment: 4096
+CHECK-NEXT:   FileAlignment: 512
+CHECK-NEXT:   MajorOperatingSystemVersion: 6
+CHECK-NEXT:   MinorOperatingSystemVersion: 0
+CHECK-NEXT:   MajorImageVersion: 0
+CHECK-NEXT:   MinorImageVersion: 0
+CHECK-NEXT:   MajorSubsystemVersion: 6
+CHECK-NEXT:   MinorSubsystemVersion: 0
+CHECK-NEXT:   SizeOfImage: 12288
+CHECK-NEXT:   SizeOfHeaders: 512
+CHECK-NEXT:   Subsystem: IMAGE_SUBSYSTEM_WINDOWS_CUI (0x3)
+CHECK-NEXT:   Subsystem [ (0x8540)
+CHECK-NEXT:     IMAGE_DLL_CHARACTERISTICS_DYNAMIC_BASE (0x40)
+CHECK-NEXT:     IMAGE_DLL_CHARACTERISTICS_NO_SEH (0x400)
+CHECK-NEXT:     IMAGE_DLL_CHARACTERISTICS_NX_COMPAT (0x100)
+CHECK-NEXT:     IMAGE_DLL_CHARACTERISTICS_TERMINAL_SERVER_AWARE (0x8000)
+CHECK-NEXT:   ]
+CHECK-NEXT:   SizeOfStackReserve: 1048576
+CHECK-NEXT:   SizeOfStackCommit: 4096
+CHECK-NEXT:   SizeOfHeapReserve: 1048576
+CHECK-NEXT:   SizeOfHeapCommit: 4096
+CHECK-NEXT:   NumberOfRvaAndSize: 16
+CHECK-NEXT:   DataDirectory {
+CHECK-NEXT:     ExportTableRVA: 0x0
+CHECK-NEXT:     ExportTableSize: 0x0
+CHECK-NEXT:     ImportTableRVA: 0x0
+CHECK-NEXT:     ImportTableSize: 0x0
+CHECK-NEXT:     ResourceTableRVA: 0x0
+CHECK-NEXT:     ResourceTableSize: 0x0
+CHECK-NEXT:     ExceptionTableRVA: 0x0
+CHECK-NEXT:     ExceptionTableSize: 0x0
+CHECK-NEXT:     CertificateTableRVA: 0x0
+CHECK-NEXT:     CertificateTableSize: 0x0
+CHECK-NEXT:     BaseRelocationTableRVA: 0x0
+CHECK-NEXT:     BaseRelocationTableSize: 0x0
+CHECK-NEXT:     DebugRVA: 0x0
+CHECK-NEXT:     DebugSize: 0x0
+CHECK-NEXT:     ArchitectureRVA: 0x0
+CHECK-NEXT:     ArchitectureSize: 0x0
+CHECK-NEXT:     GlobalPtrRVA: 0x0
+CHECK-NEXT:     GlobalPtrSize: 0x0
+CHECK-NEXT:     TLSTableRVA: 0x0
+CHECK-NEXT:     TLSTableSize: 0x0
+CHECK-NEXT:     LoadConfigTableRVA: 0x0
+CHECK-NEXT:     LoadConfigTableSize: 0x0
+CHECK-NEXT:     BoundImportRVA: 0x0
+CHECK-NEXT:     BoundImportSize: 0x0
+CHECK-NEXT:     IATRVA: 0x0
+CHECK-NEXT:     IATSize: 0x0
+CHECK-NEXT:     DelayImportDescriptorRVA: 0x0
+CHECK-NEXT:     DelayImportDescriptorSize: 0x0
+CHECK-NEXT:     CLRRuntimeHeaderRVA: 0x0
+CHECK-NEXT:     CLRRuntimeHeaderSize: 0x0
+CHECK-NEXT:     ReservedRVA: 0x0
+CHECK-NEXT:     ReservedSize: 0x0
+CHECK-NEXT:   }
+CHECK-NEXT: }