[clang][RISCV] Introduce preprocessor macro when Zicfiss-based shadow stack is enabled (#127592)

The `-fcf-protection=[full|return]` flag enables shadow stack
implementation based on RISC-V Zicfiss extension. This patch adds the
`__riscv_shadow_stack` predefined macro to preprocessing when such a
shadow stack implementation is enabled.

4 files changed, 54 insertions, 1 deletions

diff --git a/clang/include/clang/Basic/LangOptions.def b/clang/include/clang/Basic/LangOptions.def
index bfab0ba..383440d 100644
--- a/clang/include/clang/Basic/LangOptions.def
+++ b/clang/include/clang/Basic/LangOptions.def
@@ -365,6 +365,7 @@ LANGOPT(ObjCDisableDirectMethodsForTesting, 1, 0,
 LANGOPT(CFProtectionBranch , 1, 0, "Control-Flow Branch Protection enabled")
 ENUM_LANGOPT(CFBranchLabelScheme, CFBranchLabelSchemeKind, 2, CFBranchLabelSchemeKind::Default,
              "Control-Flow Branch Protection Label Scheme")
+LANGOPT(CFProtectionReturn, 1, 0, "Control-Flow Return Protection enabled")
 LANGOPT(FakeAddressSpaceMap , 1, 0, "OpenCL fake address space map")
 ENUM_LANGOPT(AddressSpaceMapMangling , AddrSpaceMapMangling, 2, ASMM_Target, "OpenCL address space map mangling mode")
 LANGOPT(IncludeDefaultHeader, 1, 0, "Include default header file for OpenCL")
diff --git a/clang/lib/Basic/Targets/RISCV.cpp b/clang/lib/Basic/Targets/RISCV.cpp
index b4aa320..dff990d 100644
--- a/clang/lib/Basic/Targets/RISCV.cpp
+++ b/clang/lib/Basic/Targets/RISCV.cpp
@@ -238,6 +238,9 @@ void RISCVTargetInfo::getTargetDefines(const LangOptions &Opts,
     else
       Builder.defineMacro("__riscv_32e");
   }
+
+  if (Opts.CFProtectionReturn && ISAInfo->hasExtension("zicfiss"))
+    Builder.defineMacro("__riscv_shadow_stack");
 }
 
 static constexpr int NumRVVBuiltins =
diff --git a/clang/lib/Frontend/CompilerInvocation.cpp b/clang/lib/Frontend/CompilerInvocation.cpp
index 014e629..b9a5c05 100644
--- a/clang/lib/Frontend/CompilerInvocation.cpp
+++ b/clang/lib/Frontend/CompilerInvocation.cpp
@@ -4048,8 +4048,13 @@ bool CompilerInvocation::ParseLangArgs(LangOptions &Opts, ArgList &Args,
 
   if (const Arg *A = Args.getLastArg(OPT_fcf_protection_EQ)) {
     StringRef Name = A->getValue();
-    if (Name == "full" || Name == "branch") {
+    if (Name == "full") {
+      Opts.CFProtectionBranch = 1;
+      Opts.CFProtectionReturn = 1;
+    } else if (Name == "branch") {
       Opts.CFProtectionBranch = 1;
+    } else if (Name == "return") {
+      Opts.CFProtectionReturn = 1;
     }
   }
 
diff --git a/clang/test/Preprocessor/riscv-cf-protection-return.c b/clang/test/Preprocessor/riscv-cf-protection-return.c
new file mode 100644
index 0000000..3a93a88f
--- /dev/null
+++ b/clang/test/Preprocessor/riscv-cf-protection-return.c
@@ -0,0 +1,44 @@
+// RUN: %clang --target=riscv32 -E -dM %s -o - | \
+// RUN: FileCheck --check-prefixes=NO-MACRO %s
+
+// RUN: %clang --target=riscv32 -fcf-protection=return -E -dM %s -o - | \
+// RUN: FileCheck --check-prefixes=NO-MACRO %s
+
+// RUN: %clang --target=riscv32 -fcf-protection=full -E -dM %s -o - | \
+// RUN: FileCheck --check-prefixes=NO-MACRO %s
+
+// RUN: %clang --target=riscv32 -march=rv32i_zicfiss1p0 \
+// RUN: -menable-experimental-extensions -E -dM %s -o - | \
+// RUN: FileCheck --check-prefixes=NO-MACRO %s
+
+// RUN: %clang --target=riscv32 -march=rv32i_zicfiss1p0 \
+// RUN: -menable-experimental-extensions -fcf-protection=return -E -dM %s \
+// RUN: -o - | FileCheck --check-prefixes=SHSTK-MACRO %s
+
+// RUN: %clang --target=riscv32 -march=rv32i_zicfiss1p0 \
+// RUN: -menable-experimental-extensions -fcf-protection=full -E -dM %s -o - \
+// RUN: | FileCheck --check-prefixes=SHSTK-MACRO %s
+
+// RUN: %clang --target=riscv64 -E -dM %s -o - | \
+// RUN: FileCheck --check-prefixes=NO-MACRO %s
+
+// RUN: %clang --target=riscv64 -fcf-protection=return -E -dM %s -o - | \
+// RUN: FileCheck --check-prefixes=NO-MACRO %s
+
+// RUN: %clang --target=riscv64 -fcf-protection=full -E -dM %s -o - | \
+// RUN: FileCheck --check-prefixes=NO-MACRO %s
+
+// RUN: %clang --target=riscv64 -march=rv64i_zicfiss1p0 \
+// RUN: -menable-experimental-extensions -E -dM %s -o - | \
+// RUN: FileCheck --check-prefixes=NO-MACRO %s
+
+// RUN: %clang --target=riscv64 -march=rv64i_zicfiss1p0 \
+// RUN: -menable-experimental-extensions -fcf-protection=return -E -dM %s \
+// RUN: -o - | FileCheck --check-prefixes=SHSTK-MACRO %s
+
+// RUN: %clang --target=riscv64 -march=rv64i_zicfiss1p0 \
+// RUN: -menable-experimental-extensions -fcf-protection=full -E -dM %s -o - \
+// RUN: | FileCheck --check-prefixes=SHSTK-MACRO %s
+
+// SHSTK-MACRO: __riscv_shadow_stack 1{{$}}
+// NO-MACRO-NOT: __riscv_shadow_stack