Implement /driver, /driver:wdm and /driver:uponly

This patch implements /driver, /driver:wdm and /driver:uponly as
described in
https://docs.microsoft.com/en-us/cpp/build/reference/driver-windows-nt-kernel-mode-driver?view=vs-2019.

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

5 files changed, 134 insertions, 3 deletions

diff --git a/lld/COFF/Config.h b/lld/COFF/Config.h
index 07fa48b3ebc..2690ea5c408 100644
--- a/lld/COFF/Config.h
+++ b/lld/COFF/Config.h
@@ -103,6 +103,9 @@ struct Configuration {
   bool debugDwarf = false;
   bool debugGHashes = false;
   bool debugSymtab = false;
+  bool driver = false;
+  bool driverUponly = false;
+  bool driverWdm = false;
   bool showTiming = false;
   bool showSummary = false;
   unsigned debugTypes = static_cast<unsigned>(DebugType::None);
diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
index 29c3fe7922f..05d60428d26 100644
--- a/lld/COFF/Driver.cpp
+++ b/lld/COFF/Driver.cpp
@@ -102,9 +102,16 @@ static std::pair<StringRef, StringRef> getOldNewOptions(opt::InputArgList &args,
   return ret;
 }
 
-// Drop directory components and replace extension with ".exe" or ".dll".
+// Drop directory components and replace extension with
+// ".exe", ".dll" or ".sys".
 static std::string getOutputPath(StringRef path) {
-  return (sys::path::stem(path) + (config->dll ? ".dll" : ".exe")).str();
+  StringRef ext = ".exe";
+  if (config->dll)
+    ext = ".dll";
+  else if (config->driver)
+    ext = ".sys";
+
+  return (sys::path::stem(path) + ext).str();
 }
 
 // Returns true if S matches /crtend.?\.o$/.
@@ -1229,6 +1236,16 @@ void LinkerDriver::link(ArrayRef<const char *> argsArr) {
   // Handle /debugtype
   config->debugTypes = parseDebugTypes(args);
 
+  // Handle /driver[:uponly|:wdm].
+  config->driverUponly = args.hasArg(OPT_driver_uponly) ||
+                         args.hasArg(OPT_driver_uponly_wdm) ||
+                         args.hasArg(OPT_driver_wdm_uponly);
+  config->driverWdm = args.hasArg(OPT_driver_wdm) ||
+                      args.hasArg(OPT_driver_uponly_wdm) ||
+                      args.hasArg(OPT_driver_wdm_uponly);
+  config->driver =
+      config->driverUponly || config->driverWdm || args.hasArg(OPT_driver);
+
   // Handle /pdb
   bool shouldCreatePDB =
       (debug == DebugKind::Full || debug == DebugKind::GHash);
@@ -1703,6 +1720,9 @@ void LinkerDriver::link(ArrayRef<const char *> argsArr) {
       StringRef s = (config->machine == I386) ? "__DllMainCRTStartup@12"
                                               : "_DllMainCRTStartup";
       config->entry = addUndefined(s);
+    } else if (config->driverWdm) {
+      // /driver:wdm implies /entry:_NtProcessStartup
+      config->entry = addUndefined(mangle("_NtProcessStartup"));
     } else {
       // Windows specific -- If entry point name is not given, we need to
       // infer that from user-defined entry name.
diff --git a/lld/COFF/Options.td b/lld/COFF/Options.td
index 9e58d92e634..468e4da7d05 100644
--- a/lld/COFF/Options.td
+++ b/lld/COFF/Options.td
@@ -104,7 +104,13 @@ def debug : F<"debug">, HelpText<"Embed a symbol table in the image">;
 def debug_opt : P<"debug", "Embed a symbol table in the image with option">;
 def debugtype : P<"debugtype", "Debug Info Options">;
 def dll : F<"dll">, HelpText<"Create a DLL">;
-def driver : P<"driver", "Generate a Windows NT Kernel Mode Driver">;
+def driver : F<"driver">, HelpText<"Generate a Windows NT Kernel Mode Driver">;
+def driver_wdm : F<"driver:wdm">,
+    HelpText<"Set IMAGE_FILE_UP_SYSTEM_ONLY bit in PE header">;
+def driver_uponly : F<"driver:uponly">,
+    HelpText<"Set IMAGE_DLL_CHARACTERISTICS_WDM_DRIVER bit in PE header">;
+def driver_wdm_uponly : F<"driver:wdm,uponly">;
+def driver_uponly_wdm : F<"driver:uponly,wdm">;
 def nodefaultlib_all : F<"nodefaultlib">,
     HelpText<"Remove all default libraries">;
 def noentry : F<"noentry">,
diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
index 9638f4202e0..9d912d339bc 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -1304,6 +1304,8 @@ template <typename PEHeaderTy> void Writer::writeHeader() {
     coff->Characteristics |= IMAGE_FILE_32BIT_MACHINE;
   if (config->dll)
     coff->Characteristics |= IMAGE_FILE_DLL;
+  if (config->driverUponly)
+    coff->Characteristics |= IMAGE_FILE_UP_SYSTEM_ONLY;
   if (!config->relocatable)
     coff->Characteristics |= IMAGE_FILE_RELOCS_STRIPPED;
   if (config->swaprunCD)
@@ -1351,6 +1353,8 @@ template <typename PEHeaderTy> void Writer::writeHeader() {
   pe->SizeOfHeapCommit = config->heapCommit;
   if (config->appContainer)
     pe->DLLCharacteristics |= IMAGE_DLL_CHARACTERISTICS_APPCONTAINER;
+  if (config->driverWdm)
+    pe->DLLCharacteristics |= IMAGE_DLL_CHARACTERISTICS_WDM_DRIVER;
   if (config->dynamicBase)
     pe->DLLCharacteristics |= IMAGE_DLL_CHARACTERISTICS_DYNAMIC_BASE;
   if (config->highEntropyVA)
diff --git a/lld/test/COFF/driver-opt.s b/lld/test/COFF/driver-opt.s
new file mode 100644
index 00000000000..72459d6bb59
--- /dev/null
+++ b/lld/test/COFF/driver-opt.s
@@ -0,0 +1,98 @@
+# RUN: mkdir -p %t.dir
+# RUN: yaml2obj < %s > %t.dir/foo.obj
+
+# RUN: rm -f %t.dir/foo.sys
+# RUN: cd %t.dir; lld-link /driver foo.obj
+# RUN: llvm-readobj --file-headers %t.dir/foo.sys | FileCheck -check-prefix=DRIVER %s
+
+# DRIVER-NOT: IMAGE_FILE_UP_SYSTEM_ONLY
+# DRIVER-NOT: IMAGE_DLL_CHARACTERISTICS_WDM_DRIVER
+# DRIVER: AddressOfEntryPoint: 0x1000
+
+# RUN: rm -f %t.dir/foo.sys
+# RUN: cd %t.dir; lld-link /driver:uponly foo.obj
+# RUN: llvm-readobj --file-headers %t.dir/foo.sys | FileCheck -check-prefix=UPONLY %s
+
+# UPONLY: IMAGE_FILE_UP_SYSTEM_ONLY
+# UPONLY: AddressOfEntryPoint: 0x1000
+
+# RUN: rm -f %t.dir/foo.sys
+# RUN: cd %t.dir; lld-link /driver:wdm foo.obj
+# RUN: llvm-readobj --file-headers %t.dir/foo.sys | FileCheck -check-prefix=WDM %s
+
+# WDM: AddressOfEntryPoint: 0x1004
+# WDM: IMAGE_DLL_CHARACTERISTICS_WDM_DRIVER
+
+# RUN: rm -f %t.dir/foo.sys
+# RUN: cd %t.dir; lld-link /driver:wdm,uponly foo.obj
+# RUN: llvm-readobj --file-headers %t.dir/foo.sys | FileCheck -check-prefix=BOTH %s
+
+# RUN: rm -f %t.dir/foo.sys
+# RUN: cd %t.dir; lld-link /driver:uponly,wdm foo.obj
+# RUN: llvm-readobj --file-headers %t.dir/foo.sys | FileCheck -check-prefix=BOTH %s
+
+# BOTH: IMAGE_FILE_UP_SYSTEM_ONLY
+# BOTH: AddressOfEntryPoint: 0x1004
+# BOTH: IMAGE_DLL_CHARACTERISTICS_WDM_DRIVER
+
+# RUN: rm -f %t.dir/foo.sys
+# RUN: cd %t.dir; lld-link /driver foo.obj
+# RUN: llvm-readobj --file-headers %t.dir/foo.sys | FileCheck -check-prefix=FIXED1 %s
+
+# RUN: rm -f %t.dir/foo.sys
+# RUN: cd %t.dir; lld-link /driver foo.obj /fixed:no
+# RUN: llvm-readobj --file-headers %t.dir/foo.sys | FileCheck -check-prefix=FIXED1 %s
+
+# FIXED1: IMAGE_DLL_CHARACTERISTICS_DYNAMIC_BASE
+
+# RUN: rm -f %t.dir/foo.sys
+# RUN: cd %t.dir; lld-link /driver foo.obj /fixed
+# RUN: llvm-readobj --file-headers %t.dir/foo.sys | FileCheck -check-prefix=FIXED2 %s
+
+# FIXED2-NOT: IMAGE_DLL_CHARACTERISTICS_DYNAMIC_BASE
+
+--- !COFF
+header:
+  Machine:         IMAGE_FILE_MACHINE_AMD64
+  Characteristics: []
+sections:
+  - Name:            .text
+    Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ]
+    Alignment:       4096
+    SectionData:     0000000000000000
+    Relocations:
+      - VirtualAddress:  0
+        SymbolName:      __ImageBase
+        Type:            IMAGE_REL_AMD64_ADDR64
+symbols:
+  - Name:            .text
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+    SectionDefinition:
+      Length:          8
+      NumberOfRelocations: 1
+      NumberOfLinenumbers: 0
+      CheckSum:        0
+      Number:          0
+  - Name:            main
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+  - Name:            mainCRTStartup
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+  - Name:            _NtProcessStartup
+    Value:           4
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+...