[PATCH][lld][PECOFF] Add Support for entry point symbol name

[Jesús Serrano García]

Currently the PECOFF writer does not take into account the entry point 
symbol name, and the resulting image starts with the firts symbol in 
.text section. Currently the AddressOfEntryPoint is assumed to be in 
0x1000, wich is not correct according to PECOFF spec. In fact, this 
address may be 0 in DLL images, this is also handled in the patch 
allowing to set the image output type.
-------------- next part --------------
diff --git include/lld/ReaderWriter/PECOFFTargetInfo.h include/lld/ReaderWriter/PECOFFTargetInfo.h
index 10337b6..7d39dc2 100644
--- include/lld/ReaderWriter/PECOFFTargetInfo.h
+++ include/lld/ReaderWriter/PECOFFTargetInfo.h
@@ -29,7 +29,7 @@ public:
         _heapReserve(1024 * 1024), _heapCommit(4096),
         _subsystem(llvm::COFF::IMAGE_SUBSYSTEM_UNKNOWN), _minOSVersion(6, 0),
         _nxCompat(true), _largeAddressAware(false), _baseRelocationEnabled(true),
-        _terminalServerAware(true) {}
+        _terminalServerAware(true), _imageType(ImageType::imageExe) {}
 
   struct OSVersion {
     OSVersion(int v1, int v2) : majorVersion(v1), minorVersion(v2) {}
@@ -37,6 +37,11 @@ public:
     int minorVersion;
   };
 
+  enum ImageType {
+    imageExe,
+    imageDll
+  };
+
   virtual error_code parseFile(
       std::unique_ptr<MemoryBuffer> &mb,
       std::vector<std::unique_ptr<File>> &result) const;
@@ -93,6 +98,9 @@ public:
   virtual ErrorOr<Reference::Kind> relocKindFromString(StringRef str) const;
   virtual ErrorOr<std::string> stringFromRelocKind(Reference::Kind kind) const;
 
+  void setImageType(ImageType type) { _imageType = type; }
+  ImageType getImageType() const { return _imageType; }
+
   StringRef allocateString(const StringRef &ref) {
     char *x = _alloc.Allocate<char>(ref.size() + 1);
     memcpy(x, ref.data(), ref.size());
@@ -115,6 +123,7 @@ private:
   bool _largeAddressAware;
   bool _baseRelocationEnabled;
   bool _terminalServerAware;
+  ImageType _imageType;
 
   std::vector<StringRef> _inputSearchPaths;
   mutable std::unique_ptr<Reader> _reader;
diff --git lib/ReaderWriter/PECOFF/PECOFFTargetInfo.cpp lib/ReaderWriter/PECOFF/PECOFFTargetInfo.cpp
index b1f599c..ef5de77 100644
--- lib/ReaderWriter/PECOFF/PECOFFTargetInfo.cpp
+++ lib/ReaderWriter/PECOFF/PECOFFTargetInfo.cpp
@@ -76,6 +76,10 @@ bool PECOFFTargetInfo::validateImpl(raw_ostream &diagnostics) {
     return true;
   }
 
+  if (_entrySymbolName.empty() && _imageType == ImageType::imageExe) {
+    _entrySymbolName = "_main";
+  }
+	  
   _reader = createReaderPECOFF(*this);
   _writer = createWriterPECOFF(*this);
   return false;
diff --git lib/ReaderWriter/PECOFF/WriterPECOFF.cpp lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
index 34dc1fb..4984f09 100644
--- lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
+++ lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
@@ -245,6 +245,8 @@ public:
 
   virtual void setSizeOfImage(uint32_t size) { _peHeader.SizeOfImage = size; }
 
+  virtual void setAddressOfEntryPoint(uint32_t address) { _peHeader.AddressOfEntryPoint = address; }
+
 private:
   llvm::object::coff_file_header _coffHeader;
   llvm::object::pe32_header _peHeader;
@@ -351,6 +353,17 @@ public:
       layout->_virtualAddr += rva;
   }
 
+  virtual bool getAtomVirtualAddress(StringRef name, uint64_t& address)
+  {
+    for (auto atomLayout : _atomLayouts) {
+      if (atomLayout->_atom->name() == name) {
+        address = atomLayout->_virtualAddr;
+        return true;
+      }
+    }
+    return false;
+  }
+
   static bool classof(const Chunk *c) {
     Kind kind = c->getKind();
     return kind == kindSection || kind == kindDataDirectory;
@@ -747,6 +760,18 @@ public:
     peHeader->setSizeOfInitializedData(rdata->size() + data->size());
     peHeader->setNumberOfSections(_numSections);
     peHeader->setSizeOfImage(_imageSizeInMemory);
+
+    // Find the virtual address of the entry point symbol if any
+    // PeCOFF spec says that entry point for dll images is optional
+    if (_PECOFFTargetInfo.entrySymbolName().empty() && _PECOFFTargetInfo.getImageType() == PECOFFTargetInfo::imageDll) {
+      peHeader->setAddressOfEntryPoint(0);
+    } else {
+      uint64_t entryPointAddress = 0; 
+      if (text->getAtomVirtualAddress(_PECOFFTargetInfo.entrySymbolName(), entryPointAddress)) {
+        // TODO: AddressOfEntryPoint in pe header is 4 byte, have to check if entry point virtual address is within 32 bit space
+        peHeader->setAddressOfEntryPoint(entryPointAddress);
+      }
+    }
   }
 
   virtual error_code writeFile(const File &linkedFile, StringRef path) {