[AArch64] Branch Protection and Return Address Signing B Key Support

- Add support for -mbranch-protection=<type>[+<type>]* where
  - <type> ::= [standard, none, bti, pac-ret[+b-key,+leaf]*]
- The protection emits relevant function attributes
  - sign-return-address=<scope>
  - sign-return-address-key=<key>
  - branch-protection

llvm-svn: 345273

clang/include/clang/Basic/DiagnosticDriverKinds.td

@@ -279,6 +279,8 @@ def warn_drv_disabling_vptr_no_rtti_default : Warning<
 def warn_drv_object_size_disabled_O0 : Warning<
   "the object size sanitizer has no effect at -O0, but is explicitly enabled: %0">,
   InGroup<InvalidCommandLineArgument>, DefaultWarnNoWerror;
+def err_invalid_branch_protection: Error <
+  "invalid branch protection option '%0' in '%1'">;
 
 def note_drv_command_failed_diag_msg : Note<
   "diagnostic msg: %0">;

clang/include/clang/Driver/CC1Options.td

@@ -354,6 +354,12 @@ def fdebug_pass_manager : Flag<["-"], "fdebug-pass-manager">,
     HelpText<"Prints debug information for the new pass manager">;
 def fno_debug_pass_manager : Flag<["-"], "fno-debug-pass-manager">,
     HelpText<"Disables debug printing for the new pass manager">;
+// The driver option takes the key as a parameter to the -msign-return-address=
+// and -mbranch-protection= options, but CC1 has a separate option so we
+// don't have to parse the parameter twice.
+def msign_return_address_key_EQ : Joined<["-"], "msign-return-address-key=">,
+    Values<"a_key,b_key">;
+def mbranch_target_enforce : Flag<["-"], "mbranch-target-enforce">;
 
 //===----------------------------------------------------------------------===//
 // Dependency Output Options

clang/include/clang/Driver/Options.td

@@ -2071,9 +2071,11 @@ foreach i = {8-15,18} in
   def fcall_saved_x#i : Flag<["-"], "fcall-saved-x"#i>, Group<m_aarch64_Features_Group>,
     HelpText<"Make the x"#i#" register call-saved (AArch64 only)">;
 
-def msign_return_address : Joined<["-"], "msign-return-address=">,
-  Flags<[CC1Option]>, Group<m_Group>,
-  HelpText<"Select return address signing scope">, Values<"none,all,non-leaf">;
+def msign_return_address_EQ : Joined<["-"], "msign-return-address=">,
+  Flags<[CC1Option]>, Group<m_Group>, Values<"none,all,non-leaf">,
+  HelpText<"Select return address signing scope">;
+def mbranch_protection_EQ : Joined<["-"], "mbranch-protection=">,
+  HelpText<"Enforce targets of indirect branches and function returns">;
 
 def msimd128 : Flag<["-"], "msimd128">, Group<m_wasm_Features_Group>;
 def mno_simd128 : Flag<["-"], "mno-simd128">, Group<m_wasm_Features_Group>;

clang/include/clang/Frontend/CodeGenOptions.def

@@ -343,6 +343,8 @@ CODEGENOPT(ForceEmitVTables, 1, 0)
 CODEGENOPT(Addrsig, 1, 0)
 
 ENUM_CODEGENOPT(SignReturnAddress, SignReturnAddressScope, 2, None)
+ENUM_CODEGENOPT(SignReturnAddressKey, SignReturnAddressKeyValue, 1, AKey)
+CODEGENOPT(BranchTargetEnforcement, 1, 0)
 
 /// Whether to emit unused static constants.
 CODEGENOPT(KeepStaticConsts, 1, 0)

clang/include/clang/Frontend/CodeGenOptions.h

@@ -114,6 +114,8 @@ public:
     All      // Sign the return address of all functions
   };
 
+  enum SignReturnAddressKeyValue { AKey, BKey };
+
   /// The code model to use (-mcmodel).
   std::string CodeModel;
 

clang/lib/CodeGen/CGDeclCXX.cpp

@@ -360,11 +360,21 @@ llvm::Function *CodeGenModule::CreateGlobalInitOrDestructFunction(
     Fn->addFnAttr(llvm::Attribute::ShadowCallStack);
 
   auto RASignKind = getCodeGenOpts().getSignReturnAddress();
-  if (RASignKind != CodeGenOptions::SignReturnAddressScope::None)
+  if (RASignKind != CodeGenOptions::SignReturnAddressScope::None) {
     Fn->addFnAttr("sign-return-address",
                   RASignKind == CodeGenOptions::SignReturnAddressScope::All
                       ? "all"
                       : "non-leaf");
+    auto RASignKey = getCodeGenOpts().getSignReturnAddressKey();
+    Fn->addFnAttr("sign-return-address-key",
+                  RASignKey == CodeGenOptions::SignReturnAddressKeyValue::AKey
+                      ? "a_key"
+                      : "b_key");
+  }
+
+  if (getCodeGenOpts().BranchTargetEnforcement)
+    Fn->addFnAttr("branch-target-enforcement");
+
   return Fn;
 }
 

clang/lib/CodeGen/TargetInfo.cpp

@@ -4978,13 +4978,21 @@ public:
     llvm::Function *Fn = cast<llvm::Function>(GV);
 
     auto Kind = CGM.getCodeGenOpts().getSignReturnAddress();
-    if (Kind == CodeGenOptions::SignReturnAddressScope::None)
-      return;
+    if (Kind != CodeGenOptions::SignReturnAddressScope::None) {
+      Fn->addFnAttr("sign-return-address",
+                    Kind == CodeGenOptions::SignReturnAddressScope::All
+                        ? "all"
+                        : "non-leaf");
 
-    Fn->addFnAttr("sign-return-address",
-                  Kind == CodeGenOptions::SignReturnAddressScope::All
-                      ? "all"
-                      : "non-leaf");
+      auto Key = CGM.getCodeGenOpts().getSignReturnAddressKey();
+      Fn->addFnAttr("sign-return-address-key",
+                    Key == CodeGenOptions::SignReturnAddressKeyValue::AKey
+                        ? "a_key"
+                        : "b_key");
+    }
+
+    if (CGM.getCodeGenOpts().BranchTargetEnforcement)
+      Fn->addFnAttr("branch-target-enforcement");
   }
 };
 

clang/lib/Driver/ToolChains/Clang.cpp

@@ -1430,6 +1430,56 @@ void Clang::RenderTargetOptions(const llvm::Triple &EffectiveTriple,
   }
 }
 
+// Parse -mbranch-protection=<protection>[+<protection>]* where
+//   <protection> ::= standard | none | [bti,pac-ret[+b-key,+leaf]*]
+// Returns a triple of (return address signing Scope, signing key, require
+// landing pads)
+static std::tuple<StringRef, StringRef, bool>
+ParseAArch64BranchProtection(const Driver &D, const ArgList &Args,
+                             const Arg *A) {
+  StringRef Scope = "none";
+  StringRef Key = "a_key";
+  bool IndirectBranches = false;
+
+  StringRef Value = A->getValue();
+  // This maps onto -mbranch-protection=<scope>+<key>
+
+  if (Value.equals("standard")) {
+    Scope = "non-leaf";
+    Key = "a_key";
+    IndirectBranches = true;
+
+  } else if (!Value.equals("none")) {
+    SmallVector<StringRef, 4> BranchProtection;
+    StringRef(A->getValue()).split(BranchProtection, '+');
+
+    auto Protection = BranchProtection.begin();
+    while (Protection != BranchProtection.end()) {
+      if (Protection->equals("bti"))
+        IndirectBranches = true;
+      else if (Protection->equals("pac-ret")) {
+        Scope = "non-leaf";
+        while (++Protection != BranchProtection.end()) {
+          // Inner loop as "leaf" and "b-key" options must only appear attached
+          // to pac-ret.
+          if (Protection->equals("leaf"))
+            Scope = "all";
+          else if (Protection->equals("b-key"))
+            Key = "b_key";
+          else
+            break;
+        }
+        Protection--;
+      } else
+        D.Diag(diag::err_invalid_branch_protection)
+            << *Protection << A->getAsString(Args);
+      Protection++;
+    }
+  }
+
+  return std::make_tuple(Scope, Key, IndirectBranches);
+}
+
 namespace {
 void RenderAArch64ABI(const llvm::Triple &Triple, const ArgList &Args,
                       ArgStringList &CmdArgs) {
@@ -1484,9 +1534,33 @@ void Clang::AddAArch64TargetArgs(const ArgList &Args,
       CmdArgs.push_back("-aarch64-enable-global-merge=true");
   }
 
-  if (Arg *A = Args.getLastArg(options::OPT_msign_return_address)) {
+  // Enable/disable return address signing and indirect branch targets.
+  if (Arg *A = Args.getLastArg(options::OPT_msign_return_address_EQ,
+                               options::OPT_mbranch_protection_EQ)) {
+
+    const Driver &D = getToolChain().getDriver();
+
+    StringRef Scope, Key;
+    bool IndirectBranches;
+
+    if (A->getOption().matches(options::OPT_msign_return_address_EQ)) {
+      Scope = A->getValue();
+      if (!Scope.equals("none") && !Scope.equals("non-leaf") &&
+          !Scope.equals("all"))
+        D.Diag(diag::err_invalid_branch_protection)
+            << Scope << A->getAsString(Args);
+      Key = "a_key";
+      IndirectBranches = false;
+    } else
+      std::tie(Scope, Key, IndirectBranches) =
+          ParseAArch64BranchProtection(D, Args, A);
+
     CmdArgs.push_back(
-        Args.MakeArgString(Twine("-msign-return-address=") + A->getValue()));
+        Args.MakeArgString(Twine("-msign-return-address=") + Scope));
+    CmdArgs.push_back(
+        Args.MakeArgString(Twine("-msign-return-address-key=") + Key));
+    if (IndirectBranches)
+      CmdArgs.push_back("-mbranch-target-enforce");
   }
 }
 

clang/lib/Frontend/CompilerInvocation.cpp

@@ -1138,8 +1138,9 @@ static bool ParseCodeGenArgs(CodeGenOptions &Opts, ArgList &Args, InputKind IK,
 
   Opts.Addrsig = Args.hasArg(OPT_faddrsig);
 
-  if (Arg *A = Args.getLastArg(OPT_msign_return_address)) {
+  if (Arg *A = Args.getLastArg(OPT_msign_return_address_EQ)) {
     StringRef SignScope = A->getValue();
+
     if (SignScope.equals_lower("none"))
       Opts.setSignReturnAddress(CodeGenOptions::SignReturnAddressScope::None);
     else if (SignScope.equals_lower("all"))
@@ -1149,9 +1150,26 @@ static bool ParseCodeGenArgs(CodeGenOptions &Opts, ArgList &Args, InputKind IK,
           CodeGenOptions::SignReturnAddressScope::NonLeaf);
     else
       Diags.Report(diag::err_drv_invalid_value)
-          << A->getAsString(Args) << A->getValue();
+          << A->getAsString(Args) << SignScope;
+
+    if (Arg *A = Args.getLastArg(OPT_msign_return_address_key_EQ)) {
+      StringRef SignKey = A->getValue();
+      if (!SignScope.empty() && !SignKey.empty()) {
+        if (SignKey.equals_lower("a_key"))
+          Opts.setSignReturnAddressKey(
+              CodeGenOptions::SignReturnAddressKeyValue::AKey);
+        else if (SignKey.equals_lower("b_key"))
+          Opts.setSignReturnAddressKey(
+              CodeGenOptions::SignReturnAddressKeyValue::BKey);
+        else
+          Diags.Report(diag::err_drv_invalid_value)
+              << A->getAsString(Args) << SignKey;
+      }
+    }
   }
 
+  Opts.BranchTargetEnforcement = Args.hasArg(OPT_mbranch_target_enforce);
+
   Opts.KeepStaticConsts = Args.hasArg(OPT_fkeep_static_consts);
 
   Opts.SpeculativeLoadHardening = Args.hasArg(OPT_mspeculative_load_hardening);

clang/test/CodeGen/aarch64-sign-return-address.c

@@ -1,14 +1,27 @@
-// RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -msign-return-address=none  %s | FileCheck %s --check-prefix=CHECK-NONE
-// RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -msign-return-address=non-leaf %s | FileCheck %s --check-prefix=CHECK-PARTIAL
-// RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -msign-return-address=all %s | FileCheck %s --check-prefix=CHECK-ALL
+// RUN: %clang -target aarch64-arm-none-eabi -march=armv8.3-a -S -emit-llvm -o - -msign-return-address=none  %s | FileCheck %s --check-prefix=CHECK --check-prefix=NONE
+// RUN: %clang -target aarch64-arm-none-eabi -march=armv8.2-a -S -emit-llvm -o - -msign-return-address=all  %s | FileCheck %s --check-prefix=CHECK --check-prefix=ALL --check-prefix=A-KEY
+// RUN: %clang -target aarch64-arm-none-eabi -march=armv8.3-a -S -emit-llvm -o - -msign-return-address=all %s   | FileCheck %s --check-prefix=CHECK --check-prefix=ALL --check-prefix=A-KEY
+// RUN: %clang -target aarch64-arm-none-eabi -march=armv8.3-a -S -emit-llvm -o - -msign-return-address=non-leaf %s | FileCheck %s --check-prefix=CHECK --check-prefix=PARTIAL --check-prefix=A-KEY
+// RUN: %clang -target aarch64-arm-none-eabi -march=armv8.3-a -S -emit-llvm -o - -msign-return-address=all %s | FileCheck %s --check-prefix=CHECK --check-prefix=ALL --check-prefix=A-KEY
+// RUN: %clang -target aarch64-arm-none-eabi -march=armv8.4-a -S -emit-llvm -o - -msign-return-address=all %s | FileCheck %s --check-prefix=CHECK --check-prefix=ALL --check-prefix=A-KEY
+// RUN: %clang -target aarch64-arm-none-eabi -march=armv8.3-a -S -emit-llvm -o - -mbranch-protection=pac-ret+b-key %s   | FileCheck %s --check-prefix=CHECK --check-prefix=PARTIAL --check-prefix=B-KEY
+// RUN: %clang -target aarch64-arm-none-eabi -march=armv8.3-a -S -emit-llvm -o - -mbranch-protection=pac-ret+b-key+leaf %s   | FileCheck %s --check-prefix=CHECK --check-prefix=ALL --check-prefix=B-KEY
+// RUN: %clang -target aarch64-arm-none-eabi -march=armv8.3-a -S -emit-llvm -o - -mbranch-protection=bti %s | FileCheck %s --check-prefix=CHECK --check-prefix=BTE
 
-// CHECK-NONE: @foo() #[[ATTR:[0-9]*]]
-// CHECK-NONE-NOT: attributes #[[ATTR]] = { {{.*}} "sign-return-address"={{.*}} }}
+// REQUIRES: aarch64-registered-target
 
-// CHECK-PARTIAL: @foo() #[[ATTR:[0-9]*]]
-// CHECK-PARTIAL: attributes #[[ATTR]] = { {{.*}} "sign-return-address"="non-leaf" {{.*}}}
+// CHECK: @foo() #[[ATTR:[0-9]*]]
+//
+// NONE-NOT: "sign-return-address"={{.*}}
 
-// CHECK-ALL: @foo() #[[ATTR:[0-9]*]]
-// CHECK-ALL: attributes #[[ATTR]] = { {{.*}} "sign-return-address"="all" {{.*}} }
+// PARTIAL: "sign-return-address"="non-leaf"
+
+// ALL: "sign-return-address"="all"
+
+// BTE: "branch-target-enforcement"
+
+// A-KEY: "sign-return-address-key"="a_key"
+
+// B-KEY: "sign-return-address-key"="b_key"
 
 void foo() {}

clang/test/CodeGenCXX/aarch64-sign-return-address-static-ctor.cpp

@@ -1,9 +1,26 @@
 // RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -msign-return-address=none  %s | \
 // RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-NONE
 // RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -msign-return-address=non-leaf %s | \
-// RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-PARTIAL
+// RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-PARTIAL  --check-prefix=CHECK-A-KEY
 // RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -msign-return-address=all %s | \
-// RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-ALL
+// RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-ALL  --check-prefix=CHECK-A-KEY
+
+// RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -mbranch-protection=none %s | \
+// RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-NONE
+// RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -mbranch-protection=standard %s | \
+// RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-PARTIAL  --check-prefix=CHECK-A-KEY  --check-prefix=CHECK-BTE
+// RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -mbranch-protection=pac-ret %s | \
+// RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-PARTIAL  --check-prefix=CHECK-A-KEY
+// RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -mbranch-protection=pac-ret+leaf %s | \
+// RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-ALL  --check-prefix=CHECK-A-KEY
+// RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -mbranch-protection=pac-ret+b-key %s | \
+// RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-PARTIAL  --check-prefix=CHECK-B-KEY
+// RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -mbranch-protection=pac-ret+b-key+leaf %s | \
+// RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-ALL  --check-prefix=CHECK-B-KEY
+// RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -mbranch-protection=bti %s | \
+// RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-BTE
+// RUN: %clang -target aarch64-arm-none-eabi -S -emit-llvm -o - -mbranch-protection=pac-ret+b-key+leaf+bti %s | \
+// RUN: FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-ALL  --check-prefix=CHECK-B-KEY --check-prefix=BTE
 
 struct Foo {
   Foo() {}
@@ -16,6 +33,9 @@ Foo f;
 
 // CHECK: @[[CTOR_FN]]() #[[ATTR:[0-9]*]]
 
-// CHECK-NONE-NOT: attributes #[[ATTR]] = { {{.*}} "sign-return-address"={{.*}} }}
-// CHECK-PARTIAL: attributes #[[ATTR]] = { {{.*}} "sign-return-address"="non-leaf" {{.*}}}
-// CHECK-ALL: attributes #[[ATTR]] = { {{.*}} "sign-return-address"="all" {{.*}} }
+// CHECK-NONE-NOT: "sign-return-address"={{.*}}
+// CHECK-PARTIAL: "sign-return-address"="non-leaf"
+// CHECK-ALL: "sign-return-address"="all"
+// CHECK-A-KEY: "sign-return-address-key"="a_key"
+// CHECK-B-KEY: "sign-return-address-key"="b_key"
+// CHECK-BTE: "branch-target-enforcement"