[PtrAuth] Add support for arm64e intrinsics

This patch provides barebones support for ARM64e instructions to be emitted like braaz and paciza.

Author: oskarwirga

compiler/rustc_codegen_llvm/src/builder.rs

@@ -431,6 +431,11 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
             bundles.push(kcfi_bundle);
         }
 
+        let pauth = self.ptrauth_operand_bundle(llfn);
+        if let Some(p) = pauth.as_ref().map(|b| b.as_ref()) {
+            bundles.push(p);
+        }
+
         let invoke = unsafe {
             llvm::LLVMBuildInvokeWithOperandBundles(
                 self.llbuilder,
@@ -1417,6 +1422,11 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
             bundles.push(kcfi_bundle);
         }
 
+        let pauth = self.ptrauth_operand_bundle(llfn);
+        if let Some(p) = pauth.as_ref().map(|b| b.as_ref()) {
+            bundles.push(p);
+        }
+
         let call = unsafe {
             llvm::LLVMBuildCallWithOperandBundles(
                 self.llbuilder,
@@ -1785,6 +1795,11 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         if let Some(kcfi_bundle) = kcfi_bundle.as_ref().map(|b| b.as_ref()) {
             bundles.push(kcfi_bundle);
         }
+        
+        let pauth = self.ptrauth_operand_bundle(llfn);
+        if let Some(p) = pauth.as_ref().map(|b| b.as_ref()) {
+            bundles.push(p);
+        }
 
         let callbr = unsafe {
             llvm::LLVMBuildCallBr(
@@ -1905,6 +1920,18 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         kcfi_bundle
     }
 
+    fn ptrauth_operand_bundle(&mut self, llfn: &'ll Value)
+        -> Option<llvm::OperandBundleBox<'ll>>
+    {
+        let on_arm64e = self.cx.sess().target.llvm_target.starts_with("arm64e-apple-");
+        if !on_arm64e { return None; }
+
+        let is_non_gv_fn_ptr = unsafe { llvm::LLVMRustIsNonGVFunctionPointerTy(llfn) };
+        if !is_non_gv_fn_ptr { return None; }
+
+        Some(llvm::OperandBundleBox::new("ptrauth", &[self.const_i32(0), self.const_u64(0)]))
+    }
+
     /// Emits a call to `llvm.instrprof.increment`. Used by coverage instrumentation.
     #[instrument(level = "debug", skip(self))]
     pub(crate) fn instrprof_increment(

compiler/rustc_codegen_llvm/src/common.rs

@@ -24,6 +24,20 @@ use crate::llvm::{self, BasicBlock, ConstantInt, FALSE, Metadata, TRUE, ToLlvmBo
 use crate::type_::Type;
 use crate::value::Value;
 
+pub(crate) fn sign_fn_ptr_if_arm64e<'ll>(
+    cx: &CodegenCx<'ll, '_>,
+    llfn: &'ll llvm::Value,
+) -> &'ll llvm::Value {
+    if !cx.sess().target.llvm_target.starts_with("arm64e-apple-") {
+        return llfn;
+    }
+
+    unsafe {
+        let authed = llvm::LLVMRustConstPtrAuth(llfn as *const _ as *mut _, /*IA*/ 0, /*disc*/ 0);
+        &*authed
+    }
+}
+
 /*
 * A note on nomenclature of linking: "extern", "foreign", and "upcall".
 *

compiler/rustc_codegen_llvm/src/context.rs

@@ -767,7 +767,14 @@ impl<'ll, 'tcx> MiscCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {
     }
 
     fn get_fn_addr(&self, instance: Instance<'tcx>) -> &'ll Value {
-        get_fn(self, instance)
+        let llfn = get_fn(self, instance);
+
+        let name_bytes = llvm::get_value_name(llfn);
+        if name_bytes.starts_with(b"llvm.") {
+            return llfn;
+        }
+
+        common::sign_fn_ptr_if_arm64e(self, llfn)
     }
 
     fn eh_personality(&self) -> &'ll Value {

compiler/rustc_codegen_llvm/src/llvm/mod.rs

@@ -437,3 +437,11 @@ pub(crate) fn append_module_inline_asm<'ll>(llmod: &'ll Module, asm: &[u8]) {
         LLVMAppendModuleInlineAsm(llmod, asm.as_ptr(), asm.len());
     }
 }
+
+unsafe extern "C" {
+    pub(crate) fn LLVMRustConstPtrAuth(
+        Fn: *mut Value,
+        Key: c_uint,
+        Disc: u64,
+    ) -> *mut Value;
+}

compiler/rustc_codegen_ssa/src/back/metadata.rs

@@ -216,7 +216,9 @@ pub(crate) fn create_object_file(sess: &Session) -> Option<write::Object<'static
     file.set_sub_architecture(sub_architecture);
     if sess.target.is_like_darwin {
         if macho_is_arm64e(&sess.target) {
-            file.set_macho_cpu_subtype(object::macho::CPU_SUBTYPE_ARM64E);
+            const ARM64E_SUBTYPE: u32 = object::macho::CPU_SUBTYPE_ARM64E | 
+                                        object::macho::CPU_SUBTYPE_PTRAUTH_ABI;
+            file.set_macho_cpu_subtype(ARM64E_SUBTYPE);
         }
 
         file.set_macho_build_version(macho_object_build_version_for_target(sess))

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -1945,3 +1945,21 @@ extern "C" void LLVMRustSetNoSanitizeHWAddress(LLVMValueRef Global) {
   MD.NoHWAddress = true;
   GV.setSanitizerMetadata(MD);
 }
+
+extern "C" LLVMValueRef LLVMRustConstPtrAuth(LLVMValueRef Ptr,
+                                             unsigned Key,
+                                             uint64_t Disc) {
+  auto *V = unwrap<Value>(Ptr);
+  auto *C = dyn_cast<Constant>(V);
+  if (!C) return Ptr;
+  auto *PTy = dyn_cast<PointerType>(C->getType());
+  if (!PTy) return Ptr;
+  if (isa<UndefValue>(C) || isa<ConstantPointerNull>(C)) return Ptr;
+
+  LLVMContext &Ctx = C->getContext();
+  auto *KeyC  = ConstantInt::get(Type::getInt32Ty(Ctx), Key);
+  auto *DiscC = ConstantInt::get(Type::getInt64Ty(Ctx), Disc);
+  auto *NullAD = ConstantPointerNull::get(PTy); // <— important
+
+  return wrap(ConstantPtrAuth::get(C, KeyC, DiscC, NullAD));
+}