[lld] fb974e8 - [LLD][COFF] Add support for custom DOS stub (#122561)
via llvm-commits
llvm-commits at lists.llvm.org
Mon Jan 20 13:39:03 PST 2025
Author: kkent030315
Date: 2025-01-20T23:38:59+02:00
New Revision: fb974e89095af01a15cf959ba0694c0beb404b9f
URL: https://github.com/llvm/llvm-project/commit/fb974e89095af01a15cf959ba0694c0beb404b9f
DIFF: https://github.com/llvm/llvm-project/commit/fb974e89095af01a15cf959ba0694c0beb404b9f.diff
LOG: [LLD][COFF] Add support for custom DOS stub (#122561)
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.
Added:
lld/test/COFF/Inputs/stub63mz
lld/test/COFF/Inputs/stub64mz
lld/test/COFF/Inputs/stub64zz
lld/test/COFF/Inputs/stub68mz
lld/test/COFF/stub.test
Modified:
lld/COFF/Config.h
lld/COFF/Driver.cpp
lld/COFF/Driver.h
lld/COFF/DriverUtils.cpp
lld/COFF/Writer.cpp
Removed:
################################################################################
diff --git a/lld/COFF/Config.h b/lld/COFF/Config.h
index 924560fef0231d..b08427f738bb48 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<MemoryBuffer> 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 898c6c17d2062a..b257071c970867 100644
--- a/lld/COFF/Driver.cpp
+++ b/lld/COFF/Driver.cpp
@@ -2298,6 +2298,10 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> 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 8ce2e13129ba67..5f65bd7f8d0974 100644
--- a/lld/COFF/Driver.h
+++ b/lld/COFF/Driver.h
@@ -218,6 +218,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 "[:<integer>]"
void parseFunctionPadMin(llvm::opt::Arg *a);
diff --git a/lld/COFF/DriverUtils.cpp b/lld/COFF/DriverUtils.cpp
index 1148be09fb10cc..19abd4806d53fc 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<MemoryBuffer> 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 536c1eef5e49c8..73b4ba3bf3e371 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<DynamicRelocsChunk>();
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 <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->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 00000000000000..2a8954d2d6917f
Binary files /dev/null and b/lld/test/COFF/Inputs/stub63mz
diff er
diff --git a/lld/test/COFF/Inputs/stub64mz b/lld/test/COFF/Inputs/stub64mz
new file mode 100644
index 00000000000000..aaeb005adb54cb
Binary files /dev/null and b/lld/test/COFF/Inputs/stub64mz
diff er
diff --git a/lld/test/COFF/Inputs/stub64zz b/lld/test/COFF/Inputs/stub64zz
new file mode 100644
index 00000000000000..fa58df18aabe77
Binary files /dev/null and b/lld/test/COFF/Inputs/stub64zz
diff er
diff --git a/lld/test/COFF/Inputs/stub68mz b/lld/test/COFF/Inputs/stub68mz
new file mode 100644
index 00000000000000..42b72259465363
Binary files /dev/null and b/lld/test/COFF/Inputs/stub68mz
diff er
diff --git a/lld/test/COFF/stub.test b/lld/test/COFF/stub.test
new file mode 100644
index 00000000000000..84de6ed84c9579
--- /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: {{.*}}
More information about the llvm-commits
mailing list