Move DOS stub data to PECOFFLinkingContext for /stub option.

llvm-svn: 194754

3 files changed, 27 insertions, 19 deletions

diff --git a/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h b/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h
index f43766f01fb..52ea0ed8ce4 100644
--- a/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h
+++ b/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h
@@ -43,7 +43,7 @@ public:
         _terminalServerAware(true), _dynamicBaseEnabled(true),
         _createManifest(true), _embedManifest(false), _manifestId(1),
         _manifestLevel("'asInvoker'"), _manifestUiAccess("'false'"),
-        _imageType(ImageType::IMAGE_EXE) {
+        _imageType(ImageType::IMAGE_EXE), _dosStub(_defaultDosStub) {
     setDeadStripping(true);
   }
 
@@ -230,6 +230,10 @@ public:
     return it == _sectionAttributeMask.end() ? 0 : it->second;
   }
 
+  const std::vector<uint8_t> &getDosStub() const {
+    return _dosStub;
+  }
+
   StringRef allocateString(StringRef ref) const {
     char *x = _allocator.Allocate<char>(ref.size() + 1);
     memcpy(x, ref.data(), ref.size());
@@ -311,6 +315,13 @@ private:
 
   // List of files that will be removed on destruction.
   std::vector<std::unique_ptr<llvm::FileRemover> > _tempFiles;
+
+  // DOS Stub. DOS stub is data located at the beginning of PE/COFF file.
+  // Windows loader do not really care about DOS stub contents, but it's usually
+  // a small DOS program that prints out a message "This program requires
+  // Microsoft Windows." This feature was somewhat useful before Windows 95.
+  std::vector<uint8_t> _dosStub;
+  static const std::vector<uint8_t> _defaultDosStub;
 };
 
 } // end namespace lld
diff --git a/lld/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp b/lld/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp
index 76c7731f8e4..9ba6750a914 100644
--- a/lld/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp
@@ -29,6 +29,13 @@
 
 namespace lld {
 
+namespace {
+uint8_t DefaultDosStub[128] = {'M', 'Z'};
+};
+
+const std::vector<uint8_t> PECOFFLinkingContext::_defaultDosStub(
+    DefaultDosStub, DefaultDosStub + sizeof(DefaultDosStub));
+
 bool PECOFFLinkingContext::validateImpl(raw_ostream &diagnostics) {
   if (_stackReserve < _stackCommit) {
     diagnostics << "Invalid stack size: reserve size must be equal to or "
diff --git a/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp b/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
index 1edf6ccee81..8ea054880cb 100644
--- a/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
@@ -106,29 +106,19 @@ public:
 /// of PE/COFF files.
 class DOSStubChunk : public HeaderChunk {
 public:
-  DOSStubChunk() : HeaderChunk() {
-    // Make the DOS stub occupy the first 128 bytes of an exe. Technically
-    // this can be as small as 64 bytes, but GNU binutil's objdump cannot
-    // parse such irregular header.
-    _size = 128;
-
-    // A DOS stub is usually a small valid DOS program that prints out a message
-    // "This program requires Microsoft Windows" to help user who accidentally
-    // run a Windows executable on DOS. That's not a technical requirement, so
-    // we don't bother to emit such code, at least for now. We simply fill the
-    // DOS stub with null bytes.
-    std::memset(&_dosHeader, 0, sizeof(_dosHeader));
-
-    _dosHeader.Magic = 'M' | ('Z' << 8);
-    _dosHeader.AddressOfNewExeHeader = _size;
+  DOSStubChunk(const PECOFFLinkingContext &ctx)
+      : HeaderChunk(), _dosStub(ctx.getDosStub()) {
+    auto *header = reinterpret_cast<llvm::object::dos_header *>(&_dosStub[0]);
+    header->AddressOfNewExeHeader = _dosStub.size();
+    _size = _dosStub.size();
   }
 
   virtual void write(uint8_t *fileBuffer) {
-    std::memcpy(fileBuffer, &_dosHeader, sizeof(_dosHeader));
+    std::memcpy(fileBuffer, &_dosStub[0], _dosStub.size());
   }
 
 private:
-  llvm::object::dos_header _dosHeader;
+  std::vector<uint8_t> _dosStub;
 };
 
 /// A PEHeaderChunk represents PE header including COFF header.
@@ -810,7 +800,7 @@ public:
   // Create all chunks that consist of the output file.
   void build(const File &linkedFile) {
     // Create file chunks and add them to the list.
-    auto *dosStub = new DOSStubChunk();
+    auto *dosStub = new DOSStubChunk(_PECOFFLinkingContext);
     auto *peHeader = new PEHeaderChunk(_PECOFFLinkingContext);
     auto *dataDirectory = new DataDirectoryChunk(linkedFile);
     auto *sectionTable = new SectionHeaderTableChunk();