Merge pull request #32791 from nate-chandler/generic-metadata-prespecialization-components/consider-conformances-ptrauth

[metadata prespecialization] Ptrauth for compared protocol conformances.

Author: nate-chandler

include/swift/ABI/MetadataValues.h

@@ -1118,6 +1118,9 @@ namespace SpecialPointerAuthDiscriminators {
   /// Runtime function variables exported by the runtime.
   const uint16_t RuntimeFunctionEntry = 0x625b;
 
+  /// Protocol conformance descriptors.
+  const uint16_t ProtocolConformanceDescriptor = 0xc6eb;
+
   /// Value witness functions.
   const uint16_t InitializeBufferWithCopyOfBuffer = 0xda4a;
   const uint16_t Destroy = 0x04f8;

include/swift/AST/IRGenOptions.h

@@ -123,6 +123,12 @@ struct PointerAuthOptions : clang::PointerAuthOptions {
   /// Type descriptor data pointers when passed as arguments.
   PointerAuthSchema TypeDescriptorsAsArguments;
 
+  /// Protocol conformance descriptors.
+  PointerAuthSchema ProtocolConformanceDescriptors;
+
+  /// Protocol conformance descriptors when passed as arguments.
+  PointerAuthSchema ProtocolConformanceDescriptorsAsArguments;
+
   /// Resumption functions from yield-once coroutines.
   PointerAuthSchema YieldOnceResumeFunctions;
 

lib/IRGen/GenPointerAuth.cpp

@@ -309,6 +309,9 @@ PointerAuthEntity::getDeclDiscriminator(IRGenModule &IGM) const {
       case Special::TypeDescriptor:
       case Special::TypeDescriptorAsArgument:
         return SpecialPointerAuthDiscriminators::TypeDescriptor;
+      case Special::ProtocolConformanceDescriptor:
+      case Special::ProtocolConformanceDescriptorAsArgument:
+        return SpecialPointerAuthDiscriminators::ProtocolConformanceDescriptor;
       case Special::PartialApplyCapture:
         return PointerAuthDiscriminator_PartialApplyCapture;
       case Special::KeyPathDestroy:

lib/IRGen/GenPointerAuth.h

@@ -62,6 +62,8 @@ class PointerAuthEntity {
     KeyPathInitializer,
     KeyPathMetadataAccessor,
     DynamicReplacementKey,
+    ProtocolConformanceDescriptor,
+    ProtocolConformanceDescriptorAsArgument,
   };
 
 private:

lib/IRGen/IRGen.cpp

@@ -656,6 +656,11 @@ static void setPointerAuthOptions(PointerAuthOptions &opts,
   opts.SwiftDynamicReplacementKeys =
     PointerAuthSchema(dataKey, /*address*/ true, Discrimination::Decl);
 
+  opts.ProtocolConformanceDescriptors =
+      PointerAuthSchema(dataKey, /*address*/ true, Discrimination::Decl);
+  opts.ProtocolConformanceDescriptorsAsArguments =
+      PointerAuthSchema(dataKey, /*address*/ false, Discrimination::Decl);
+
   // Coroutine resumption functions are never stored globally in the ABI,
   // so we can do some things that aren't normally okay to do.  However,
   // we can't use ASIB because that would break ARM64 interoperation.

lib/IRGen/MetadataRequest.cpp

@@ -1971,7 +1971,7 @@ static void emitCanonicalSpecializationsForGenericTypeMetadataAccessFunction(
           } else {
             RootProtocolConformance *rootConformance =
                 concreteConformance->getRootConformance();
-            auto *expectedDescriptor =
+            llvm::Value *expectedDescriptor =
                 IGF.IGM.getAddrOfProtocolConformanceDescriptor(rootConformance);
             auto *witnessTable = valueAtIndex(requirementIndex);
             auto *witnessBuffer =
@@ -1982,6 +1982,32 @@ static void emitCanonicalSpecializationsForGenericTypeMetadataAccessFunction(
                 uncastProvidedDescriptor,
                 IGM.ProtocolConformanceDescriptorPtrTy);
 
+            // Auth the stored descriptor.
+            auto storedScheme =
+                IGF.IGM.getOptions().PointerAuth.ProtocolConformanceDescriptors;
+            if (storedScheme) {
+              auto authInfo = PointerAuthInfo::emit(
+                  IGF, storedScheme, witnessTable,
+                  PointerAuthEntity::Special::ProtocolConformanceDescriptor);
+              providedDescriptor =
+                  emitPointerAuthAuth(IGF, providedDescriptor, authInfo);
+            }
+
+            // Sign the descriptors.
+            auto argScheme =
+                IGF.IGM.getOptions()
+                    .PointerAuth.ProtocolConformanceDescriptorsAsArguments;
+            if (argScheme) {
+              auto authInfo = PointerAuthInfo::emit(
+                  IGF, argScheme, nullptr,
+                  PointerAuthEntity::Special::
+                      ProtocolConformanceDescriptorAsArgument);
+              expectedDescriptor =
+                  emitPointerAuthSign(IGF, expectedDescriptor, authInfo);
+              providedDescriptor =
+                  emitPointerAuthSign(IGF, providedDescriptor, authInfo);
+            }
+
             auto *call = IGF.Builder.CreateCall(
                 IGF.IGM.getCompareProtocolConformanceDescriptorsFn(),
                 {providedDescriptor, expectedDescriptor});

stdlib/public/runtime/Metadata.cpp

@@ -4996,6 +4996,11 @@ const WitnessTable *swift::swift_getAssociatedConformanceWitness(
 bool swift::swift_compareProtocolConformanceDescriptors(
     const ProtocolConformanceDescriptor *lhs,
     const ProtocolConformanceDescriptor *rhs) {
+  lhs = swift_auth_data_non_address(
+      lhs, SpecialPointerAuthDiscriminators::ProtocolConformanceDescriptor);
+  rhs = swift_auth_data_non_address(
+      rhs, SpecialPointerAuthDiscriminators::ProtocolConformanceDescriptor);
+
   return MetadataCacheKey::compareProtocolConformanceDescriptors(lhs, rhs) == 0;
 }