[AArch64][PAC] Lower direct authenticated calls to ptrauth constants.

ahmedbougacha · ahmedbougacha · commit adffb5b0ceef · 2024-07-03T20:10:08.000-07:00
diff --git a/llvm/lib/CodeGen/GlobalISel/CallLowering.cpp b/llvm/lib/CodeGen/GlobalISel/CallLowering.cpp
@@ -149,6 +149,14 @@ bool CallLowering::lowerCall(MachineIRBuilder &MIRBuilder, const CallBase &CB,
   // Try looking through a bitcast from one function type to another.
   // Commonly happens with calls to objc_msgSend().
   const Value *CalleeV = CB.getCalledOperand()->stripPointerCasts();
+
+  // If IRTranslator chose to drop the ptrauth info, we can turn this into
+  // a direct call.
+  if (!PAI && CB.countOperandBundlesOfType(LLVMContext::OB_ptrauth)) {
+    CalleeV = cast<ConstantPtrAuth>(CalleeV)->getPointer();
+    assert(isa<Function>(CalleeV));
+  }
+
   if (const Function *F = dyn_cast<Function>(CalleeV)) {
     if (F->hasFnAttribute(Attribute::NonLazyBind)) {
       LLT Ty = getLLTForType(*F->getType(), DL);
diff --git a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
@@ -2649,17 +2649,24 @@ bool IRTranslator::translateCallBase(const CallBase &CB,
   }
 
   std::optional<CallLowering::PtrAuthInfo> PAI;
-  if (CB.countOperandBundlesOfType(LLVMContext::OB_ptrauth)) {
+  if (auto Bundle = CB.getOperandBundle(LLVMContext::OB_ptrauth)) {
     // Functions should never be ptrauth-called directly.
     assert(!CB.getCalledFunction() && "invalid direct ptrauth call");
 
-    auto PAB = CB.getOperandBundle("ptrauth");
-    const Value *Key = PAB->Inputs[0];
-    const Value *Discriminator = PAB->Inputs[1];
-
-    Register DiscReg = getOrCreateVReg(*Discriminator);
-    PAI = CallLowering::PtrAuthInfo{cast<ConstantInt>(Key)->getZExtValue(),
-                                    DiscReg};
+    const Value *Key = Bundle->Inputs[0];
+    const Value *Discriminator = Bundle->Inputs[1];
+
+    // Look through ptrauth constants to try to eliminate the matching bundle
+    // and turn this into a direct call with no ptrauth.
+    // CallLowering will use the raw pointer if it doesn't find the PAI.
+    auto *CalleeCPA = dyn_cast<ConstantPtrAuth>(CB.getCalledOperand());
+    if (!CalleeCPA || !isa<Function>(CalleeCPA->getPointer()) ||
+        !CalleeCPA->isKnownCompatibleWith(Key, Discriminator, *DL)) {
+      // If we can't make it direct, package the bundle into PAI.
+      Register DiscReg = getOrCreateVReg(*Discriminator);
+      PAI = CallLowering::PtrAuthInfo{cast<ConstantInt>(Key)->getZExtValue(),
+                                      DiscReg};
+    }
   }
 
   Register ConvergenceCtrlToken = 0;
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -9374,6 +9374,14 @@ void SelectionDAGBuilder::LowerCallSiteWithPtrAuthBundle(
   assert(Discriminator->getType()->isIntegerTy(64) &&
          "Invalid ptrauth discriminator");
 
+  // Look through ptrauth constants to find the raw callee.
+  // Do a direct unauthenticated call if we found it and everything matches.
+  if (auto *CalleeCPA = dyn_cast<ConstantPtrAuth>(CalleeV))
+    if (CalleeCPA->isKnownCompatibleWith(Key, Discriminator,
+                                         DAG.getDataLayout()))
+      return LowerCallTo(CB, getValue(CalleeCPA->getPointer()), CB.isTailCall(),
+                         CB.isMustTailCall(), EHPadBB);
+
   // Functions should never be ptrauth-called directly.
   assert(!isa<Function>(CalleeV) && "invalid direct ptrauth call");
 
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-call.ll b/llvm/test/CodeGen/AArch64/ptrauth-call.ll
@@ -269,4 +269,138 @@ define i32 @test_tailcall_ib_arg_ind(ptr %arg0, i64 %arg1) #0 {
   ret i32 %tmp1
 }
 
+; Test direct calls
+
+define i32 @test_direct_call() #0 {
+; DARWIN-LABEL: test_direct_call:
+; DARWIN-NEXT:   stp x29, x30, [sp, #-16]!
+; DARWIN-NEXT:   bl _f
+; DARWIN-NEXT:   ldp x29, x30, [sp], #16
+; DARWIN-NEXT:   ret
+;
+; ELF-LABEL: test_direct_call:
+; ELF-NEXT:   str x30, [sp, #-16]!
+; ELF-NEXT:   bl f
+; ELF-NEXT:   ldr x30, [sp], #16
+; ELF-NEXT:   ret
+  %tmp0 = call i32 ptrauth(ptr @f, i32 0, i64 42)() [ "ptrauth"(i32 0, i64 42) ]
+  ret i32 %tmp0
+}
+
+define i32 @test_direct_tailcall(ptr %arg0) #0 {
+; DARWIN-LABEL: test_direct_tailcall:
+; DARWIN:    b _f
+;
+; ELF-LABEL: test_direct_tailcall:
+; ELF-NEXT:   b f
+  %tmp0 = tail call i32 ptrauth(ptr @f, i32 0, i64 42)() [ "ptrauth"(i32 0, i64 42) ]
+  ret i32 %tmp0
+}
+
+define i32 @test_direct_call_mismatch() #0 {
+; DARWIN-LABEL: test_direct_call_mismatch:
+; DARWIN-NEXT:   stp x29, x30, [sp, #-16]!
+; DARWIN-NEXT:   adrp x16, _f@GOTPAGE
+; DARWIN-NEXT:   ldr x16, [x16, _f@GOTPAGEOFF]
+; DARWIN-NEXT:   mov x17, #42
+; DARWIN-NEXT:   pacia x16, x17
+; DARWIN-NEXT:   mov x8, x16
+; DARWIN-NEXT:   mov x17, #42
+; DARWIN-NEXT:   blrab x8, x17
+; DARWIN-NEXT:   ldp x29, x30, [sp], #16
+; DARWIN-NEXT:   ret
+;
+; ELF-LABEL: test_direct_call_mismatch:
+; ELF-NEXT:   str x30, [sp, #-16]!
+; ELF-NEXT:   adrp x16, :got:f
+; ELF-NEXT:   ldr x16, [x16, :got_lo12:f]
+; ELF-NEXT:   mov x17, #42
+; ELF-NEXT:   pacia x16, x17
+; ELF-NEXT:   mov x8, x16
+; ELF-NEXT:   mov x17, #42
+; ELF-NEXT:   blrab x8, x17
+; ELF-NEXT:   ldr x30, [sp], #16
+; ELF-NEXT:   ret
+  %tmp0 = call i32 ptrauth(ptr @f, i32 0, i64 42)() [ "ptrauth"(i32 1, i64 42) ]
+  ret i32 %tmp0
+}
+
+define i32 @test_direct_call_addr() #0 {
+; DARWIN-LABEL: test_direct_call_addr:
+; DARWIN-NEXT:   stp x29, x30, [sp, #-16]!
+; DARWIN-NEXT:   bl _f
+; DARWIN-NEXT:   ldp x29, x30, [sp], #16
+; DARWIN-NEXT:   ret
+;
+; ELF-LABEL: test_direct_call_addr:
+; ELF-NEXT:   str x30, [sp, #-16]!
+; ELF-NEXT:   bl f
+; ELF-NEXT:   ldr x30, [sp], #16
+; ELF-NEXT:   ret
+  %tmp0 = call i32 ptrauth(ptr @f, i32 1, i64 0, ptr @f.ref.ib.0.addr)() [ "ptrauth"(i32 1, i64 ptrtoint (ptr @f.ref.ib.0.addr to i64)) ]
+  ret i32 %tmp0
+}
+
+define i32 @test_direct_call_addr_blend() #0 {
+; DARWIN-LABEL: test_direct_call_addr_blend:
+; DARWIN-NEXT:   stp x29, x30, [sp, #-16]!
+; DARWIN-NEXT:   bl _f
+; DARWIN-NEXT:   ldp x29, x30, [sp], #16
+; DARWIN-NEXT:   ret
+;
+; ELF-LABEL: test_direct_call_addr_blend:
+; ELF-NEXT:   str x30, [sp, #-16]!
+; ELF-NEXT:   bl f
+; ELF-NEXT:   ldr x30, [sp], #16
+; ELF-NEXT:   ret
+  %tmp0 = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr @f.ref.ib.42.addr to i64), i64 42)
+  %tmp1 = call i32 ptrauth(ptr @f, i32 1, i64 42, ptr @f.ref.ib.42.addr)() [ "ptrauth"(i32 1, i64 %tmp0) ]
+  ret i32 %tmp1
+}
+
+define i32 @test_direct_call_addr_gep_different_index_types() #0 {
+; DARWIN-LABEL: test_direct_call_addr_gep_different_index_types:
+; DARWIN-NEXT:   stp x29, x30, [sp, #-16]!
+; DARWIN-NEXT:   bl _f
+; DARWIN-NEXT:   ldp x29, x30, [sp], #16
+; DARWIN-NEXT:   ret
+;
+; ELF-LABEL: test_direct_call_addr_gep_different_index_types:
+; ELF-NEXT:   str x30, [sp, #-16]!
+; ELF-NEXT:   bl f
+; ELF-NEXT:   ldr x30, [sp], #16
+; ELF-NEXT:   ret
+  %tmp0 = call i32 ptrauth(ptr @f, i32 1, i64 0, ptr getelementptr ({ ptr }, ptr @f_struct.ref.ib.0.addr, i64 0, i32 0))() [ "ptrauth"(i32 1, i64 ptrtoint (ptr getelementptr ({ ptr }, ptr @f_struct.ref.ib.0.addr, i32 0, i32 0) to i64)) ]
+  ret i32 %tmp0
+}
+
+define i32 @test_direct_call_addr_blend_gep_different_index_types() #0 {
+; DARWIN-LABEL: test_direct_call_addr_blend_gep_different_index_types:
+; DARWIN-NEXT:   stp x29, x30, [sp, #-16]!
+; DARWIN-NEXT:   bl _f
+; DARWIN-NEXT:   ldp x29, x30, [sp], #16
+; DARWIN-NEXT:   ret
+;
+; ELF-LABEL: test_direct_call_addr_blend_gep_different_index_types:
+; ELF-NEXT:   str x30, [sp, #-16]!
+; ELF-NEXT:   bl f
+; ELF-NEXT:   ldr x30, [sp], #16
+; ELF-NEXT:   ret
+  %tmp0 = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr getelementptr ({ ptr }, ptr @f_struct.ref.ib.123.addr, i32 0, i32 0) to i64), i64 123)
+  %tmp1 = call i32 ptrauth(ptr @f, i32 1, i64 123, ptr getelementptr ({ ptr }, ptr @f_struct.ref.ib.123.addr, i64 0, i32 0))() [ "ptrauth"(i32 1, i64 %tmp0) ]
+  ret i32 %tmp1
+}
+
+attributes #0 = { nounwind }
+
+@f.ref.ib.42.addr = external global ptr
+@f.ref.ib.0.addr = external global ptr
+@f_struct.ref.ib.0.addr = external global ptr
+@f_struct.ref.ib.123.addr = external global ptr
+
+declare void @f()
+
+declare i64 @llvm.ptrauth.auth(i64, i32, i64)
+declare i64 @llvm.ptrauth.blend(i64, i64)
+
 attributes #0 = { nounwind }
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-invoke.ll b/llvm/test/CodeGen/AArch64/ptrauth-invoke.ll
@@ -96,6 +96,90 @@ continuebb:
   ret i32 %tmp0
 }
 
+; DARWIN-LABEL: _test_invoke_ia_0_direct:
+; DARWIN-NEXT: [[FNBEGIN:L.*]]:
+; DARWIN-NEXT:  .cfi_startproc
+; DARWIN-NEXT:  .cfi_personality 155, ___gxx_personality_v0
+; DARWIN-NEXT:  .cfi_lsda 16, [[EXCEPT:Lexception[0-9]+]]
+; DARWIN-NEXT: ; %bb.0:
+; DARWIN-NEXT:  stp x20, x19, [sp, #-32]!
+; DARWIN-NEXT:  stp x29, x30, [sp, #16]
+; DARWIN-NEXT:  .cfi_def_cfa_offset 32
+; DARWIN-NEXT:  .cfi_offset w30, -8
+; DARWIN-NEXT:  .cfi_offset w29, -16
+; DARWIN-NEXT:  .cfi_offset w19, -24
+; DARWIN-NEXT:  .cfi_offset w20, -32
+; DARWIN-NEXT: [[PRECALL:L.*]]:
+; DARWIN-NEXT:  bl _baz
+
+; DARWIN-SDAG-NEXT: [[POSTCALL:L.*]]:
+; DARWIN-SDAG-NEXT: ; %bb.1:
+; DARWIN-SDAG-NEXT:  mov x19, x0
+
+; DARWIN-GISEL-NEXT:  mov x19, x0
+; DARWIN-GISEL-NEXT: [[POSTCALL:L.*]]:
+
+; DARWIN-NEXT: [[CALLBB:L.*]]:
+; DARWIN-NEXT:  bl _foo
+; DARWIN-NEXT:  mov x0, x19
+; DARWIN-NEXT:  ldp x29, x30, [sp, #16]
+; DARWIN-NEXT:  ldp x20, x19, [sp], #32
+; DARWIN-NEXT:  ret
+; DARWIN-NEXT: [[LPADBB:LBB[0-9_]+]]:
+; DARWIN-NEXT: [[LPAD:L.*]]:
+; DARWIN-NEXT:  mov w19, #-1
+; DARWIN-NEXT:  b [[CALLBB]]
+
+; ELF-LABEL: test_invoke_ia_0_direct:
+; ELF-NEXT: [[FNBEGIN:.L.*]]:
+; ELF-NEXT:  .cfi_startproc
+; ELF-NEXT:  .cfi_personality 156, DW.ref.__gxx_personality_v0
+; ELF-NEXT:  .cfi_lsda 28, [[EXCEPT:.Lexception[0-9]+]]
+; ELF-NEXT: // %bb.0:
+; ELF-NEXT:  stp x30, x19, [sp, #-16]!
+; ELF-NEXT:  .cfi_def_cfa_offset 16
+; ELF-NEXT:  .cfi_offset w19, -8
+; ELF-NEXT:  .cfi_offset w30, -16
+; ELF-NEXT: [[PRECALL:.L.*]]:
+; ELF-NEXT:  bl baz
+
+; ELF-SDAG-NEXT: [[POSTCALL:.L.*]]:
+; ELF-SDAG-NEXT: // %bb.1:
+; ELF-SDAG-NEXT:  mov w19, w0
+
+; ELF-GISEL-NEXT:  mov w19, w0
+; ELF-GISEL-NEXT: [[POSTCALL:.L.*]]:
+
+; ELF-NEXT: [[CALLBB:.L.*]]:
+; ELF-NEXT:  bl foo
+; ELF-NEXT:  mov w0, w19
+; ELF-NEXT:  ldp x30, x19, [sp], #16
+; ELF-NEXT:  ret
+; ELF-NEXT: [[LPADBB:.LBB[0-9_]+]]:
+; ELF-NEXT: [[LPAD:.L.*]]:
+; ELF-NEXT:  mov w19, #-1
+; ELF-NEXT:  b [[CALLBB]]
+
+; CHECK-LABEL: GCC_except_table{{.*}}:
+; CHECK-NEXT: [[EXCEPT]]:
+; CHECK:       .uleb128 [[POSTCALL]]-[[PRECALL]] {{.*}} Call between [[PRECALL]] and [[POSTCALL]]
+; CHECK-NEXT:  .uleb128 [[LPAD]]-[[FNBEGIN]]     {{.*}}   jumps to [[LPAD]]
+; CHECK-NEXT:  .byte 0                           {{.*}} On action: cleanup
+
+define i32 @test_invoke_ia_0_direct() #0 personality ptr @__gxx_personality_v0 {
+  %tmp0 = invoke i32 ptrauth (ptr @baz, i32 0)() [ "ptrauth"(i32 0, i64 0) ] to label %continuebb
+            unwind label %unwindbb
+
+unwindbb:
+  %tmp1 = landingpad { ptr, i32 } cleanup
+  call void @foo()
+  ret i32 -1
+
+continuebb:
+  call void @foo()
+  ret i32 %tmp0
+}
+
 @_ZTIPKc = external constant ptr
 @hello_str = private unnamed_addr constant [6 x i8] c"hello\00", align 1
 
@@ -265,6 +349,7 @@ continuebb:
 
 declare void @foo()
 declare void @bar(ptr)
+declare i32 @baz()
 
 declare i32 @__gxx_personality_v0(...)
 declare ptr @__cxa_allocate_exception(i64)
