[LLD][COFF] Add support for custom DOS stub

This change implements support for the /stub flag to align with MS link.exe. This option is useful when a program needs to optimize the DOS program that executes when the PE runs on DOS, avoiding the traditional hardcoded DOS program in LLD.

Author: kkent030315

lld/COFF/Config.h

@@ -115,6 +115,7 @@ struct Configuration {
   enum ManifestKind { Default, SideBySide, Embed, No };
   bool is64() const { return llvm::COFF::is64Bit(machine); }
 
+  llvm::SmallVector<uint8_t> stub;
   llvm::COFF::MachineTypes machine = IMAGE_FILE_MACHINE_UNKNOWN;
   bool machineInferred = false;
   size_t wordsize;

lld/COFF/Driver.cpp

@@ -2418,6 +2418,10 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
     config->noSEH = args.hasArg(OPT_noseh);
   }
 
+  // Handle /stub
+  if (auto *arg = args.getLastArg(OPT_stub))
+    parseStub(arg);
+
   // Handle /functionpadmin
   for (auto *arg : args.filtered(OPT_functionpadmin, OPT_functionpadmin_opt))
     parseFunctionPadMin(arg);

lld/COFF/Driver.h

@@ -236,6 +236,9 @@ class LinkerDriver {
   void parseSection(StringRef);
   void parseAligncomm(StringRef);
 
+  // Parses a MS-DOS stub file
+  void parseStub(llvm::opt::Arg *a);
+
   // Parses a string in the form of "[:<integer>]"
   void parseFunctionPadMin(llvm::opt::Arg *a);
 

lld/COFF/DriverUtils.cpp

@@ -246,6 +246,21 @@ void LinkerDriver::parseAligncomm(StringRef s) {
       std::max(ctx.config.alignComm[std::string(name)], 1 << v);
 }
 
+void LinkerDriver::parseStub(llvm::opt::Arg *a) {
+  StringRef arg = a->getValue();
+  std::unique_ptr<MemoryBuffer> stub = check(MemoryBuffer::getFile(arg));
+  size_t bufferSize = stub->getBufferSize();
+  const char *bufferStart = stub->getBufferStart();
+  // MS link.exe compatibility:
+  // 1. stub must be greater or equal than 64 bytes
+  // 2. stub must be 8-byte aligned
+  // 3. stub must be start with a valid dos signature 'MZ'
+  if (bufferSize < 0x40 || bufferSize % 8 != 0 ||
+      (bufferStart[0] != 'M' || bufferStart[1] != 'Z'))
+    Err(ctx) << "/stub: invalid format for MS-DOS stub file: " << arg;
+  ctx.config.stub.append(bufferStart, bufferStart + bufferSize);
+}
+
 // Parses /functionpadmin option argument.
 void LinkerDriver::parseFunctionPadMin(llvm::opt::Arg *a) {
   StringRef arg = a->getNumValues() ? a->getValue() : "";

lld/COFF/Writer.cpp

@@ -79,11 +79,6 @@ static_assert(sizeof(dosProgram) % 8 == 0,
 
 static const int dosStubSize = sizeof(dos_header) + sizeof(dosProgram);
 static_assert(dosStubSize % 8 == 0, "DOSStub size must be multiple of 8");
-static const uint32_t coffHeaderOffset = dosStubSize + sizeof(PEMagic);
-static const uint32_t peHeaderOffset =
-    coffHeaderOffset + sizeof(coff_file_header);
-static const uint32_t dataDirOffset64 =
-    peHeaderOffset + sizeof(pe32plus_header);
 
 static const int numberOfDataDirectory = 16;
 
@@ -315,6 +310,10 @@ class Writer {
   uint64_t sizeOfImage;
   uint64_t sizeOfHeaders;
 
+  uint32_t coffHeaderOffset;
+  uint32_t peHeaderOffset;
+  uint32_t dataDirOffset64;
+
   OutputSection *textSec;
   OutputSection *hexpthkSec;
   OutputSection *rdataSec;
@@ -1668,21 +1667,35 @@ template <typename PEHeaderTy> void Writer::writeHeader() {
   // When run under Windows, the loader looks at AddressOfNewExeHeader and uses
   // the PE header instead.
   Configuration *config = &ctx.config;
+
   uint8_t *buf = buffer->getBufferStart();
   auto *dos = reinterpret_cast<dos_header *>(buf);
-  buf += sizeof(dos_header);
-  dos->Magic[0] = 'M';
-  dos->Magic[1] = 'Z';
-  dos->UsedBytesInTheLastPage = dosStubSize % 512;
-  dos->FileSizeInPages = divideCeil(dosStubSize, 512);
-  dos->HeaderSizeInParagraphs = sizeof(dos_header) / 16;
-
-  dos->AddressOfRelocationTable = sizeof(dos_header);
-  dos->AddressOfNewExeHeader = dosStubSize;
 
   // Write DOS program.
-  memcpy(buf, dosProgram, sizeof(dosProgram));
-  buf += sizeof(dosProgram);
+  if (config->stub.size()) {
+    memcpy(buf, config->stub.data(), config->stub.size());
+    // MS link.exe accepts an invalid `e_lfanew` and updates it automatically.
+    // Replicate the same behaviour.
+    dos->AddressOfNewExeHeader = config->stub.size();
+    buf += config->stub.size();
+  } else {
+    buf += sizeof(dos_header);
+    dos->Magic[0] = 'M';
+    dos->Magic[1] = 'Z';
+    dos->UsedBytesInTheLastPage = dosStubSize % 512;
+    dos->FileSizeInPages = divideCeil(dosStubSize, 512);
+    dos->HeaderSizeInParagraphs = sizeof(dos_header) / 16;
+
+    dos->AddressOfRelocationTable = sizeof(dos_header);
+    dos->AddressOfNewExeHeader = dosStubSize;
+
+    memcpy(buf, dosProgram, sizeof(dosProgram));
+    buf += sizeof(dosProgram);
+  }
+
+  coffHeaderOffset = buf - buffer->getBufferStart() + sizeof(PEMagic);
+  peHeaderOffset = coffHeaderOffset + sizeof(coff_file_header);
+  dataDirOffset64 = peHeaderOffset + sizeof(pe32plus_header);
 
   // Write PE magic
   memcpy(buf, PEMagic, sizeof(PEMagic));

lld/test/COFF/Inputs/stub511mz

@@ -0,0 +1,31 @@
+# RUN: yaml2obj %p/Inputs/ret42.yaml -o %t.obj
+
+# RUN: lld-link /out:%t.exe /entry:main /stub:%p/Inputs/stub512mz %t.obj
+# RUN: llvm-readobj --file-headers %t.exe | FileCheck -check-prefix=CHECK1 %s
+
+CHECK1: Magic: MZ
+CHECK1: UsedBytesInTheLastPage: 144
+CHECK1: FileSizeInPages: 3
+CHECK1: NumberOfRelocationItems: 0
+CHECK1: HeaderSizeInParagraphs: 4
+CHECK1: MinimumExtraParagraphs: 0
+CHECK1: MaximumExtraParagraphs: 65535
+CHECK1: InitialRelativeSS: 0
+CHECK1: InitialSP: 184
+CHECK1: Checksum: 0
+CHECK1: InitialIP: 0
+CHECK1: InitialRelativeCS: 0
+CHECK1: AddressOfRelocationTable: 64
+CHECK1: OverlayNumber: 0
+CHECK1: OEMid: 0
+CHECK1: OEMinfo: 0
+CHECK1: AddressOfNewExeHeader: 64
+
+# RUN: not lld-link /out:%t.exe /entry:main /stub:%p/Inputs/stub512zz %t.obj 2>&1 | FileCheck -check-prefix=CHECK2 %s
+# CHECK2: lld-link: error: /stub: invalid format for MS-DOS stub file: {{.*}}
+
+# RUN: not lld-link /out:%t.exe /entry:main /stub:%p/Inputs/stub516mz %t.obj 2>&1 | FileCheck -check-prefix=CHECK3 %s
+# CHECK3: lld-link: error: /stub: invalid format for MS-DOS stub file: {{.*}}
+
+# RUN: not lld-link /out:%t.exe /entry:main /stub:%p/Inputs/stub511mz %t.obj 2>&1 | FileCheck -check-prefix=CHECK4 %s
+# CHECK4: lld-link: error: /stub: invalid format for MS-DOS stub file: {{.*}}