[arm64e] Upstream signing for coroutine continuations

[varungandhi-apple]

Summary

Clients can create fake continuation prototypes to pass down signing information for pointer authentication to the coroutine-splitting pass.

Unblocks https://github.com/apple/swift/pull/30413.

[rjmccall]

Unfortunately, you can't upstream this this far yet because LLVM ptrauth hasn't been upstreamed this far. You need to upstream it to apple/llvm-project.

[varungandhi-apple]

Closing. I have a PR up against apple/llvm-project. https://github.com/apple/llvm-project/pull/979

llvm/lib/Transforms/Coroutines/CoroInstr.h

 // The helpful comment above is borrowed from llvm/IntrinsicInst.h, we keep
 // coroutine intrinsic wrappers here since they are only used by the passes in
 // the Coroutine library.
 //===----------------------------------------------------------------------===//
 
 #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"
 
 namespace llvm {
 
 /// This class represents the llvm.coro.subfn.addr instruction.
 class LLVM_LIBRARY_VISIBILITY CoroSubFnInst : public IntrinsicInst {
public:
   Value *getStorage() const {
     return getArgOperand(StorageArg);
   }
 
   /// Return the prototype for the continuation function.  The type,
   /// 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);
+  }
+
+  Optional<GlobalPtrAuthInfo> getPtrAuthInfo() const {
+    return GlobalPtrAuthInfo::analyze(getArgOperand(PrototypeArg));
   }
 
   /// Return the function to use for allocating memory.
   Function *getAllocFunction() const {
     return cast<Function>(getArgOperand(AllocArg)->stripPointerCasts());
   }
 
   /// Return the function to use for deallocating memory.

llvm/lib/Transforms/Coroutines/CoroInternal.h

struct LLVM_LIBRARY_VISIBILITY Shape {
   struct SwitchLoweringStorage {
     SwitchInst *ResumeSwitch;
     AllocaInst *PromiseAlloca;
     BasicBlock *ResumeEntryBlock;
     bool HasFinalSuspend;
   };
 
   struct RetconLoweringStorage {
+    Optional<GlobalPtrAuthInfo> ResumePtrAuthInfo;
     Function *ResumePrototype;
     Function *Alloc;
     Function *Dealloc;
     BasicBlock *ReturnBlock;
     bool IsFrameInlineInStorage;
   };
 
   union {
IntegerType *getIndexType() const {
     return cast<IntegerType>(FrameTy->getElementType(SwitchFieldIndex::Index));
   }
   ConstantInt *getIndex(uint64_t Value) const {
     return ConstantInt::get(getIndexType(), Value);
   }
 
   PointerType *getSwitchResumePointerType() const {
     assert(ABI == coro::ABI::Switch);
     assert(FrameTy && "frame type not assigned");
     return cast<PointerType>(FrameTy->getElementType(SwitchFieldIndex::Resume));
   }
 
+  Optional<GlobalPtrAuthInfo> getResumePtrAuthInfo() const {
+    switch (ABI) {
+    case coro::ABI::Switch:
+      return None;
+    case coro::ABI::Retcon:
+    case coro::ABI::RetconOnce:
+      return RetconLowering.ResumePtrAuthInfo;
+    }
+  }
+
   FunctionType *getResumeFunctionType() const {
     switch (ABI) {
     case coro::ABI::Switch: {
       auto *FnPtrTy = getSwitchResumePointerType();
       return cast<FunctionType>(FnPtrTy->getPointerElementType());
     }
     case coro::ABI::Retcon:
     case coro::ABI::RetconOnce:
struct LLVM_LIBRARY_VISIBILITY Shape {
   /// \param CG - if non-null, will be updated for the new call
   Value *emitAlloc(IRBuilder<> &Builder, Value *Size, CallGraph *CG) const;
 
   /// Deallocate memory according to the rules of the active lowering.
   ///
   /// \param CG - if non-null, will be updated for the new call
   void emitDealloc(IRBuilder<> &Builder, Value *Ptr, CallGraph *CG) const;
 
-  Shape() = default;
   explicit Shape(Function &F) { buildFrom(F); }
   void buildFrom(Function &F);
 };
 
 void buildCoroutineFrame(Function &F, Shape &Shape);
 
 } // End namespace coro.
 } // End namespace llvm
 
 #endif

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

if (!ReturnBB) {
         RetV = Builder.CreateInsertValue(RetV, CastedContinuation, 0);
         for (size_t I = 1, E = ReturnPHIs.size(); I != E; ++I)
           RetV = Builder.CreateInsertValue(RetV, ReturnPHIs[I], I);
       }
 
       Builder.CreateRet(RetV);
     }
 
+    // Sign the continuation value if requested.
+    Constant *ContinuationValue = Continuation;
+    if (auto PtrAuthInfo = Shape.getResumePtrAuthInfo()) {
+      ContinuationValue =
+          PtrAuthInfo->createWithSameSchema(*F.getParent(), ContinuationValue);
+    }
+
     // 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);
     assert(NextPHIIndex == ReturnPHIs.size());
   }
 
   assert(Clones.size() == Shape.CoroSuspends.size());
   for (size_t i = 0, e = Shape.CoroSuspends.size(); i != e; ++i) {

llvm/lib/Transforms/Coroutines/Coroutines.cpp

#endif
   case Intrinsic::coro_id_retcon_once: {
     auto ContinuationId = cast<AnyCoroIdRetconInst>(Id);
     ContinuationId->checkWellFormed();
     this->ABI = (IdIntrinsic == Intrinsic::coro_id_retcon
                   ? coro::ABI::Retcon
                   : coro::ABI::RetconOnce);
     auto Prototype = ContinuationId->getPrototype();
     this->RetconLowering.ResumePrototype = Prototype;
+    this->RetconLowering.ResumePtrAuthInfo = ContinuationId->getPtrAuthInfo();
     this->RetconLowering.Alloc = ContinuationId->getAllocFunction();
     this->RetconLowering.Dealloc = ContinuationId->getDeallocFunction();
     this->RetconLowering.ReturnBlock = nullptr;
     this->RetconLowering.IsFrameInlineInStorage = false;
 
+    if (RetconLowering.ResumePtrAuthInfo &&
+        RetconLowering.ResumePtrAuthInfo->hasAddressDiversity()) {
+#ifndef NDEBUG
+      RetconLowering.ResumePtrAuthInfo->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();
     auto ResumeTys = getRetconResumeTypes();
 
     for (auto AnySuspend : CoroSuspends) {
       auto Suspend = dyn_cast<CoroSuspendRetconInst>(AnySuspend);
       if (!Suspend) {
#ifndef NDEBUG
   }
 #endif
   report_fatal_error(Reason);
 }
 
 /// 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);
 
   auto FT = F->getFunctionType();
 
   if (isa<CoroIdRetconInst>(I)) {
     bool ResultOkay;

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

+; RUN: opt -coro-early -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 }>
+
[e-mail] = private constant { i8*, i32, i64, i64 } { i8* bitcast (i8* (i8*, i1)* @prototype to i8*), i32 2, i64 0, i64 12867 }, section "llvm.ptrauth"
+
+define i8* @f(i8* %buffer, i32 %n) {
+entry:
+  %id = call token @llvm.coro.id.retcon(i32 8, i32 4, i8* %buffer, i8* bitcast ({ i8*, i32, i64, i64 }* @prototype.signed to i8*), i8* bitcast (i8* (i32)* @allocate to i8*), i8* bitcast (void (i8*)* @deallocate to i8*))
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* 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(i8* %hdl, i1 0)
+  unreachable
+}
+
+; CHECK:       @prototype.signed = private constant
+; CHECK:       [[GLOBAL:@.*]] = private constant { i8*, i32, i64, i64 } { i8* bitcast (i8* (i8*, i1)* [[RESUME:@.*]] to i8*), i32 2, i64 0, i64 12867 }, section "llvm.ptrauth"
+
+; CHECK-LABEL: define i8* @f(i8* %buffer, i32 %n)
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[T0:%.*]] = bitcast i8* %buffer to [[FRAME_T:%.*]]*
+; CHECK-NEXT:    [[T1:%.*]] = getelementptr inbounds [[FRAME_T]], [[FRAME_T]]* [[T0]], i32 0, i32 0
+; CHECK-NEXT:    store i32 %n, i32* [[T1]]
+; CHECK-NEXT:    call void @print(i32 %n)
+; CHECK-NEXT:    ret i8* bitcast ({ i8*, i32, i64, i64 }* [[GLOBAL]] to i8*)
+
+; CHECK:      define internal i8* [[RESUME]](i8* noalias nonnull, i1 zeroext) {
+; CHECK:         bitcast ({ i8*, i32, i64, i64 }* [[GLOBAL]] to
+
+declare noalias i8* @malloc(i64) #5
+declare void @free(i8* nocapture) #5
+
+declare token @llvm.coro.id.retcon(i32, i32, i8*, i8*, i8*, i8*)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.suspend.retcon.i1(...)
+declare i1 @llvm.coro.end(i8*, i1)
+declare i8* @llvm.coro.prepare.retcon(i8*)
+
+declare i8* @prototype(i8*, i1 zeroext)
+
+declare noalias i8* @allocate(i32 %size)
+declare void @deallocate(i8* %ptr)
+
+declare void @print(i32)