Merge pull request swiftlang#34928 from nate-chandler/concurrency/irgen/ptrauth

[Async CC] Pointer authentication.

Author: nate-chandler

include/swift/AST/IRGenOptions.h

@@ -140,6 +140,9 @@ struct PointerAuthOptions : clang::PointerAuthOptions {
 
   /// The parent async context stored within a child async context.
   PointerAuthSchema AsyncContextParent;
+
+  /// The function to call to resume running in the parent context.
+  PointerAuthSchema AsyncContextResume;
 };
 
 enum class JITDebugArtifact : unsigned {

lib/IRGen/GenCall.cpp

@@ -1867,12 +1867,15 @@ std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
     FunctionPointer functionPointer, llvm::Value *thickContext,
     std::pair<bool, bool> values, Size initialContextSize) {
   assert(values.first || values.second);
+  assert(functionPointer.getKind() ==
+         FunctionPointer::KindTy::AsyncFunctionPointer);
   bool emitFunction = values.first;
   bool emitSize = values.second;
   // TODO: This calculation should be extracted out into standalone functions
   //       emitted on-demand per-module to improve codesize.
   switch (representation) {
   case SILFunctionTypeRepresentation::Thick: {
+    assert(!functionPointer.useStaticContextSize());
     // If the called function is thick, the size of the called function's
     // async context is not statically knowable.
     //
@@ -1941,16 +1944,25 @@ std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
     SmallVector<std::pair<llvm::BasicBlock *, llvm::Value *>, 2> sizePhiValues;
     { // thin
       IGF.Builder.emitBlock(thinBlock);
+      auto *ptr = functionPointer.getRawPointer();
+      if (auto authInfo = functionPointer.getAuthInfo()) {
+        ptr = emitPointerAuthAuth(IGF, ptr, authInfo);
+      }
+      auto *afpPtr =
+          IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
       if (emitFunction) {
-        auto *uncastFnPtr = functionPointer.getPointer(IGF);
+        llvm::Value *addrPtr = IGF.Builder.CreateStructGEP(afpPtr, 0);
+        auto *uncastFnPtr = IGF.emitLoadOfRelativePointer(
+            Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
+            /*expectedType*/ functionPointer.getFunctionType()->getPointerTo());
         auto *fnPtr = IGF.Builder.CreateBitCast(uncastFnPtr, IGF.IGM.Int8PtrTy);
+        if (auto authInfo = functionPointer.getAuthInfo()) {
+          fnPtr = emitPointerAuthSign(IGF, fnPtr, authInfo);
+        }
         fnPhiValues.push_back({thinBlock, fnPtr});
       }
       if (emitSize) {
-        auto *ptr = functionPointer.getRawPointer();
-        auto *descriptorPtr =
-            IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
-        auto *sizePtr = IGF.Builder.CreateStructGEP(descriptorPtr, 1);
+        auto *sizePtr = IGF.Builder.CreateStructGEP(afpPtr, 1);
         auto *size =
             IGF.Builder.CreateLoad(sizePtr, IGF.IGM.getPointerAlignment());
         sizePhiValues.push_back({thinBlock, size});
@@ -2026,20 +2038,34 @@ std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
   case SILFunctionTypeRepresentation::WitnessMethod:
   case SILFunctionTypeRepresentation::Closure:
   case SILFunctionTypeRepresentation::Block: {
+    auto *ptr = functionPointer.getRawPointer();
+    if (auto authInfo = functionPointer.getAuthInfo()) {
+      ptr = emitPointerAuthAuth(IGF, ptr, authInfo);
+    }
+    auto *afpPtr =
+        IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
     llvm::Value *fn = nullptr;
     if (emitFunction) {
-      fn = functionPointer.getPointer(IGF);
+      if (functionPointer.useStaticContextSize()) {
+        fn = functionPointer.getRawPointer();
+      } else {
+        llvm::Value *addrPtr = IGF.Builder.CreateStructGEP(afpPtr, 0);
+        fn = IGF.emitLoadOfRelativePointer(
+            Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
+            /*expectedType*/ functionPointer.getFunctionType()->getPointerTo());
+      }
+      if (auto authInfo = functionPointer.getAuthInfo()) {
+        fn = emitPointerAuthSign(IGF, fn, authInfo);
+      }
     }
     llvm::Value *size = nullptr;
     if (emitSize) {
       if (functionPointer.useStaticContextSize()) {
         size = llvm::ConstantInt::get(IGF.IGM.Int32Ty,
                                       initialContextSize.getValue());
-      }  else {
-        auto *ptr = functionPointer.getRawPointer();
-        auto *descriptorPtr =
-            IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
-        auto *sizePtr = IGF.Builder.CreateStructGEP(descriptorPtr, 1);
+      } else {
+        assert(!functionPointer.useStaticContextSize());
+        auto *sizePtr = IGF.Builder.CreateStructGEP(afpPtr, 1);
         size = IGF.Builder.CreateLoad(sizePtr, IGF.IGM.getPointerAlignment());
       }
     }
@@ -2336,7 +2362,8 @@ class AsyncCallEmission final : public CallEmission {
   }
   FunctionPointer getCalleeFunctionPointer() override {
     return FunctionPointer(
-        FunctionPointer::KindTy::Function, calleeFunction, PointerAuthInfo(),
+        FunctionPointer::KindTy::Function, calleeFunction,
+        CurCallee.getFunctionPointer().getAuthInfo(),
         IGF.IGM.getSignature(getCallee().getSubstFunctionType()));
   }
   SILType getParameterType(unsigned index) override {
@@ -2374,8 +2401,7 @@ class AsyncCallEmission final : public CallEmission {
           llvm::Intrinsic::coro_async_resume, {});
       auto fnVal = currentResumeFn;
       // Sign the pointer.
-      // TODO: use a distinct schema.
-      if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextParent) {
+      if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextResume) {
         Address fieldAddr =
             fieldLayout.project(IGF, this->context, /*offsets*/ llvm::None);
         auto authInfo = PointerAuthInfo::emit(
@@ -2454,6 +2480,9 @@ class AsyncCallEmission final : public CallEmission {
         Builder.CreateBitOrPointerCast(dispatchFn, IGM.Int8PtrTy));
     arguments.push_back(
         Builder.CreateBitOrPointerCast(fn.getRawPointer(), IGM.Int8PtrTy));
+    if (auto authInfo = fn.getAuthInfo()) {
+      arguments.push_back(fn.getAuthInfo().getDiscriminator());
+    }
     for (auto arg: args)
       arguments.push_back(arg);
     return IGF.emitSuspendAsyncCall(arguments);
@@ -3640,14 +3669,34 @@ llvm::Value *irgen::emitTaskCreate(
   auto layout = getAsyncContextLayout(
       IGF.IGM, taskFunctionCanSILType, taskFunctionCanSILType, subs);
 
+  CanSILFunctionType taskContinuationFunctionTy = [&]() {
+    ASTContext &ctx = IGF.IGM.IRGen.SIL.getASTContext();
+    auto extInfo =
+        ASTExtInfoBuilder()
+            .withRepresentation(FunctionTypeRepresentation::CFunctionPointer)
+            .build();
+    // FIXME: Use the appropriate signature for TaskContinuationFunction:
+    //
+    //            using TaskContinuationFunction =
+    //              SWIFT_CC(swift)
+    //              void (AsyncTask *, ExecutorRef, AsyncContext *);
+    auto ty = FunctionType::get({}, ctx.TheEmptyTupleType, extInfo);
+    return IGF.IGM.getLoweredType(ty).castTo<SILFunctionType>();
+  }();
+
   // Call the function.
   llvm::CallInst *result;
   llvm::Value *theSize, *theFunction;
+  auto taskFunctionPointer = FunctionPointer::forExplosionValue(
+      IGF, taskFunction, taskFunctionCanSILType);
   std::tie(theFunction, theSize) =
       getAsyncFunctionAndSize(IGF, SILFunctionTypeRepresentation::Thick,
-                              FunctionPointer::forExplosionValue(
-                                  IGF, taskFunction, taskFunctionCanSILType),
-                              localContextInfo);
+                              taskFunctionPointer, localContextInfo);
+  if (auto authInfo = PointerAuthInfo::forFunctionPointer(
+          IGF.IGM, taskContinuationFunctionTy)) {
+    theFunction = emitPointerAuthResign(
+        IGF, theFunction, taskFunctionPointer.getAuthInfo(), authInfo);
+  }
   theFunction = IGF.Builder.CreateBitOrPointerCast(
       theFunction, IGF.IGM.TaskContinuationFunctionPtrTy);
   theSize = IGF.Builder.CreateZExtOrBitCast(theSize, IGF.IGM.SizeTy);
@@ -4519,18 +4568,25 @@ llvm::Value *FunctionPointer::getPointer(IRGenFunction &IGF) const {
     return Value;
   case KindTy::Value::AsyncFunctionPointer: {
     if (!isFunctionPointerWithoutContext) {
+      auto *fnPtr = Value;
+      if (auto authInfo = AuthInfo) {
+        fnPtr = emitPointerAuthAuth(IGF, fnPtr, authInfo);
+      }
       auto *descriptorPtr =
-          IGF.Builder.CreateBitCast(Value, IGF.IGM.AsyncFunctionPointerPtrTy);
+          IGF.Builder.CreateBitCast(fnPtr, IGF.IGM.AsyncFunctionPointerPtrTy);
       auto *addrPtr = IGF.Builder.CreateStructGEP(descriptorPtr, 0);
-      return IGF.emitLoadOfRelativePointer(
+      auto *result = IGF.emitLoadOfRelativePointer(
           Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
           /*expectedType*/ getFunctionType()->getPointerTo());
+      if (auto authInfo = AuthInfo) {
+        result = emitPointerAuthSign(IGF, result, authInfo);
+      }
+      return result;
     } else {
       return IGF.Builder.CreateBitOrPointerCast(
           Value, getFunctionType()->getPointerTo());
     }
   }
-
   }
 }
 
@@ -4577,8 +4633,7 @@ void irgen::emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &asyncLayout,
       .loadAsCopy(IGF, returnToCallerAddr, fn);
   llvm::Value *fnVal = fn.claimNext();
 
-  // TODO: use distinct schema
-  if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextParent) {
+  if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextResume) {
     Address fieldAddr =
         returnToCallerLayout.project(IGF, contextAddr, /*offsets*/ llvm::None);
     auto authInfo = PointerAuthInfo::emit(IGF, schema, fieldAddr.getAddress(),

lib/IRGen/GenFunc.cpp

@@ -769,7 +769,7 @@ class PartialApplicationForwarderEmission {
   virtual void addArgument(llvm::Value *argValue, unsigned index) = 0;
   virtual SILParameterInfo getParameterInfo(unsigned index) = 0;
   virtual llvm::Value *getContext() = 0;
-  virtual llvm::Value *getDynamicFunctionPointer() = 0;
+  virtual llvm::Value *getDynamicFunctionPointer(PointerAuthInfo &authInfo) = 0;
   virtual llvm::Value *getDynamicFunctionContext() = 0;
   virtual void addDynamicFunctionContext(Explosion &explosion,
                                          DynamicFunctionKind kind) = 0;
@@ -931,7 +931,9 @@ class SyncPartialApplicationForwarderEmission
     return substType->getParameters()[index];
   }
   llvm::Value *getContext() override { return origParams.claimNext(); }
-  llvm::Value *getDynamicFunctionPointer() override { return args.takeLast(); }
+  llvm::Value *getDynamicFunctionPointer(PointerAuthInfo &authInfo) override {
+    return args.takeLast();
+  }
   llvm::Value *getDynamicFunctionContext() override { return args.takeLast(); }
   void addDynamicFunctionContext(Explosion &explosion,
                                  DynamicFunctionKind kind) override {
@@ -1127,15 +1129,14 @@ class AsyncPartialApplicationForwarderEmission
   llvm::Value *getContext() override {
     return loadValue(layout.getLocalContextLayout());
   }
-  llvm::Value *getDynamicFunctionPointer() override {
+  llvm::Value *getDynamicFunctionPointer(PointerAuthInfo &authInfo) override {
     assert(dynamicFunction && dynamicFunction->pointer);
     auto *context = dynamicFunction->context;
     if (!context) {
       return dynamicFunction->pointer;
     }
     auto *rawFunction = subIGF.Builder.CreateBitCast(
         dynamicFunction->pointer, origSig.getType()->getPointerTo());
-    auto authInfo = PointerAuthInfo::forFunctionPointer(IGM, origType);
     auto functionPointer =
         FunctionPointer(FunctionPointer::KindTy::AsyncFunctionPointer,
                         rawFunction, authInfo, origSig);
@@ -1719,19 +1720,19 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
 
     // Otherwise, it was the last thing we added to the layout.
 
-    // The dynamic function pointer is packed "last" into the context,
-    // and we pulled it out as an argument.  Just pop it off.
-    auto fnPtr = emission->getDynamicFunctionPointer();
-
-    // It comes out of the context as an i8*. Cast to the function type.
-    fnPtr = subIGF.Builder.CreateBitCast(fnPtr, fnTy);
-
     assert(lastCapturedFieldPtr);
     auto authInfo = PointerAuthInfo::emit(subIGF,
                             IGM.getOptions().PointerAuth.PartialApplyCapture,
                             lastCapturedFieldPtr,
                             PointerAuthEntity::Special::PartialApplyCapture);
 
+    // The dynamic function pointer is packed "last" into the context,
+    // and we pulled it out as an argument.  Just pop it off.
+    auto fnPtr = emission->getDynamicFunctionPointer(authInfo);
+
+    // It comes out of the context as an i8*. Cast to the function type.
+    fnPtr = subIGF.Builder.CreateBitCast(fnPtr, fnTy);
+
     return FunctionPointer(FunctionPointer::KindTy::Function, fnPtr, authInfo,
                            origSig);
   }();
@@ -2391,7 +2392,12 @@ llvm::Function *IRGenFunction::getOrCreateResumePrjFn() {
         auto &Builder = IGF.Builder;
         auto addr = Builder.CreateBitOrPointerCast(&(*it), IGF.IGM.Int8PtrPtrTy);
         Address callerContextAddr(addr, IGF.IGM.getPointerAlignment());
-        auto callerContext = Builder.CreateLoad(callerContextAddr);
+        llvm::Value *callerContext = Builder.CreateLoad(callerContextAddr);
+        if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextParent) {
+          auto authInfo =
+              PointerAuthInfo::emit(IGF, schema, addr, PointerAuthEntity());
+          callerContext = emitPointerAuthAuth(IGF, callerContext, authInfo);
+        }
         Builder.CreateRet(callerContext);
       },
       false /*isNoInline*/));
@@ -2412,6 +2418,10 @@ IRGenFunction::createAsyncDispatchFn(const FunctionPointer &fnPtr,
                                      ArrayRef<llvm::Type *> argTypes) {
   SmallVector<llvm::Type*, 8> argTys;
   argTys.push_back(IGM.Int8PtrTy); // Function pointer to be called.
+  auto originalAuthInfo = fnPtr.getAuthInfo();
+  if (fnPtr.getAuthInfo()) {
+    argTys.push_back(IGM.Int64Ty); // Discriminator for the function pointer.
+  }
   for (auto ty : argTypes) {
     argTys.push_back(ty);
   }
@@ -2431,13 +2441,18 @@ IRGenFunction::createAsyncDispatchFn(const FunctionPointer &fnPtr,
     IGM.DebugInfo->emitArtificialFunction(dispatchIGF, dispatch);
   auto &Builder = dispatchIGF.Builder;
   auto it = dispatchIGF.CurFn->arg_begin(), end = dispatchIGF.CurFn->arg_end();
-  llvm::Value *ptrArg = &*(it++);
+  llvm::Value *fnPtrArg = &*(it++);
+  llvm::Value *discriminatorArg = ((bool)originalAuthInfo) ? &*(it++) : nullptr;
   SmallVector<llvm::Value *, 8> callArgs;
   for (; it != end; ++it) {
     callArgs.push_back(&*it);
   }
-  ptrArg = Builder.CreateBitOrPointerCast(ptrArg, calleeFnPtrType);
-  auto callee = FunctionPointer(fnPtr.getKind(), ptrArg, fnPtr.getAuthInfo(),
+  fnPtrArg = Builder.CreateBitOrPointerCast(fnPtrArg, calleeFnPtrType);
+  PointerAuthInfo newAuthInfo =
+      ((bool)originalAuthInfo)
+          ? PointerAuthInfo(fnPtr.getAuthInfo().getKey(), discriminatorArg)
+          : originalAuthInfo;
+  auto callee = FunctionPointer(fnPtr.getKind(), fnPtrArg, newAuthInfo,
                                 fnPtr.getSignature());
   auto call = Builder.CreateCall(callee, callArgs);
   call->setTailCall();

lib/IRGen/GenOpaque.cpp

@@ -484,7 +484,7 @@ IRGenFunction::emitValueWitnessFunctionRef(SILType type,
 
   auto vwtable = emitValueWitnessTableRef(type, &metadataSlot);
   auto witness = emitLoadOfValueWitnessFunction(*this, vwtable, index);
-  setScopedLocalTypeDataForLayout(type, key, witness.getPointer(*this));
+  setScopedLocalTypeDataForLayout(type, key, witness.getRawPointer());
   if (auto &authInfo = witness.getAuthInfo()) {
     setScopedLocalTypeDataForLayout(type,
                         LocalTypeDataKind::forValueWitnessDiscriminator(index),

lib/IRGen/GenPointerAuth.cpp

@@ -71,9 +71,7 @@ llvm::Value *irgen::emitPointerAuthStrip(IRGenFunction &IGF,
 FunctionPointer irgen::emitPointerAuthResign(IRGenFunction &IGF,
                                              const FunctionPointer &fn,
                                           const PointerAuthInfo &newAuthInfo) {
-  // TODO: Handle resigning AsyncFunctionPointers.
-  assert(fn.getKind().value == FunctionPointer::KindTy::Value::Function);
-  llvm::Value *fnPtr = emitPointerAuthResign(IGF, fn.getPointer(IGF),
+  llvm::Value *fnPtr = emitPointerAuthResign(IGF, fn.getRawPointer(),
                                              fn.getAuthInfo(), newAuthInfo);
   return FunctionPointer(fn.getKind(), fnPtr, newAuthInfo, fn.getSignature());
 }

lib/IRGen/IRGen.cpp

@@ -690,8 +690,12 @@ static void setPointerAuthOptions(PointerAuthOptions &opts,
       SpecialPointerAuthDiscriminators::ResilientClassStubInitCallback);
 
   opts.AsyncContextParent =
-      PointerAuthSchema(codeKey, /*address*/ true, Discrimination::Constant,
+      PointerAuthSchema(dataKey, /*address*/ true, Discrimination::Constant,
                         SpecialPointerAuthDiscriminators::AsyncContextParent);
+
+  opts.AsyncContextResume =
+      PointerAuthSchema(codeKey, /*address*/ true, Discrimination::Constant,
+                        SpecialPointerAuthDiscriminators::AsyncContextResume);
 }
 
 std::unique_ptr<llvm::TargetMachine>

test/Concurrency/Runtime/async_task_priority_basic.swift

@@ -2,7 +2,6 @@
 // REQUIRES: executable_test
 // REQUIRES: concurrency
 // REQUIRES: libdispatch
-// XFAIL: CPU=arm64e
 
 import Dispatch
 

test/Concurrency/Runtime/basic_future.swift

@@ -3,7 +3,6 @@
 // REQUIRES: executable_test
 // REQUIRES: concurrency
 // REQUIRES: libdispatch
-// XFAIL: CPU=arm64e
 
 import Dispatch
 

test/Concurrency/Runtime/future_fibonacci.swift

@@ -3,7 +3,6 @@
 // REQUIRES: executable_test
 // REQUIRES: concurrency
 // REQUIRES: libdispatch
-// XFAIL: CPU=arm64e
 
 import Dispatch
 

test/IRGen/async/partial_apply.sil

@@ -2,7 +2,6 @@
 // RUN: %target-swift-frontend -emit-module -enable-library-evolution -emit-module-path=%t/resilient_struct.swiftmodule -module-name=resilient_struct %S/../../Inputs/resilient_struct.swift
 // RUN: %target-swift-frontend -I %t -emit-ir %s | %FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-%target-ptrsize
 
-// REQUIRES: CPU=x86_64
 // REQUIRES: concurrency
 
 import Builtin