diff --git a/lld/COFF/Config.h b/lld/COFF/Config.h index 9e6b17e87c9e7..687c7b9529251 100644 --- a/lld/COFF/Config.h +++ b/lld/COFF/Config.h @@ -115,6 +115,7 @@ struct Configuration { enum ManifestKind { Default, SideBySide, Embed, No }; bool is64() const { return llvm::COFF::is64Bit(machine); } + std::unique_ptr dosStub; llvm::COFF::MachineTypes machine = IMAGE_FILE_MACHINE_UNKNOWN; bool machineInferred = false; size_t wordsize; diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp index 791382fd9bdd4..2ee04a16b9815 100644 --- a/lld/COFF/Driver.cpp +++ b/lld/COFF/Driver.cpp @@ -2418,6 +2418,10 @@ void LinkerDriver::linkerMain(ArrayRef argsArr) { config->noSEH = args.hasArg(OPT_noseh); } + // Handle /stub + if (auto *arg = args.getLastArg(OPT_stub)) + parseDosStub(arg->getValue()); + // Handle /functionpadmin for (auto *arg : args.filtered(OPT_functionpadmin, OPT_functionpadmin_opt)) parseFunctionPadMin(arg); diff --git a/lld/COFF/Driver.h b/lld/COFF/Driver.h index 5132568904298..9ac7f5dd96d27 100644 --- a/lld/COFF/Driver.h +++ b/lld/COFF/Driver.h @@ -236,6 +236,9 @@ class LinkerDriver { void parseSection(StringRef); void parseAligncomm(StringRef); + // Parses a MS-DOS stub file + void parseDosStub(StringRef path); + // Parses a string in the form of "[:]" void parseFunctionPadMin(llvm::opt::Arg *a); diff --git a/lld/COFF/DriverUtils.cpp b/lld/COFF/DriverUtils.cpp index 1148be09fb10c..19abd4806d53f 100644 --- a/lld/COFF/DriverUtils.cpp +++ b/lld/COFF/DriverUtils.cpp @@ -246,6 +246,22 @@ void LinkerDriver::parseAligncomm(StringRef s) { std::max(ctx.config.alignComm[std::string(name)], 1 << v); } +void LinkerDriver::parseDosStub(StringRef path) { + std::unique_ptr stub = + CHECK(MemoryBuffer::getFile(path), "could not open " + path); + size_t bufferSize = stub->getBufferSize(); + const char *bufferStart = stub->getBufferStart(); + // MS link.exe compatibility: + // 1. stub must be greater than or equal to 64 bytes + // 2. stub must start with a valid dos signature 'MZ' + if (bufferSize < 64) + Err(ctx) << "/stub: stub must be greater than or equal to 64 bytes: " + << path; + if (bufferStart[0] != 'M' || bufferStart[1] != 'Z') + Err(ctx) << "/stub: invalid DOS signature: " << path; + ctx.config.dosStub = std::move(stub); +} + // Parses /functionpadmin option argument. void LinkerDriver::parseFunctionPadMin(llvm::opt::Arg *a) { StringRef arg = a->getNumValues() ? a->getValue() : ""; diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp index eb82a9cc01593..907c291aa37e6 100644 --- a/lld/COFF/Writer.cpp +++ b/lld/COFF/Writer.cpp @@ -76,14 +76,8 @@ static unsigned char dosProgram[] = { }; static_assert(sizeof(dosProgram) % 8 == 0, "DOSProgram size must be multiple of 8"); - -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_assert((sizeof(dos_header) + sizeof(dosProgram)) % 8 == 0, + "DOSStub size must be multiple of 8"); static const int numberOfDataDirectory = 16; @@ -214,6 +208,7 @@ class Writer { void run(); private: + void calculateStubDependentSizes(); void createSections(); void createMiscChunks(); void createImportTables(); @@ -315,6 +310,11 @@ class Writer { uint64_t sizeOfImage; uint64_t sizeOfHeaders; + uint32_t dosStubSize; + uint32_t coffHeaderOffset; + uint32_t peHeaderOffset; + uint32_t dataDirOffset64; + OutputSection *textSec; OutputSection *hexpthkSec; OutputSection *rdataSec; @@ -728,10 +728,8 @@ void Writer::writePEChecksum() { uint32_t *buf = (uint32_t *)buffer->getBufferStart(); uint32_t size = (uint32_t)(buffer->getBufferSize()); - coff_file_header *coffHeader = - (coff_file_header *)((uint8_t *)buf + dosStubSize + sizeof(PEMagic)); - pe32_header *peHeader = - (pe32_header *)((uint8_t *)coffHeader + sizeof(coff_file_header)); + pe32_header *peHeader = (pe32_header *)((uint8_t *)buf + coffHeaderOffset + + sizeof(coff_file_header)); uint64_t sum = 0; uint32_t count = size; @@ -762,6 +760,7 @@ void Writer::run() { llvm::TimeTraceScope timeScope("Write PE"); ScopedTimer t1(ctx.codeLayoutTimer); + calculateStubDependentSizes(); if (ctx.config.machine == ARM64X) ctx.dynamicRelocs = make(); createImportTables(); @@ -1035,6 +1034,17 @@ void Writer::sortSections() { sortBySectionOrder(it.second->chunks); } +void Writer::calculateStubDependentSizes() { + if (ctx.config.dosStub) + dosStubSize = alignTo(ctx.config.dosStub->getBufferSize(), 8); + else + dosStubSize = sizeof(dos_header) + sizeof(dosProgram); + + coffHeaderOffset = dosStubSize + sizeof(PEMagic); + peHeaderOffset = coffHeaderOffset + sizeof(coff_file_header); + dataDirOffset64 = peHeaderOffset + sizeof(pe32plus_header); +} + // Create output section objects and add them to OutputSections. void Writer::createSections() { llvm::TimeTraceScope timeScope("Output sections"); @@ -1668,21 +1678,37 @@ template 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(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->dosStub) { + memcpy(buf, config->dosStub->getBufferStart(), + config->dosStub->getBufferSize()); + // MS link.exe accepts an invalid `e_lfanew` (AddressOfNewExeHeader) and + // updates it automatically. Replicate the same behaviour. + dos->AddressOfNewExeHeader = alignTo(config->dosStub->getBufferSize(), 8); + // Unlike MS link.exe, LLD accepts non-8-byte-aligned stubs. + // In that case, we add zero paddings ourselves. + buf += alignTo(config->dosStub->getBufferSize(), 8); + } 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); + } + + // Make sure DOS stub is aligned to 8 bytes at this point + assert((buf - buffer->getBufferStart()) % 8 == 0); // Write PE magic memcpy(buf, PEMagic, sizeof(PEMagic)); diff --git a/lld/test/COFF/Inputs/stub63mz b/lld/test/COFF/Inputs/stub63mz new file mode 100644 index 0000000000000..2a8954d2d6917 Binary files /dev/null and b/lld/test/COFF/Inputs/stub63mz differ diff --git a/lld/test/COFF/Inputs/stub64mz b/lld/test/COFF/Inputs/stub64mz new file mode 100644 index 0000000000000..aaeb005adb54c Binary files /dev/null and b/lld/test/COFF/Inputs/stub64mz differ diff --git a/lld/test/COFF/Inputs/stub64zz b/lld/test/COFF/Inputs/stub64zz new file mode 100644 index 0000000000000..fa58df18aabe7 Binary files /dev/null and b/lld/test/COFF/Inputs/stub64zz differ diff --git a/lld/test/COFF/Inputs/stub68mz b/lld/test/COFF/Inputs/stub68mz new file mode 100644 index 0000000000000..42b7225946536 Binary files /dev/null and b/lld/test/COFF/Inputs/stub68mz differ diff --git a/lld/test/COFF/stub.test b/lld/test/COFF/stub.test new file mode 100644 index 0000000000000..84de6ed84c957 --- /dev/null +++ b/lld/test/COFF/stub.test @@ -0,0 +1,55 @@ +# RUN: yaml2obj %p/Inputs/ret42.yaml -o %t.obj + +# RUN: lld-link /out:%t.exe /entry:main /stub:%p/Inputs/stub64mz %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 + +## Invalid DOS signature (must be `MZ`) +# RUN: not lld-link /out:%t.exe /entry:main /stub:%p/Inputs/stub64zz %t.obj 2>&1 | FileCheck -check-prefix=CHECK2 %s + +CHECK2: lld-link: error: /stub: invalid DOS signature: {{.*}} + +## Unlike MS linker, we accept non-8byte-aligned stubs and we add paddings ourselves +# RUN: lld-link /out:%t.exe /entry:main /stub:%p/Inputs/stub68mz %t.obj +# RUN: llvm-readobj --file-headers %t.exe | FileCheck -check-prefix=CHECK3 %s + +CHECK3: Magic: MZ +CHECK3: UsedBytesInTheLastPage: 144 +CHECK3: FileSizeInPages: 3 +CHECK3: NumberOfRelocationItems: 0 +CHECK3: HeaderSizeInParagraphs: 4 +CHECK3: MinimumExtraParagraphs: 0 +CHECK3: MaximumExtraParagraphs: 65535 +CHECK3: InitialRelativeSS: 0 +CHECK3: InitialSP: 184 +CHECK3: Checksum: 0 +CHECK3: InitialIP: 0 +CHECK3: InitialRelativeCS: 0 +CHECK3: AddressOfRelocationTable: 64 +CHECK3: OverlayNumber: 0 +CHECK3: OEMid: 0 +CHECK3: OEMinfo: 0 +## 68 is unaligned and rounded up to 72 by LLD +CHECK3: AddressOfNewExeHeader: 72 + +## Too-small stub (must be at least 64 bytes long) && Unaligned +# RUN: not lld-link /out:%t.exe /entry:main /stub:%p/Inputs/stub63mz %t.obj 2>&1 | FileCheck -check-prefix=CHECK4 %s + +CHECK4: lld-link: error: /stub: stub must be greater than or equal to 64 bytes: {{.*}}