[AArch64] Return Address Signing B Key Support

[LukeCheeseman]

Summary

• Add command line options for selecting the B key when enabling return address signing (also tidy up the msign_return_address tablegen definiton by suffixing the option with EQ)
• This expands the -msign-return-address=<scope> to now be -msign-return-address=<scope>[+<key>] where <key> is one of a_key or b_key. The default key is a_key.
• Generate the "sign-return-address-key" function attribute with the relevant key

[SjoerdMeijer]

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

1–9

If the default is the a_key, does this test need to check the a_key attribute?

[olista01]

This looks like it has the same problem as D51418 (doesn't get applied to C++ static constructor functions).

[LukeCheeseman]

• updated tests to check the default a_key attribute is included in functions that have the return address signing scope attribute
• setting the return-address-sign-key attribute of global static constructors

[LukeCheeseman]

• Introduce the -mbranch-protection option. This is to be used for both pointer authentication and branch protection security features.
  • The branch protection compiler support will follow later
• The options available are -mbranch-protection=<type>(+<type>)* where <type> ::= [bti,pac-ret[+leaf,b-key]*]
• This should be the primary way of using branch protection features and the -msign-return-address option should be deprecated
• Remove the b-key selection from the earlier patch for -msign-return-address

[javed.absar]

include/clang/Frontend/CodeGenOptions.h

117

Perhaps a line of comment on each enum-value would be useful (much like the code around).

[LukeCheeseman]

include/clang/Frontend/CodeGenOptions.h

117

I'm not sure what value there would be in adding extra comments here, I think the values are fairly self explanatory.

[olista01]

lib/Frontend/CompilerInvocation.cpp

1142

The driver code is emitting a separate -msign-return-address-key= option, but this is expecting the key to still be part of -msign-return-address=, so we can never end up selecting the B key. Out of the two approaches, I prefer having multiple simpler options in CC1, so that we don't have to repeat all of the parsing that happens in the driver.

Also, you're emitting the -mbranch-target-enforce option in the driver, but not handling it here. We should either reject it in the driver for now, or add the CC1 part of that in this patch.

[LukeCheeseman]

• Stop parsing msign-return-address as a scope and key pair
• pass bti value through to the CC1 driver and handle it, this adds the branch-target-enforce attribute to functions

[olista01]

lib/CodeGen/CGDeclCXX.cpp

364

LLVM style has the opening brace on the same line as the if. There's a git-clang-format script in the clang repo which can fix up style issues like this for a whole patch, changing touching lines you haven't changes.

374

I think we need to handle branch-target-enforcement here too.

lib/CodeGen/TargetInfo.cpp

4982

Style nit: brace on same line as if.

lib/Driver/ToolChains/Clang.cpp

1433

Please add a comment about what the parts of the return value are.

1440

I'm not sure what this is trying to say. Is it referring to the "standard" case below?

1448

I'd make this 4, because that's the longest sensible argument ("pac-ret+leaf+b-key+bti").

1460

Add a comment about the fact that "leaf" and "b-key" must follow "pac-ret", to explain why we don't just parse this in one loop.

1543

Nit: brace on same line as if.

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

7

This should also test branch target enforcement (it's also missing from the code).

[olista01]

LGTM, thanks!

[LukeCheeseman]

This was committed under 345273. (Forgot to mention the revision in the commit message)

include/clang/Basic/DiagnosticDriverKinds.td

 def warn_debug_compression_unavailable : Warning<"cannot compress debug sections (zlib not installed)">,
   InGroup<DiagGroup<"debug-compression-unavailable">>;
 def warn_drv_disabling_vptr_no_rtti_default : Warning<
   "implicitly disabling vptr sanitizer because rtti wasn't enabled">,
   InGroup<AutoDisableVptrSanitizer>;
 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">;
 def note_drv_t_option_is_global : Note<
   "The last /TC or /TP option takes precedence over earlier instances">;
 def note_drv_address_sanitizer_debug_runtime : Note<
   "AddressSanitizer doesn't support linking with debug runtime libraries yet">;
 def note_drv_use_standard : Note<"use '%0'"

include/clang/Driver/CC1Options.td

def flto_unit: Flag<["-"], "flto-unit">,
     HelpText<"Emit IR to support LTO unit features (CFI, whole program vtable opt)">;
 def fno_lto_unit: Flag<["-"], "fno-lto-unit">;
 def fthin_link_bitcode_EQ : Joined<["-"], "fthin-link-bitcode=">,
     HelpText<"Write minimized bitcode to <file> for the ThinLTO thin link only">;
 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
 //===----------------------------------------------------------------------===//
 
 def sys_header_deps : Flag<["-"], "sys-header-deps">,
   HelpText<"Include system headers in dependency output">;
 def module_file_deps : Flag<["-"], "module-file-deps">,

include/clang/Driver/Options.td

 foreach i = {1-7,18,20} in
   def ffixed_x#i : Flag<["-"], "ffixed-x"#i>, Group<m_aarch64_Features_Group>,
     HelpText<"Reserve the "#i#" register (AArch64 only)">;
 
 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>;
 def mnontrapping_fptoint : Flag<["-"], "mnontrapping-fptoint">, Group<m_wasm_Features_Group>;
 def mno_nontrapping_fptoint : Flag<["-"], "mno-nontrapping-fptoint">, Group<m_wasm_Features_Group>;
 def msign_ext : Flag<["-"], "msign-ext">, Group<m_wasm_Features_Group>;
 def mno_sign_ext : Flag<["-"], "mno-sign-ext">, Group<m_wasm_Features_Group>;
 def mexception_handing : Flag<["-"], "mexception-handling">, Group<m_wasm_Features_Group>;

include/clang/Frontend/CodeGenOptions.h

public:
   };
 
   enum SignReturnAddressScope {
     None,    // No signing for any function
     NonLeaf, // Sign the return address of functions that spill LR
     All      // Sign the return address of all functions
   };
 
+  enum SignReturnAddressKeyValue { AKey, BKey };
+
   /// The code model to use (-mcmodel).
   std::string CodeModel;
 
   /// The filename with path we use for coverage data files. The runtime
   /// allows further manipulation with the GCOV_PREFIX and GCOV_PREFIX_STRIP
   /// environment variables.
   std::string CoverageDataFile;
 

include/clang/Frontend/CodeGenOptions.def

 
 /// Whether to emit all vtables
 CODEGENOPT(ForceEmitVTables, 1, 0)
 
 /// Whether to emit an address-significance table into the object file.
 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)
 
 #undef CODEGENOPT
 #undef ENUM_CODEGENOPT
 #undef VALUE_CODEGENOPT
 

lib/CodeGen/CGDeclCXX.cpp

if (getLangOpts().Sanitize.has(SanitizerKind::SafeStack) &&
     Fn->addFnAttr(llvm::Attribute::SafeStack);
 
   if (getLangOpts().Sanitize.has(SanitizerKind::ShadowCallStack) &&
       !isInSanitizerBlacklist(SanitizerKind::ShadowCallStack, Fn, Loc))
     Fn->addFnAttr(llvm::Attribute::ShadowCallStack);
 
   auto RASignKind = getCodeGenOpts().getSignReturnAddress();
   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");
+  }
+
   return Fn;
 }
 
 /// Create a global pointer to a function that will initialize a global
 /// variable.  The user has requested that this pointer be emitted in a specific
 /// section.
 void CodeGenModule::EmitPointerToInitFunc(const VarDecl *D,
                                           llvm::GlobalVariable *GV,

lib/CodeGen/TargetInfo.cpp

public:
   void setTargetAttributes(const Decl *D, llvm::GlobalValue *GV,
                            CodeGen::CodeGenModule &CGM) const override {
     const FunctionDecl *FD = dyn_cast_or_null<FunctionDecl>(D);
     if (!FD)
       return;
     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");
+
+      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");
   }
 };
 
 class WindowsAArch64TargetCodeGenInfo : public AArch64TargetCodeGenInfo {
 public:
   WindowsAArch64TargetCodeGenInfo(CodeGenTypes &CGT, AArch64ABIInfo::ABIKind K)
       : AArch64TargetCodeGenInfo(CGT, K) {}
 

lib/Driver/ToolChains/Clang.cpp

void Clang::RenderTargetOptions(const llvm::Triple &EffectiveTriple,
 
   case llvm::Triple::wasm32:
   case llvm::Triple::wasm64:
     AddWebAssemblyTargetArgs(Args, CmdArgs);
     break;
   }
 }
 
+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, 3> BranchProtection;
+    StringRef(A->getValue()).split(BranchProtection, '+');
+
+    auto Protection = BranchProtection.begin();
+    // Check for the -mbranch-protection options
+    // looking for -mbranch-protection=<type>(+<type>)* where
+    // <type> ::= [bti,pac-ret[+leaf,+b-key]*]
+    while (Protection != BranchProtection.end()) {
+      if (Protection->equals("bti"))
+        IndirectBranches = true;
+      else if (Protection->equals("pac-ret")) {
+        Scope = "non-leaf";
+        while (++Protection != BranchProtection.end()) {
+          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) {
   const char *ABIName = nullptr;
   if (Arg *A = Args.getLastArg(options::OPT_mabi_EQ))
     ABIName = A->getValue();
   else if (Triple.isOSDarwin())
     ABIName = "darwinpcs";
if (Arg *A = Args.getLastArg(options::OPT_mglobal_merge,
                                options::OPT_mno_global_merge)) {
     CmdArgs.push_back("-mllvm");
     if (A->getOption().matches(options::OPT_mno_global_merge))
       CmdArgs.push_back("-aarch64-enable-global-merge=false");
     else
       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=") + Scope));
     CmdArgs.push_back(
-        Args.MakeArgString(Twine("-msign-return-address=") + A->getValue()));
+        Args.MakeArgString(Twine("-msign-return-address-key=") + Key));
+    if (IndirectBranches)
+      CmdArgs.push_back("-mbranch-target-enforce");
   }
 }
 
 void Clang::AddMIPSTargetArgs(const ArgList &Args,
                               ArgStringList &CmdArgs) const {
   const Driver &D = getToolChain().getDriver();
   StringRef CPUName;
   StringRef ABIName;

lib/Frontend/CompilerInvocation.cpp

static bool ParseCodeGenArgs(CodeGenOptions &Opts, ArgList &Args, InputKind IK,
 
   Opts.EmitCheckPathComponentsToStrip = getLastArgIntValue(
       Args, OPT_fsanitize_undefined_strip_path_components_EQ, 0, Diags);
 
   Opts.EmitVersionIdentMetadata = Args.hasFlag(OPT_Qy, OPT_Qn, true);
 
   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"))
       Opts.setSignReturnAddress(CodeGenOptions::SignReturnAddressScope::All);
     else if (SignScope.equals_lower("non-leaf"))
       Opts.setSignReturnAddress(
           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);
 
   return Success;
 }
 

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

-// 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() {}

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

 // 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
 
 struct Foo {
   Foo() {}
   ~Foo() {}
 };
 
 Foo f;
 
 // CHECK: @llvm.global_ctors {{.*}}i32 65535, void ()* @[[CTOR_FN:.*]], i8* null
 
 // 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"