Merge pull request #36729 from nate-chandler/rdar76118522

[ptrauth] Sign AsyncFunctionPointers as data.

Author: nate-chandler

include/swift/AST/IRGenOptions.h

@@ -142,6 +142,23 @@ struct PointerAuthOptions : clang::PointerAuthOptions {
   /// Resilient class stub initializer callbacks.
   PointerAuthSchema ResilientClassStubInitCallbacks;
 
+  /// Like SwiftFunctionPointers but for use with AsyncFunctionPointer values.
+  PointerAuthSchema AsyncSwiftFunctionPointers;
+
+  /// Like SwiftClassMethods but for use with AsyncFunctionPointer values.
+  PointerAuthSchema AsyncSwiftClassMethods;
+
+  /// Like ProtocolWitnesses but for use with AsyncFunctionPointer values.
+  PointerAuthSchema AsyncProtocolWitnesses;
+
+  /// Like SwiftClassMethodPointers but for use with AsyncFunctionPointer
+  /// values.
+  PointerAuthSchema AsyncSwiftClassMethodPointers;
+
+  /// Like SwiftDynamicReplacements but for use with AsyncFunctionPointer
+  /// values.
+  PointerAuthSchema AsyncSwiftDynamicReplacements;
+
   /// The parent async context stored within a child async context.
   PointerAuthSchema AsyncContextParent;
 

include/swift/IRGen/Linking.h

@@ -456,6 +456,11 @@ class LinkEntity {
     /// An async function pointer to a partial apply forwarder.
     /// The pointer is the llvm::Function* for a partial apply forwarder.
     PartialApplyForwarderAsyncFunctionPointer,
+
+    /// An async function pointer to a function which is known to exist whose
+    /// name is known.
+    /// The pointer is a const char* of the name.
+    KnownAsyncFunctionPointer,
   };
   friend struct llvm::DenseMapInfo<LinkEntity>;
 
@@ -1197,6 +1202,15 @@ class LinkEntity {
     return entity;
   }
 
+  static LinkEntity forKnownAsyncFunctionPointer(const char *name) {
+    LinkEntity entity;
+    entity.Pointer = const_cast<char *>(name);
+    entity.SecondaryPointer = nullptr;
+    entity.Data =
+        LINKENTITY_SET_FIELD(Kind, unsigned(Kind::KnownAsyncFunctionPointer));
+    return entity;
+  }
+
   LinkEntity getUnderlyingEntityForAsyncFunctionPointer() const {
     LinkEntity entity;
     entity.Pointer = Pointer;

include/swift/Runtime/Config.h

@@ -303,9 +303,39 @@ static inline void swift_ptrauth_copy(T *dest, const T *src, unsigned extra) {
 #endif
 }
 
-/// Initialize the destination with an address-discriminated signed pointer.
-/// This does not authenticate the source value, so be careful about how
-/// you construct it.
+/// Copy an address-discriminated signed data pointer from the source
+/// to the destination.
+template <class T>
+SWIFT_RUNTIME_ATTRIBUTE_ALWAYS_INLINE
+static inline void swift_ptrauth_copy_data(T *dest, const T *src,
+                                           unsigned extra) {
+#if SWIFT_PTRAUTH
+  *dest = ptrauth_auth_and_resign(*src,
+                                  ptrauth_key_process_independent_data,
+                                  ptrauth_blend_discriminator(src, extra),
+                                  ptrauth_key_process_independent_data,
+                                  ptrauth_blend_discriminator(dest, extra));
+#else
+  *dest = *src;
+#endif
+}
+
+/// Copy an address-discriminated signed pointer from the source
+/// to the destination.
+template <class T>
+SWIFT_RUNTIME_ATTRIBUTE_ALWAYS_INLINE static inline void
+swift_ptrauth_copy_code_or_data(T *dest, const T *src, unsigned extra,
+                                bool isCode) {
+  if (isCode) {
+    return swift_ptrauth_copy(dest, src, extra);
+  } else {
+    return swift_ptrauth_copy_data(dest, src, extra);
+  }
+}
+
+/// Initialize the destination with an address-discriminated signed
+/// function pointer.  This does not authenticate the source value, so be
+/// careful about how you construct it.
 template <class T>
 SWIFT_RUNTIME_ATTRIBUTE_ALWAYS_INLINE
 static inline void swift_ptrauth_init(T *dest, T value, unsigned extra) {
@@ -319,6 +349,35 @@ static inline void swift_ptrauth_init(T *dest, T value, unsigned extra) {
 #endif
 }
 
+/// Initialize the destination with an address-discriminated signed
+/// data pointer.  This does not authenticate the source value, so be
+/// careful about how you construct it.
+template <class T>
+SWIFT_RUNTIME_ATTRIBUTE_ALWAYS_INLINE
+static inline void swift_ptrauth_init_data(T *dest, T value, unsigned extra) {
+  // FIXME: assert that T is not a function-pointer type?
+#if SWIFT_PTRAUTH
+  *dest = ptrauth_sign_unauthenticated(value,
+                                  ptrauth_key_process_independent_data,
+                                  ptrauth_blend_discriminator(dest, extra));
+#else
+  *dest = value;
+#endif
+}
+
+/// Initialize the destination with an address-discriminated signed
+/// pointer.  This does not authenticate the source value, so be
+/// careful about how you construct it.
+template <class T>
+SWIFT_RUNTIME_ATTRIBUTE_ALWAYS_INLINE static inline void
+swift_ptrauth_init_code_or_data(T *dest, T value, unsigned extra, bool isCode) {
+  if (isCode) {
+    return swift_ptrauth_init(dest, value, extra);
+  } else {
+    return swift_ptrauth_init_data(dest, value, extra);
+  }
+}
+
 template <typename T>
 SWIFT_RUNTIME_ATTRIBUTE_ALWAYS_INLINE
 static inline T swift_auth_data_non_address(T value, unsigned extra) {

lib/IRGen/Callee.h

@@ -101,6 +101,36 @@ namespace irgen {
       assert(isSigned());
       return Key;
     }
+    bool hasCodeKey() const {
+      assert(isSigned());
+      return (getKey() == (unsigned)PointerAuthSchema::ARM8_3Key::ASIA) ||
+             (getKey() == (unsigned)PointerAuthSchema::ARM8_3Key::ASIB);
+    }
+    bool hasDataKey() const {
+      assert(isSigned());
+      return (getKey() == (unsigned)PointerAuthSchema::ARM8_3Key::ASDA) ||
+             (getKey() == (unsigned)PointerAuthSchema::ARM8_3Key::ASDB);
+    }
+    bool getCorrespondingCodeKey() const {
+      assert(hasDataKey());
+      switch (getKey()) {
+      case (unsigned)PointerAuthSchema::ARM8_3Key::ASDA:
+        return (unsigned)PointerAuthSchema::ARM8_3Key::ASIA;
+      case (unsigned)PointerAuthSchema::ARM8_3Key::ASDB:
+        return (unsigned)PointerAuthSchema::ARM8_3Key::ASIB;
+      }
+      llvm_unreachable("unhandled case");
+    }
+    bool getCorrespondingDataKey() const {
+      assert(hasCodeKey());
+      switch (getKey()) {
+      case (unsigned)PointerAuthSchema::ARM8_3Key::ASIA:
+        return (unsigned)PointerAuthSchema::ARM8_3Key::ASDA;
+      case (unsigned)PointerAuthSchema::ARM8_3Key::ASIB:
+        return (unsigned)PointerAuthSchema::ARM8_3Key::ASDB;
+      }
+      llvm_unreachable("unhandled case");
+    }
     llvm::Value *getDiscriminator() const {
       assert(isSigned());
       return Discriminator;
@@ -219,6 +249,13 @@ namespace irgen {
       // The function pointer should have function type.
       assert(value->getType()->getPointerElementType()->isFunctionTy());
       // TODO: maybe assert similarity to signature.getType()?
+      if (authInfo) {
+        if (kind == Kind::Function) {
+          assert(authInfo.hasCodeKey());
+        } else {
+          assert(authInfo.hasDataKey());
+        }
+      }
     }
 
     // Temporary only!
@@ -295,7 +332,7 @@ namespace irgen {
     llvm::Value *getExplosionValue(IRGenFunction &IGF,
                                    CanSILFunctionType fnType) const;
 
-    /// Form a FunctionPointer whose KindTy is ::Function.
+    /// Form a FunctionPointer whose Kind is ::Function.
     FunctionPointer getAsFunction(IRGenFunction &IGF) const;
 
     bool useStaticContextSize() const {

lib/IRGen/GenCall.cpp

@@ -2002,7 +2002,9 @@ std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
           /*expectedType*/ functionPointer.getFunctionType()->getPointerTo());
     }
     if (auto authInfo = functionPointer.getAuthInfo()) {
-      fn = emitPointerAuthSign(IGF, fn, authInfo);
+      auto newAuthInfo = PointerAuthInfo(authInfo.getCorrespondingCodeKey(),
+                                         authInfo.getDiscriminator());
+      fn = emitPointerAuthSign(IGF, fn, newAuthInfo);
     }
   }
   llvm::Value *size = nullptr;
@@ -2330,9 +2332,13 @@ class AsyncCallEmission final : public CallEmission {
   }
 
   FunctionPointer getCalleeFunctionPointer() override {
+    PointerAuthInfo newAuthInfo;
+    if (auto authInfo = CurCallee.getFunctionPointer().getAuthInfo()) {
+      newAuthInfo = PointerAuthInfo(authInfo.getCorrespondingCodeKey(),
+                                    authInfo.getDiscriminator());
+    }
     return FunctionPointer(
-        FunctionPointer::Kind::Function, calleeFunction,
-        CurCallee.getFunctionPointer().getAuthInfo(),
+        FunctionPointer::Kind::Function, calleeFunction, newAuthInfo,
         Signature::forAsyncAwait(IGF.IGM, getCallee().getOrigFunctionType()));
   }
 
@@ -4695,7 +4701,9 @@ llvm::Value *FunctionPointer::getPointer(IRGenFunction &IGF) const {
         Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
         /*expectedType*/ getFunctionType()->getPointerTo());
     if (auto authInfo = AuthInfo) {
-      result = emitPointerAuthSign(IGF, result, authInfo);
+      auto newAuthInfo = PointerAuthInfo(authInfo.getCorrespondingCodeKey(),
+                                         authInfo.getDiscriminator());
+      result = emitPointerAuthSign(IGF, result, newAuthInfo);
     }
     return result;
   }
@@ -4732,7 +4740,19 @@ FunctionPointer::getExplosionValue(IRGenFunction &IGF,
 }
 
 FunctionPointer FunctionPointer::getAsFunction(IRGenFunction &IGF) const {
-  return FunctionPointer(Kind::Function, getPointer(IGF), AuthInfo, Sig);
+  switch (getBasicKind()) {
+  case FunctionPointer::BasicKind::Function:
+    return *this;
+  case FunctionPointer::BasicKind::AsyncFunctionPointer: {
+    auto authInfo = AuthInfo;
+    if (authInfo) {
+      authInfo = PointerAuthInfo(AuthInfo.getCorrespondingCodeKey(),
+                                 AuthInfo.getDiscriminator());
+    }
+    return FunctionPointer(Kind::Function, getPointer(IGF), authInfo, Sig);
+  }
+  }
+  llvm_unreachable("unhandled case");
 }
 
 void irgen::emitAsyncReturn(

lib/IRGen/GenClass.cpp

@@ -2643,7 +2643,9 @@ FunctionPointer irgen::emitVirtualMethodValue(IRGenFunction &IGF,
                                         signature.getType()->getPointerTo(),
                                         IGF.IGM.getPointerAlignment());
     auto fnPtr = IGF.emitInvariantLoad(slot);
-    auto &schema = IGF.getOptions().PointerAuth.SwiftClassMethods;
+    auto &schema = methodType->isAsync()
+                       ? IGF.getOptions().PointerAuth.AsyncSwiftClassMethods
+                       : IGF.getOptions().PointerAuth.SwiftClassMethods;
     auto authInfo =
       PointerAuthInfo::emit(IGF, schema, slot.getAddress(), method);
     return FunctionPointer(methodType, fnPtr, authInfo, signature);

lib/IRGen/GenDecl.cpp

@@ -2573,7 +2573,9 @@ void IRGenModule::createReplaceableProlog(IRGenFunction &IGF, SILFunction *f) {
   auto *FnAddr = llvm::ConstantExpr::getPointerBitCastOrAddrSpaceCast(
       IGF.CurFn, FunctionPtrTy);
 
-  auto &schema = getOptions().PointerAuth.SwiftDynamicReplacements;
+  auto &schema = f->isAsync()
+                     ? getOptions().PointerAuth.AsyncSwiftDynamicReplacements
+                     : getOptions().PointerAuth.SwiftDynamicReplacements;
   llvm::Value *ReplFn = nullptr, *hasReplFn = nullptr;
 
   if (UseBasicDynamicReplacement) {

lib/IRGen/GenMeta.cpp

@@ -262,7 +262,10 @@ static Flags getMethodDescriptorFlags(ValueDecl *fn) {
     }
     llvm_unreachable("bad kind");
   }();
-  return Flags(kind).withIsInstance(!fn->isStatic());
+  bool hasAsync = false;
+  if (auto *afd = dyn_cast<AbstractFunctionDecl>(fn))
+    hasAsync = afd->hasAsync();
+  return Flags(kind).withIsInstance(!fn->isStatic()).withIsAsync(hasAsync);
 }
 
 static void buildMethodDescriptorFields(IRGenModule &IGM,
@@ -279,7 +282,9 @@ static void buildMethodDescriptorFields(IRGenModule &IGM,
     flags = flags.withIsDynamic(true);
 
   // Include the pointer-auth discriminator.
-  if (auto &schema = IGM.getOptions().PointerAuth.SwiftClassMethods) {
+  if (auto &schema = func->hasAsync()
+                         ? IGM.getOptions().PointerAuth.AsyncSwiftClassMethods
+                         : IGM.getOptions().PointerAuth.SwiftClassMethods) {
     auto discriminator =
       PointerAuthInfo::getOtherDiscriminator(IGM, schema, fn);
     flags = flags.withExtraDiscriminator(discriminator->getZExtValue());
@@ -3159,6 +3164,7 @@ namespace {
       // Find the vtable entry.
       assert(VTable && "no vtable?!");
       auto entry = VTable->getEntry(IGM.getSILModule(), fn);
+      auto *afd = cast<AbstractFunctionDecl>(fn.getDecl());
 
       // The class is fragile. Emit a direct reference to the vtable entry.
       llvm::Constant *ptr;
@@ -3172,11 +3178,19 @@ namespace {
       } else {
         // The method is removed by dead method elimination.
         // It should be never called. We add a pointer to an error function.
-        ptr = llvm::ConstantExpr::getBitCast(IGM.getDeletedMethodErrorFn(),
-                                             IGM.FunctionPtrTy);
+        if (afd->hasAsync()) {
+          ptr = llvm::ConstantExpr::getBitCast(
+              IGM.getDeletedAsyncMethodErrorAsyncFunctionPointer(),
+              IGM.FunctionPtrTy);
+        } else {
+          ptr = llvm::ConstantExpr::getBitCast(IGM.getDeletedMethodErrorFn(),
+                                               IGM.FunctionPtrTy);
+        }
       }
 
-      auto &schema = IGM.getOptions().PointerAuth.SwiftClassMethods;
+      PointerAuthSchema schema =
+          afd->hasAsync() ? IGM.getOptions().PointerAuth.AsyncSwiftClassMethods
+                          : IGM.getOptions().PointerAuth.SwiftClassMethods;
       B.addSignedPointer(ptr, schema, fn);
     }
 

lib/IRGen/GenPointerAuth.cpp

@@ -198,6 +198,10 @@ static const PointerAuthSchema &getFunctionPointerSchema(IRGenModule &IGM,
   case SILFunctionTypeRepresentation::Method:
   case SILFunctionTypeRepresentation::WitnessMethod:
   case SILFunctionTypeRepresentation::Closure:
+    if (fnType->isAsync()) {
+      return options.AsyncSwiftFunctionPointers;
+    }
+
     return options.SwiftFunctionPointers;
 
   case SILFunctionTypeRepresentation::ObjCMethod:

lib/IRGen/GenProto.cpp

@@ -1333,6 +1333,8 @@ class AccessorConformanceInfo : public ConformanceInfo {
 #endif
 
       SILFunction *Func = entry.getMethodWitness().Witness;
+      auto *afd = cast<AbstractFunctionDecl>(
+          entry.getMethodWitness().Requirement.getDecl());
       llvm::Constant *witness = nullptr;
       if (Func) {
         if (Func->isAsync()) {
@@ -1343,11 +1345,20 @@ class AccessorConformanceInfo : public ConformanceInfo {
       } else {
         // The method is removed by dead method elimination.
         // It should be never called. We add a pointer to an error function.
-        witness = IGM.getDeletedMethodErrorFn();
+        if (afd->hasAsync()) {
+          witness = llvm::ConstantExpr::getBitCast(
+              IGM.getDeletedAsyncMethodErrorAsyncFunctionPointer(),
+              IGM.FunctionPtrTy);
+        } else {
+          witness = llvm::ConstantExpr::getBitCast(
+              IGM.getDeletedMethodErrorFn(), IGM.FunctionPtrTy);
+        }
       }
       witness = llvm::ConstantExpr::getBitCast(witness, IGM.Int8PtrTy);
 
-      auto &schema = IGM.getOptions().PointerAuth.ProtocolWitnesses;
+      PointerAuthSchema schema =
+          afd->hasAsync() ? IGM.getOptions().PointerAuth.AsyncProtocolWitnesses
+                          : IGM.getOptions().PointerAuth.ProtocolWitnesses;
       Table.addSignedPointer(witness, schema, requirement);
       return;
     }
@@ -3353,7 +3364,9 @@ FunctionPointer irgen::emitWitnessMethodValue(IRGenFunction &IGF,
   witnessFnPtr = IGF.Builder.CreateBitCast(witnessFnPtr,
                                            signature.getType()->getPointerTo());
 
-  auto &schema = IGF.getOptions().PointerAuth.ProtocolWitnesses;
+  auto &schema = fnType->isAsync()
+                     ? IGF.getOptions().PointerAuth.AsyncProtocolWitnesses
+                     : IGF.getOptions().PointerAuth.ProtocolWitnesses;
   auto authInfo = PointerAuthInfo::emit(IGF, schema, slot, member);
 
   return FunctionPointer(fnType, witnessFnPtr, authInfo, signature);

lib/IRGen/GenThunk.cpp

@@ -472,7 +472,10 @@ void IRGenModule::emitMethodLookupFunction(ClassDecl *classDecl) {
                                                    NotForDefinition);
       // Sign using the discriminator we would include in the method
       // descriptor.
-      if (auto &schema = IGM.getOptions().PointerAuth.SwiftClassMethods) {
+      if (auto &schema =
+              entry->getImplementation()->getLoweredFunctionType()->isAsync()
+                  ? IGM.getOptions().PointerAuth.AsyncSwiftClassMethods
+                  : IGM.getOptions().PointerAuth.SwiftClassMethods) {
         auto discriminator =
           PointerAuthInfo::getOtherDiscriminator(IGM, schema, method);