[PECOFF] Add /stub option.

llvm-svn: 194757

3 files changed, 22 insertions, 3 deletions

diff --git a/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h b/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h
index 52ea0ed..481354c 100644
--- a/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h
+++ b/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h
@@ -230,9 +230,8 @@ public:
     return it == _sectionAttributeMask.end() ? 0 : it->second;
   }
 
-  const std::vector<uint8_t> &getDosStub() const {
-    return _dosStub;
-  }
+  void setDosStub(std::vector<uint8_t> &data) { _dosStub = std::move(data); }
+  const std::vector<uint8_t> &getDosStub() const { return _dosStub; }
 
   StringRef allocateString(StringRef ref) const {
     char *x = _allocator.Allocate<char>(ref.size() + 1);
diff --git a/lld/lib/Driver/WinLinkDriver.cpp b/lld/lib/Driver/WinLinkDriver.cpp
index 35cf6dd..1b5fd49 100644
--- a/lld/lib/Driver/WinLinkDriver.cpp
+++ b/lld/lib/Driver/WinLinkDriver.cpp
@@ -203,6 +203,14 @@ bool parseSection(StringRef option, std::string &section,
   return true;
 }
 
+bool readFile(StringRef path, std::vector<uint8_t> &result) {
+  OwningPtr<MemoryBuffer> buf;
+  if (MemoryBuffer::getFile(path, buf))
+    return false;
+  result = std::vector<uint8_t>(buf->getBufferStart(), buf->getBufferEnd());
+  return true;
+}
+
 // Parse /manifest:EMBED[,ID=#]|NO.
 bool parseManifest(StringRef option, bool &enable, bool &embed, int &id) {
   if (option.equals_lower("no")) {
@@ -828,6 +836,17 @@ WinLinkDriver::parse(int argc, const char *argv[], PECOFFLinkingContext &ctx,
       ctx.setSwapRunFromNet(true);
       break;
 
+    case OPT_stub: {
+      std::vector<uint8_t> contents;
+      if (!readFile(inputArg->getValue(), contents)) {
+        diagnostics << "Failed to read DOS stub file "
+                    << inputArg->getValue() << "\n";
+        return false;
+      }
+      ctx.setDosStub(contents);
+      break;
+    }
+
     case OPT_incl:
       ctx.addInitialUndefinedSymbol(ctx.allocateString(inputArg->getValue()));
       break;
diff --git a/lld/lib/Driver/WinLinkOptions.td b/lld/lib/Driver/WinLinkOptions.td
index 5cf7a55..14500d7 100644
--- a/lld/lib/Driver/WinLinkOptions.td
+++ b/lld/lib/Driver/WinLinkOptions.td
@@ -33,6 +33,7 @@ def version : P<"version", "Specify a version number in the PE header">;
 def merge   : P<"merge", "Combine sections">;
 def section : P<"section", "Specify section attributes">;
 def subsystem : P<"subsystem", "Specify subsystem">;
+def stub    : P<"stub", "Specify DOS stub file">;
 
 def manifest : F<"manifest">;
 def manifest_colon : P<"manifest", "Create manifest file">;