[X86] Support -mstack-protector-guard-symbol

Reviewed By: nickdesaulniers

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

12 files changed, 111 insertions, 2 deletions

diff --git a/clang/docs/ClangCommandLineReference.rst b/clang/docs/ClangCommandLineReference.rst
index ae110df..129e70d 100644
--- a/clang/docs/ClangCommandLineReference.rst
+++ b/clang/docs/ClangCommandLineReference.rst
@@ -3285,6 +3285,10 @@ Use the given offset for addressing the stack-protector guard
 
 Use the given reg for addressing the stack-protector guard
 
+.. option:: -mstack-protector-guard-symbol=<arg>
+
+Use the given symbol for addressing the stack-protector guard
+
 .. option:: -mstack-protector-guard=<arg>
 
 Use the given guard (global, tls) for addressing the stack-protector guard
diff --git a/clang/docs/ReleaseNotes.rst b/clang/docs/ReleaseNotes.rst
index 27ba4f1..86b9b7e 100644
--- a/clang/docs/ReleaseNotes.rst
+++ b/clang/docs/ReleaseNotes.rst
@@ -548,6 +548,8 @@ X86 Support in Clang
 - Added the ``-m[no-]rdpru`` flag to enable/disable the RDPRU instruction
   provided by AMD Zen2 and later processors. Defined intrinsics for using
   this instruction (see rdpruintrin.h).
+- Support ``-mstack-protector-guard-symbol=[SymbolName]`` to use the given
+  symbol for addressing the stack protector guard.
 
 DWARF Support in Clang
 ----------------------
diff --git a/clang/include/clang/Basic/CodeGenOptions.h b/clang/include/clang/Basic/CodeGenOptions.h
index 23d76c3..cd204e5 100644
--- a/clang/include/clang/Basic/CodeGenOptions.h
+++ b/clang/include/clang/Basic/CodeGenOptions.h
@@ -389,6 +389,9 @@ public:
   /// On AArch64 this can only be "sp_el0".
   std::string StackProtectorGuardReg;
 
+  /// Specify a symbol to be the guard value.
+  std::string StackProtectorGuardSymbol;
+
   /// Path to ignorelist file specifying which objects
   /// (files, functions) listed for instrumentation by sanitizer
   /// coverage pass should actually not be instrumented.
diff --git a/clang/include/clang/Driver/Options.td b/clang/include/clang/Driver/Options.td
index 8ae9145..a1329e7 100644
--- a/clang/include/clang/Driver/Options.td
+++ b/clang/include/clang/Driver/Options.td
@@ -3771,6 +3771,9 @@ def mstack_protector_guard_EQ : Joined<["-"], "mstack-protector-guard=">, Group<
 def mstack_protector_guard_offset_EQ : Joined<["-"], "mstack-protector-guard-offset=">, Group<m_Group>, Flags<[CC1Option]>,
   HelpText<"Use the given offset for addressing the stack-protector guard">,
   MarshallingInfoInt<CodeGenOpts<"StackProtectorGuardOffset">, "INT_MAX", "int">;
+def mstack_protector_guard_symbol_EQ : Joined<["-"], "mstack-protector-guard-symbol=">, Group<m_Group>, Flags<[CC1Option]>,
+  HelpText<"Use the given symbol for addressing the stack-protector guard">,
+  MarshallingInfoString<CodeGenOpts<"StackProtectorGuardSymbol">>;
 def mstack_protector_guard_reg_EQ : Joined<["-"], "mstack-protector-guard-reg=">, Group<m_Group>, Flags<[CC1Option]>,
   HelpText<"Use the given reg for addressing the stack-protector guard">,
   MarshallingInfoString<CodeGenOpts<"StackProtectorGuardReg">>;
diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
index 56ed59d..cf954b7 100644
--- a/clang/lib/CodeGen/CodeGenModule.cpp
+++ b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -900,6 +900,9 @@ void CodeGenModule::Release() {
   if (!getCodeGenOpts().StackProtectorGuardReg.empty())
     getModule().setStackProtectorGuardReg(
         getCodeGenOpts().StackProtectorGuardReg);
+  if (!getCodeGenOpts().StackProtectorGuardSymbol.empty())
+    getModule().setStackProtectorGuardSymbol(
+        getCodeGenOpts().StackProtectorGuardSymbol);
   if (getCodeGenOpts().StackProtectorGuardOffset != INT_MAX)
     getModule().setStackProtectorGuardOffset(
         getCodeGenOpts().StackProtectorGuardOffset);
diff --git a/clang/lib/Driver/ToolChains/Clang.cpp b/clang/lib/Driver/ToolChains/Clang.cpp
index aae416f..ec1a31f 100644
--- a/clang/lib/Driver/ToolChains/Clang.cpp
+++ b/clang/lib/Driver/ToolChains/Clang.cpp
@@ -3231,6 +3231,16 @@ static void RenderAnalyzerOptions(const ArgList &Args, ArgStringList &CmdArgs,
   Args.AddAllArgValues(CmdArgs, options::OPT_Xanalyzer);
 }
 
+static bool isValidSymbolName(StringRef S) {
+  if (S.empty())
+    return false;
+
+  if (std::isdigit(S[0]))
+    return false;
+
+  return llvm::all_of(S, [](char C) { return std::isalnum(C) || C == '_'; });
+}
+
 static void RenderSSPOptions(const Driver &D, const ToolChain &TC,
                              const ArgList &Args, ArgStringList &CmdArgs,
                              bool KernelOrKext) {
@@ -3362,6 +3372,16 @@ static void RenderSSPOptions(const Driver &D, const ToolChain &TC,
     }
     A->render(Args, CmdArgs);
   }
+
+  if (Arg *A = Args.getLastArg(options::OPT_mstack_protector_guard_symbol_EQ)) {
+    StringRef Value = A->getValue();
+    if (!isValidSymbolName(Value)) {
+      D.Diag(diag::err_drv_argument_only_allowed_with)
+          << A->getOption().getName() << "legal symbol name";
+      return;
+    }
+    A->render(Args, CmdArgs);
+  }
 }
 
 static void RenderSCPOptions(const ToolChain &TC, const ArgList &Args,
diff --git a/clang/test/Driver/stack-protector-guard.c b/clang/test/Driver/stack-protector-guard.c
index 2c215b1..de8776b 100644
--- a/clang/test/Driver/stack-protector-guard.c
+++ b/clang/test/Driver/stack-protector-guard.c
@@ -13,6 +13,8 @@
 // RUN:   FileCheck -check-prefix=CHECK-FS %s
 // RUN: %clang -### -target x86_64-unknown-unknown -mstack-protector-guard-reg=gs %s 2>&1 | \
 // RUN:   FileCheck -check-prefix=CHECK-GS %s
+// RUN: %clang -### -target x86_64-unknown-unknown -mstack-protector-guard-symbol=sym %s 2>&1 | \
+// RUN:   FileCheck -check-prefix=CHECK-SYM %s
 
 // Invalid arch
 // RUN: not %clang -target powerpc64le-linux-gnu -mstack-protector-guard=tls %s 2>&1 | \
@@ -32,10 +34,16 @@
 // RUN:   FileCheck -check-prefix=INVALID-REG %s
 // RUN: not %clang -target x86_64-unknown-unknown -c -mstack-protector-guard-reg=ds %s 2>&1 | \
 // RUN:   FileCheck -check-prefix=INVALID-REG %s
+// RUN: not %clang -target x86_64-unknown-unknown -c -mstack-protector-guard-symbol=2s %s 2>&1 | \
+// RUN:   FileCheck -check-prefix=INVALID-SYM %s
+// RUN: not %clang -target x86_64-unknown-unknown -c -mstack-protector-guard-symbol= %s 2>&1 | \
+// RUN:   FileCheck -check-prefix=INVALID-SYM %s
 
 // CHECK-FS: "-cc1" {{.*}}"-mstack-protector-guard-reg=fs"
 // CHECK-GS: "-cc1" {{.*}}"-mstack-protector-guard-reg=gs"
 // INVALID-REG: error: invalid value {{.*}} in 'mstack-protector-guard-reg=', expected one of: fs gs
+// CHECK-SYM: "-cc1" {{.*}}"-mstack-protector-guard-symbol=sym"
+// INVALID-SYM: error: invalid argument 'mstack-protector-guard-symbol=' only allowed with 'legal symbol name'
 
 // RUN: not %clang -target arm-eabi-c -mstack-protector-guard=tls %s 2>&1 | \
 // RUN:   FileCheck -check-prefix=MISSING-OFFSET %s
diff --git a/llvm/docs/ReleaseNotes.rst b/llvm/docs/ReleaseNotes.rst
index 1268cc7..03d964e 100644
--- a/llvm/docs/ReleaseNotes.rst
+++ b/llvm/docs/ReleaseNotes.rst
@@ -221,6 +221,12 @@ Changes to the DAG infrastructure
 ---------------------------------
 
 
+Changes to the Metadata Info
+---------------------------------
+
+* Add Module Flags Metadata ``stack-protector-guard-symbol`` which specify a
+  symbol for addressing the stack-protector guard.
+
 Changes to the Debug Info
 ---------------------------------
 
diff --git a/llvm/include/llvm/IR/Module.h b/llvm/include/llvm/IR/Module.h
index fc2d609..24da08d 100644
--- a/llvm/include/llvm/IR/Module.h
+++ b/llvm/include/llvm/IR/Module.h
@@ -911,6 +911,10 @@ public:
   StringRef getStackProtectorGuardReg() const;
   void setStackProtectorGuardReg(StringRef Reg);
 
+  /// Get/set a symbol to use as the stack protector guard.
+  StringRef getStackProtectorGuardSymbol() const;
+  void setStackProtectorGuardSymbol(StringRef Symbol);
+
   /// Get/set what offset from the stack protector to use.
   int getStackProtectorGuardOffset() const;
   void setStackProtectorGuardOffset(int Offset);
diff --git a/llvm/lib/IR/Module.cpp b/llvm/lib/IR/Module.cpp
index 5cd74d5..b51ea45 100644
--- a/llvm/lib/IR/Module.cpp
+++ b/llvm/lib/IR/Module.cpp
@@ -714,6 +714,18 @@ void Module::setStackProtectorGuardReg(StringRef Reg) {
   addModuleFlag(ModFlagBehavior::Error, "stack-protector-guard-reg", ID);
 }
 
+StringRef Module::getStackProtectorGuardSymbol() const {
+  Metadata *MD = getModuleFlag("stack-protector-guard-symbol");
+  if (auto *MDS = dyn_cast_or_null<MDString>(MD))
+    return MDS->getString();
+  return {};
+}
+
+void Module::setStackProtectorGuardSymbol(StringRef Symbol) {
+  MDString *ID = MDString::get(getContext(), Symbol);
+  addModuleFlag(ModFlagBehavior::Error, "stack-protector-guard-symbol", ID);
+}
+
 int Module::getStackProtectorGuardOffset() const {
   Metadata *MD = getModuleFlag("stack-protector-guard-offset");
   if (auto *CI = mdconst::dyn_extract_or_null<ConstantInt>(MD))
diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 047ac0d..6dc823f 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -2845,6 +2845,21 @@ Value *X86TargetLowering::getIRStackGuard(IRBuilderBase &IRB) const {
         AddressSpace = X86AS::FS;
       else if (GuardReg == "gs")
         AddressSpace = X86AS::GS;
+
+      // Use symbol guard if user specify.
+      StringRef GuardSymb = M->getStackProtectorGuardSymbol();
+      if (!GuardSymb.empty()) {
+        GlobalVariable *GV = M->getGlobalVariable(GuardSymb);
+        if (!GV) {
+          Type *Ty = Subtarget.is64Bit() ? Type::getInt64Ty(M->getContext())
+                                         : Type::getInt32Ty(M->getContext());
+          GV = new GlobalVariable(*M, Ty, false, GlobalValue::ExternalLinkage,
+                                  nullptr, GuardSymb, nullptr,
+                                  GlobalValue::NotThreadLocal, AddressSpace);
+        }
+        return GV;
+      }
+
       return SegmentOffset(IRB, Offset, AddressSpace);
     }
   }
diff --git a/llvm/test/CodeGen/X86/stack-protector-3.ll b/llvm/test/CodeGen/X86/stack-protector-3.ll
index 82e1157..59f583b 100644
--- a/llvm/test/CodeGen/X86/stack-protector-3.ll
+++ b/llvm/test/CodeGen/X86/stack-protector-3.ll
@@ -6,6 +6,8 @@
 ; RUN: cat %t/main.ll %t/e.ll > %t/e2.ll
 ; RUN: cat %t/main.ll %t/f.ll > %t/f2.ll
 ; RUN: cat %t/main.ll %t/g.ll > %t/g2.ll
+; RUN: cat %t/main.ll %t/h.ll > %t/h2.ll
+; RUN: cat %t/existedGV.ll %t/main.ll %t/h.ll > %t/i2.ll
 ; RUN: llc -mtriple=x86_64-pc-linux-gnu -o - < %t/a2.ll | FileCheck --check-prefix=CHECK-TLS-FS-40 %s
 ; RUN: llc -mtriple=x86_64-pc-linux-gnu -o - < %t/b2.ll | FileCheck --check-prefix=CHECK-TLS-FS-40 %s
 ; RUN: llc -mtriple=x86_64-pc-linux-gnu -o - < %t/c2.ll | FileCheck --check-prefix=CHECK-GLOBAL %s
@@ -13,8 +15,8 @@
 ; RUN: llc -mtriple=x86_64-pc-linux-gnu -o - < %t/e2.ll | FileCheck --check-prefix=CHECK-GS %s
 ; RUN: llc -mtriple=x86_64-pc-linux-gnu -o - < %t/f2.ll | FileCheck --check-prefix=CHECK-OFFSET %s
 ; RUN: llc -mtriple=x86_64-pc-linux-gnu -o - < %t/g2.ll | FileCheck --check-prefix=CHECK-NEGATIVE-OFFSET %s
-
-;--- main.ll
+; RUN: llc -mtriple=x86_64-pc-linux-gnu -o - < %t/h2.ll | FileCheck --check-prefix=CHECK-SYM %s
+; RUN: llc -mtriple=x86_64-pc-linux-gnu -o - < %t/i2.ll | FileCheck --check-prefix=CHECK-SYMGV %s
 
 ; CHECK-TLS-FS-40:       movq    %fs:40, %rax
 ; CHECK-TLS-FS-40:       movq    %fs:40, %rax
@@ -57,7 +59,31 @@
 ; CHECK-GLOBAL-NEXT:  .cfi_def_cfa_offset 32
 ; CHECK-GLOBAL-NEXT:  callq   __stack_chk_fail
 
+; CHECK-SYM:         movq    __woof@GOTPCREL(%rip), %rax
+; CHECK-SYM-NEXT:    movq    %fs:(%rax), %rcx
+; CHECK-SYM-NEXT:    movq    %rcx, 16(%rsp)
+; CHECK-SYM:         movq    %fs:(%rax), %rax
+; CHECK-SYM-NEXT:    cmpq    16(%rsp), %rax
+; CHECK-SYM-NEXT:    jne     .LBB0_2
+; CHECK-SYM:         .LBB0_2:
+; CHECK-SYM-NEXT:    .cfi_def_cfa_offset 32
+; CHECK-SYM-NEXT:    callq   __stack_chk_fai
+
+; CHECK-SYMGV:       movq    __woof(%rip), %rax
+; CHECK-SYMGV-NEXT:  movq    %rax, 16(%rsp)
+; CHECK-SYMGV:       cmpq    16(%rsp), %rax
+; CHECK-SYMGV-NEXT:  jne     .LBB0_2
+; CHECK-SYMGV:       .LBB0_2:
+; CHECK-SYMGV-NEXT:  .cfi_def_cfa_offset 32
+; CHECK-SYMGV-NEXT:  callq   __stack_chk_fail
+
 ; ModuleID = 't.c'
+;--- existedGV.ll
+
+@__woof = dso_local local_unnamed_addr global ptr null, align 8
+
+;--- main.ll
+
 @.str = private unnamed_addr constant [14 x i8] c"stackoverflow\00", align 1
 @a = dso_local local_unnamed_addr global ptr null, align 8
 
@@ -104,3 +130,6 @@ attributes #2 = { nounwind }
 ;--- g.ll
 !llvm.module.flags = !{!1}
 !1 = !{i32 2, !"stack-protector-guard-offset", i32 -20}
+;--- h.ll
+!llvm.module.flags = !{!1}
+!1 = !{i32 2, !"stack-protector-guard-symbol", !"__woof"}