[lld] 3f55853 - [LLD][Windows]Feature "checksum" for Windows PE

[via llvm-commits]

Author: Qfrost
Date: 2023-01-02T17:20:15+08:00
New Revision: 3f55853edf93e1bb2e3c9b38fedfb7b8837464dc

URL: https://github.com/llvm/llvm-project/commit/3f55853edf93e1bb2e3c9b38fedfb7b8837464dc
DIFF: https://github.com/llvm/llvm-project/commit/3f55853edf93e1bb2e3c9b38fedfb7b8837464dc.diff

LOG: [LLD][Windows]Feature "checksum" for Windows PE

Differential Revision: https://reviews.llvm.org/D139184

Added: 
    lld/test/COFF/checksum.test

Modified: 
    lld/COFF/Config.h
    lld/COFF/Driver.cpp
    lld/COFF/Options.td
    lld/COFF/Writer.cpp

Removed: 
    


################################################################################
diff  --git a/lld/COFF/Config.h b/lld/COFF/Config.h
index b49b21e773519..c7f10f253b8a6 100644
--- a/lld/COFF/Config.h
+++ b/lld/COFF/Config.h
@@ -289,6 +289,7 @@ struct Configuration {
   bool autoImport = false;
   bool pseudoRelocs = false;
   bool stdcallFixup = false;
+  bool writeCheckSum = false;
 };
 
 extern std::unique_ptr<Configuration> config;

diff  --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
index 2c2b0dbad822b..8c73e25d56305 100644
--- a/lld/COFF/Driver.cpp
+++ b/lld/COFF/Driver.cpp
@@ -415,6 +415,9 @@ void LinkerDriver::parseDirectives(InputFile *file) {
     case OPT_nodefaultlib:
       config->noDefaultLibs.insert(doFindLib(arg->getValue()).lower());
       break;
+    case OPT_release:
+      config->writeCheckSum = true;
+      break;
     case OPT_section:
       parseSection(arg->getValue());
       break;
@@ -2021,6 +2024,10 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   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
index fb4c65ef074b0..df527bb48ae72 100644
--- a/lld/COFF/Options.td
+++ b/lld/COFF/Options.td
@@ -161,6 +161,8 @@ def swaprun_net : F<"swaprun:net">, Alias<swaprun>, AliasArgs<["net"]>,
 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
index 9875eee45eb94..ccadc26d27e83 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -229,6 +229,7 @@ class Writer {
   void setSectionPermissions();
   void writeSections();
   void writeBuildId();
+  void writePEChecksum();
   void sortSections();
   void sortExceptionTable();
   void sortCRTSectionChunks(std::vector<Chunk *> &chunks);
@@ -598,6 +599,43 @@ void Writer::finalizeAddresses() {
   }
 }
 
+void Writer::writePEChecksum() {
+  if (!config->writeCheckSum) {
+    return;
+  }
+
+  // https://docs.microsoft.com/en-us/windows/win32/debug/pe-format#checksum
+  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));
+
+  uint64_t sum = 0;
+  uint32_t count = size;
+  ulittle16_t *addr = (ulittle16_t *)buf;
+
+  // The PE checksum algorithm, implemented as suggested in RFC1071
+  while (count > 1) {
+    sum += *addr++;
+    count -= 2;
+  }
+
+  // Add left-over byte, if any
+  if (count > 0)
+    sum += *(unsigned char *)addr;
+
+  // Fold 32-bit sum to 16 bits
+  while (sum >> 16) {
+    sum = (sum & 0xffff) + (sum >> 16);
+  }
+
+  sum += size;
+  peHeader->CheckSum = sum;
+}
+
 // The main function of the writer.
 void Writer::run() {
   ScopedTimer t1(ctx.codeLayoutTimer);
@@ -646,6 +684,8 @@ void Writer::run() {
   writeLLDMapFile(ctx);
   writeMapFile(ctx);
 
+  writePEChecksum();
+
   if (errorCount())
     return;
 

diff  --git a/lld/test/COFF/checksum.test b/lld/test/COFF/checksum.test
new file mode 100644
index 0000000000000..6e67f78382a32
--- /dev/null
+++ b/lld/test/COFF/checksum.test
@@ -0,0 +1,11 @@
+# RUN: yaml2obj %p/Inputs/hello32.yaml -o %t.obj
+# RUN: lld-link /out:%t.exe /entry:main /timestamp:0 %p/Inputs/std32.lib %t.obj
+# RUN: llvm-readobj --file-headers %t.exe | FileCheck -check-prefix=CHECKSUM1 %s
+
+CHECKSUM1: CheckSum: 0x0
+
+# RUN: yaml2obj %p/Inputs/hello32.yaml -o %t.obj
+# RUN: lld-link /out:%t.exe /entry:main /timestamp:0 /RELEASE %p/Inputs/std32.lib %t.obj
+# RUN: llvm-readobj --file-headers %t.exe | FileCheck -check-prefix=CHECKSUM2 %s
+
+CHECKSUM2: CheckSum: 0x9196