[Coroutines] Support ptrauth.

Author: ahmedbougacha

llvm/lib/Transforms/Coroutines/CoroInstr.h

@@ -25,6 +25,7 @@
 #ifndef LLVM_LIB_TRANSFORMS_COROUTINES_COROINSTR_H
 #define LLVM_LIB_TRANSFORMS_COROUTINES_COROINSTR_H
 
+#include "llvm/IR/GlobalPtrAuthInfo.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/Support/raw_ostream.h"
@@ -223,7 +224,16 @@ class LLVM_LIBRARY_VISIBILITY AnyCoroIdRetconInst : public AnyCoroIdInst {
   /// attributes, and calling convention of the continuation function(s)
   /// are taken from this declaration.
   Function *getPrototype() const {
-    return cast<Function>(getArgOperand(PrototypeArg)->stripPointerCasts());
+    Value *rawPrototype = getArgOperand(PrototypeArg)->stripPointerCasts();
+    if (auto global = GlobalPtrAuthInfo::analyze(rawPrototype)) {
+      rawPrototype =
+        const_cast<Constant*>(global->getPointer()->stripPointerCasts());
+    }
+    return cast<Function>(rawPrototype);
+  }
+
+  std::optional<GlobalPtrAuthInfo> getPtrAuthInfo() const {
+    return GlobalPtrAuthInfo::analyze(getArgOperand(PrototypeArg));
   }
 
   /// Return the function to use for allocating memory.

llvm/lib/Transforms/Coroutines/CoroInternal.h

@@ -119,6 +119,7 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
   };
 
   struct RetconLoweringStorage {
+    const Constant *ResumePtrAuthInfo;
     Function *ResumePrototype;
     Function *Alloc;
     Function *Dealloc;
@@ -175,10 +176,22 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
     return ConstantInt::get(getIndexType(), Value);
   }
 
+  std::optional<GlobalPtrAuthInfo> getResumePtrAuthInfo() const {
+    switch (ABI) {
+    case coro::ABI::Switch:
+    case coro::ABI::Async:
+      return std::nullopt;
+    case coro::ABI::Retcon:
+    case coro::ABI::RetconOnce:
+      if (!RetconLowering.ResumePtrAuthInfo) return std::nullopt;
+      return GlobalPtrAuthInfo::analyze(RetconLowering.ResumePtrAuthInfo);
+    }
+  }
+
   PointerType *getSwitchResumePointerType() const {
     assert(ABI == coro::ABI::Switch);
-  assert(FrameTy && "frame type not assigned");
-  return cast<PointerType>(FrameTy->getElementType(SwitchFieldIndex::Resume));
+    assert(FrameTy && "frame type not assigned");
+    return cast<PointerType>(FrameTy->getElementType(SwitchFieldIndex::Resume));
   }
 
   FunctionType *getResumeFunctionType() const {

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -937,6 +937,19 @@ void CoroCloner::create() {
     VMap[&A] = DummyArgs.back();
   }
 
+  switch (Shape.ABI) {
+  case coro::ABI::Retcon:
+  case coro::ABI::RetconOnce:
+    // In retcon coroutines, replace any remaining uses of the original
+    // storage value with the storage parameter.
+    VMap[Shape.getRetconCoroId()->getStorage()] = &*NewF->arg_begin();
+    break;
+
+  case coro::ABI::Switch:
+  case coro::ABI::Async:
+    break;
+  }
+
   SmallVector<ReturnInst *, 4> Returns;
 
   // Ignore attempts to change certain attributes of the function.
@@ -1091,8 +1104,10 @@ void CoroCloner::create() {
 
   // Remap frame pointer.
   Value *OldFramePtr = VMap[Shape.FramePtr];
-  NewFramePtr->takeName(OldFramePtr);
-  OldFramePtr->replaceAllUsesWith(NewFramePtr);
+  if (OldFramePtr != NewFramePtr) {
+    NewFramePtr->takeName(OldFramePtr);
+    OldFramePtr->replaceAllUsesWith(NewFramePtr);
+  }
 
   // Remap vFrame pointer.
   auto *NewVFrame = Builder.CreateBitCast(
@@ -1846,6 +1861,46 @@ static void splitAsyncCoroutine(Function &F, coro::Shape &Shape,
   }
 }
 
+/// Produced a signed pointer value with this same schema.  If the address
+/// is provided, and this schema uses address discrimination, use that
+/// address instead of the stored address discriminator.
+static Value *emitSignedPointer(IRBuilder<> &Builder,
+                                const GlobalPtrAuthInfo &Schema,
+                                Constant *ConstPointer,
+                                Value *Address) {
+  auto Module = Builder.GetInsertBlock()->getModule();
+  if (!Schema.hasAddressDiversity()) {
+    return Schema.createWithSameSchema(*Module, ConstPointer);
+  }
+
+  auto IntPtrTy = Schema.getAddrDiscriminator()->getType();
+
+  // Blend the provided address with the extra discriminator if the extra
+  // discriminator is non-zero.
+  Value *Discriminator = Builder.CreatePtrToInt(Address, IntPtrTy);
+  auto ExtraDiscriminator = const_cast<ConstantInt*>(Schema.getDiscriminator());
+  if (!ExtraDiscriminator->isNullValue()) {
+    auto Blend = Intrinsic::getDeclaration(Module, Intrinsic::ptrauth_blend);
+    Discriminator =
+      Builder.CreateCall(Blend, {Discriminator, ExtraDiscriminator});
+  }
+
+  // Sign the pointer with the discriminator.
+
+  // The intrinsic wants an integer.
+  Value *Pointer = Builder.CreatePtrToInt(ConstPointer, IntPtrTy);
+
+  // call i64 @llvm.ptrauth.sign.i64(i64 %pointer, i32 %key, i64 %discriminator)
+  auto Sign = Intrinsic::getDeclaration(Module, Intrinsic::ptrauth_sign);
+  auto Key = const_cast<ConstantInt*>(Schema.getKey());
+  Pointer = Builder.CreateCall(Sign, { Pointer, Key, Discriminator });
+
+  // Convert back to the original pointer type.
+  Pointer = Builder.CreateIntToPtr(Pointer, ConstPointer->getType());
+
+  return Pointer;
+}
+
 static void splitRetconCoroutine(Function &F, coro::Shape &Shape,
                                  SmallVectorImpl<Function *> &Clones) {
   assert(Shape.ABI == coro::ABI::Retcon ||
@@ -1957,9 +2012,21 @@ static void splitRetconCoroutine(Function &F, coro::Shape &Shape,
       Builder.CreateRet(RetV);
     }
 
+    // Sign the continuation value if requested.
+    Value *ContinuationValue = Continuation;
+    if (auto PtrAuthInfo = Shape.getResumePtrAuthInfo()) {
+      assert(!PtrAuthInfo->hasAddressDiversity() ||
+             PtrAuthInfo->hasSpecialAddressDiscriminator(
+               GlobalPtrAuthInfo::AddrDiscriminator_UseCoroStorage));
+      IRBuilder<> Builder(Branch);
+      ContinuationValue =
+        emitSignedPointer(Builder, *PtrAuthInfo, Continuation,
+                          Id->getStorage());
+    }
+
     // Branch to the return block.
     Branch->setSuccessor(0, ReturnBB);
-    ReturnPHIs[0]->addIncoming(Continuation, SuspendBB);
+    ReturnPHIs[0]->addIncoming(ContinuationValue, SuspendBB);
     size_t NextPHIIndex = 1;
     for (auto &VUse : Suspend->value_operands())
       ReturnPHIs[NextPHIIndex++]->addIncoming(&*VUse, SuspendBB);

llvm/lib/Transforms/Coroutines/Coroutines.cpp

@@ -337,11 +337,23 @@ void coro::Shape::buildFrom(Function &F) {
                   : coro::ABI::RetconOnce);
     auto Prototype = ContinuationId->getPrototype();
     this->RetconLowering.ResumePrototype = Prototype;
+    auto AuthInfo = ContinuationId->getPtrAuthInfo();
+    this->RetconLowering.ResumePtrAuthInfo = (AuthInfo ? AuthInfo->getGV() : nullptr);
     this->RetconLowering.Alloc = ContinuationId->getAllocFunction();
     this->RetconLowering.Dealloc = ContinuationId->getDeallocFunction();
     this->RetconLowering.ReturnBlock = nullptr;
     this->RetconLowering.IsFrameInlineInStorage = false;
 
+    if (AuthInfo && AuthInfo->hasAddressDiversity() &&
+        !AuthInfo->hasSpecialAddressDiscriminator(
+          GlobalPtrAuthInfo::AddrDiscriminator_UseCoroStorage)) {
+#ifndef NDEBUG
+      AuthInfo->getGV()->dump();
+#endif
+      report_fatal_error("ptrauth-signed prototype must not have address "
+                         "diversity");
+    }
+
     // Determine the result value types, and make sure they match up with
     // the values passed to the suspends.
     auto ResultTys = getRetconResultTypes();
@@ -512,6 +524,9 @@ void coro::Shape::emitDealloc(IRBuilder<> &Builder, Value *Ptr,
 /// Check that the given value is a well-formed prototype for the
 /// llvm.coro.id.retcon.* intrinsics.
 static void checkWFRetconPrototype(const AnyCoroIdRetconInst *I, Value *V) {
+  if (auto PtrAuth = GlobalPtrAuthInfo::analyze(V)) {
+    V = const_cast<Constant*>(PtrAuth->getPointer());
+  }
   auto F = dyn_cast<Function>(V->stripPointerCasts());
   if (!F)
     fail(I, "llvm.coro.id.retcon.* prototype not a Function", V);

llvm/test/Transforms/Coroutines/coro-retcon-ptrauth.ll

@@ -0,0 +1,105 @@
+; RUN: opt -passes='module(coro-early),cgscc(coro-split),coro-cleanup' -S %s | FileCheck %s
+target datalayout = "E-p:64:64"
+
+%swift.type = type { i64 }
+%swift.opaque = type opaque
+%T4red215EmptyCollectionV = type opaque
+%TSi = type <{ i64 }>
+
+@prototype.signed = private constant { ptr, i32, i64, i64 } { ptr @prototype, i32 2, i64 0, i64 12867 }, section "llvm.ptrauth"
+
+define ptr @f(ptr %buffer, i32 %n) {
+entry:
+  %id = call token @llvm.coro.id.retcon(i32 8, i32 4, ptr %buffer, ptr @prototype.signed, ptr @allocate, ptr @deallocate)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
+  br label %loop
+
+loop:
+  %n.val = phi i32 [ %n, %entry ], [ %inc, %resume ]
+  call void @print(i32 %n.val)
+  %unwind0 = call i1 (...) @llvm.coro.suspend.retcon.i1()
+  br i1 %unwind0, label %cleanup, label %resume
+
+resume:
+  %inc = add i32 %n.val, 1
+  br label %loop
+
+cleanup:
+  call i1 @llvm.coro.end(ptr %hdl, i1 0)
+  unreachable
+}
+
+; CHECK:       @prototype.signed = private constant
+; CHECK:       [[GLOBAL:@.*]] = private constant { ptr, i32, i64, i64 } { ptr [[RESUME:@.*]], i32 2, i64 0, i64 12867 }, section "llvm.ptrauth"
+
+; CHECK-LABEL: define ptr @f(ptr %buffer, i32 %n)
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[T1:%.*]] = getelementptr inbounds [[FRAME_T:%.*]], ptr %buffer, i32 0, i32 0
+; CHECK-NEXT:    store i32 %n, ptr [[T1]]
+; CHECK-NEXT:    call void @print(i32 %n)
+; CHECK-NEXT:    ret ptr [[GLOBAL]]
+
+; CHECK:      define internal ptr [[RESUME]](ptr noalias noundef nonnull align 4 dereferenceable(8) %0, i1 zeroext %1) {
+; CHECK:         [[GLOBAL]]
+
+@g.prototype.signed = private constant { ptr, i32, i64, i64 } { ptr @prototype, i32 2, i64 1, i64 8723 }, section "llvm.ptrauth"
+
+define ptr @g(ptr %buffer, i32 %n) {
+entry:
+  %id = call token @llvm.coro.id.retcon(i32 8, i32 4, ptr %buffer, ptr @g.prototype.signed, ptr @allocate, ptr @deallocate)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
+  br label %loop
+
+loop:
+  %n.val = phi i32 [ %n, %entry ], [ %inc, %resume ]
+  call void @print(i32 %n.val)
+  %unwind0 = call i1 (...) @llvm.coro.suspend.retcon.i1()
+  br i1 %unwind0, label %cleanup, label %resume
+
+resume:
+  %inc = add i32 %n.val, 1
+  br label %loop
+
+cleanup:
+  call i1 @llvm.coro.end(ptr %hdl, i1 0)
+  unreachable
+}
+
+; CHECK-LABEL: define ptr @g(ptr %buffer, i32 %n)
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[T1:%.*]] = getelementptr inbounds [[FRAME_T:%.*]], ptr %buffer, i32 0, i32 0
+; CHECK-NEXT:    store i32 %n, ptr [[T1]]
+; CHECK-NEXT:    call void @print(i32 %n)
+; CHECK-NEXT:    [[T0:%.*]] = ptrtoint ptr %buffer to i64
+; CHECK-NEXT:    [[T1:%.*]] = call i64 @llvm.ptrauth.blend(i64 [[T0]], i64 8723)
+; CHECK-NEXT:    [[T2:%.*]] = call i64 @llvm.ptrauth.sign(i64 ptrtoint (ptr [[RESUME:@.*]] to i64), i32 2, i64 [[T1]])
+; CHECK-NEXT:    [[T3:%.*]] = inttoptr i64 [[T2]] to ptr
+; CHECK-NEXT:    ret ptr [[T3]]
+
+; CHECK:      define internal ptr [[RESUME]](ptr noalias noundef nonnull align 4 dereferenceable(8) %0, i1 zeroext %1) {
+
+; CHECK: common.ret:
+; CHECK: [[RETVAL:%.*]] = phi ptr [ [[T4:%.*]], %resume ], [ null,
+; CHECK: ret ptr [[RETVAL]]
+
+; CHECK:         call void @print(i32 %inc)
+; CHECK-NEXT:    [[T0:%.*]] = ptrtoint ptr %0 to i64
+; CHECK-NEXT:    [[T1:%.*]] = call i64 @llvm.ptrauth.blend(i64 [[T0]], i64 8723)
+; CHECK-NEXT:    [[T2:%.*]] = call i64 @llvm.ptrauth.sign(i64 ptrtoint (ptr [[RESUME]] to i64), i32 2, i64 [[T1]])
+; CHECK-NEXT:    [[T3:%.*]] = inttoptr i64 [[T2]] to ptr
+
+declare noalias ptr @malloc(i64) #5
+declare void @free(ptr nocapture) #5
+
+declare token @llvm.coro.id.retcon(i32, i32, ptr, ptr, ptr, ptr)
+declare ptr @llvm.coro.begin(token, ptr)
+declare i1 @llvm.coro.suspend.retcon.i1(...)
+declare i1 @llvm.coro.end(ptr, i1)
+declare ptr @llvm.coro.prepare.retcon(ptr)
+
+declare ptr @prototype(ptr, i1 zeroext)
+
+declare noalias ptr @allocate(i32 %size)
+declare void @deallocate(ptr %ptr)
+
+declare void @print(i32)