diff --git a/lld/COFF/Config.h b/lld/COFF/Config.h --- a/lld/COFF/Config.h +++ b/lld/COFF/Config.h @@ -289,6 +289,7 @@ bool autoImport = false; bool pseudoRelocs = false; bool stdcallFixup = false; + bool writeCheckSum = false; }; extern std::unique_ptr config; diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp --- a/lld/COFF/Driver.cpp +++ b/lld/COFF/Driver.cpp @@ -432,6 +432,10 @@ } break; } + case OPT_release: { + config->writeCheckSum = true; + break; + } // Only add flags here that link.exe accepts in // `#pragma comment(linker, "/flag")`-generated sections. case OPT_editandcontinue: @@ -2020,6 +2024,10 @@ if (errorCount()) return; + // Handle /RELEASE + if (args.hasArg(OPT_release)) + config->writeCheckSum = true; + // Handle /safeseh, x86 only, on by default, except for mingw. if (config->machine == I386) { config->safeSEH = args.hasFlag(OPT_safeseh, OPT_safeseh_no, !config->mingw); diff --git a/lld/COFF/Options.td b/lld/COFF/Options.td --- a/lld/COFF/Options.td +++ b/lld/COFF/Options.td @@ -161,6 +161,8 @@ def verbose : F<"verbose">; def wholearchive_flag : F<"wholearchive">, HelpText<"Include all object files from all libraries">; +def release : F<"release">, + HelpText<"Set the Checksum in the header of an PE file">; def force : F<"force">, HelpText<"Allow undefined and multiply defined symbols">; diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp --- a/lld/COFF/Writer.cpp +++ b/lld/COFF/Writer.cpp @@ -230,6 +230,7 @@ void setSectionPermissions(); void writeSections(); void writeBuildId(); + void writePEChecksum(); void sortSections(); void sortExceptionTable(); void sortCRTSectionChunks(std::vector &chunks); @@ -599,6 +600,54 @@ } } +void Writer::writePEChecksum() { + if (!config->writeCheckSum) { + return; + } + + // https://docs.microsoft.com/en-us/windows/win32/debug/pe-format#checksum + uint32_t checkSum = 0; + 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)); + uint32_t oldCheckSum = peHeader->CheckSum; + + auto CalcCheckSum = [](uint32_t StartValue, void *BaseAddress, + uint32_t WordCount) -> uint16_t { + uint16_t *p = (uint16_t *)BaseAddress; + uint32_t sum = StartValue; + for (uint32_t i = 0; i < WordCount; i++) { + sum += *p; + if (((sum >> 16) & 0xffff) != 0) { + sum = (sum & 0xffff) + ((sum >> 16) & 0xffff); + } + p++; + } + return (uint16_t)((sum & 0xffff) + ((sum >> 16) & 0xffff)); + }; + + checkSum = CalcCheckSum(0, buf, (size + 1) / sizeof(uint16_t)); + if ((checkSum & 0xffff) >= (oldCheckSum & 0xffff)) { + checkSum -= (oldCheckSum & 0xffff); + } else { + checkSum = (((checkSum & 0xffff) - (oldCheckSum & 0xffff)) & 0xFFFF) - 1; + } + + if ((checkSum & 0xffff) >= ((oldCheckSum >> 16) & 0xffff)) { + checkSum -= ((oldCheckSum >> 16) & 0xffff); + } else { + checkSum = + (((checkSum & 0xffff) - ((oldCheckSum >> 16) & 0xffff)) & 0xFFFF) - 1; + } + + checkSum += size; + peHeader->CheckSum = checkSum; +} + // The main function of the writer. void Writer::run() { ScopedTimer t1(ctx.codeLayoutTimer); @@ -647,6 +696,8 @@ writeLLDMapFile(ctx); writeMapFile(ctx); + writePEChecksum(); + if (errorCount()) return;