Support for address discriminated pointers

Author: meg-gupta

include/swift/AST/Decl.h

@@ -613,13 +613,16 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
     IsActor : 1
   );
 
-  SWIFT_INLINE_BITFIELD(
-      StructDecl, NominalTypeDecl, 1 + 1,
-      /// True if this struct has storage for fields that aren't accessible in
-      /// Swift.
-      HasUnreferenceableStorage : 1,
-      /// True if this struct is imported from C++ and does not have trivial value witness functions.
-      IsCxxNonTrivial : 1);
+  SWIFT_INLINE_BITFIELD(StructDecl, NominalTypeDecl, 1 + 1 + 1,
+                        /// True if this struct has storage for fields that
+                        /// aren't accessible in Swift.
+                        HasUnreferenceableStorage : 1,
+                        /// True if this struct is imported from C++ and does
+                        /// not have trivial value witness functions.
+                        IsCxxNonTrivial : 1,
+                        /// True if this struct is imported from C and has
+                        /// address diversified ptrauth qualified field.
+                        IsNonTrivialPtrAuth : 1);
 
   SWIFT_INLINE_BITFIELD(EnumDecl, NominalTypeDecl, 2+1,
     /// True if the enum has cases and at least one case has associated values.
@@ -4267,6 +4270,14 @@ class StructDecl final : public NominalTypeDecl {
 
   void setIsCxxNonTrivial(bool v) { Bits.StructDecl.IsCxxNonTrivial = v; }
 
+  bool isNonTrivialPtrAuth() const {
+    return Bits.StructDecl.IsNonTrivialPtrAuth;
+  }
+
+  void setHasNonTrivialPtrAuth(bool v) {
+    Bits.StructDecl.IsNonTrivialPtrAuth = v;
+  }
+
   Type getTemplateInstantiationType() const { return TemplateInstantiationType; }
   void setTemplateInstantiationType(Type t) { TemplateInstantiationType = t; }
 };

lib/ClangImporter/ImportDecl.cpp

@@ -2024,26 +2024,45 @@ namespace {
         return nullptr;
       }
 
+      auto isNonTrivialDueToAddressDiversifiedPtrAuth =
+          [](const clang::RecordDecl *decl) {
+            for (auto *field : decl->fields()) {
+              if (!field->getType().isNonTrivialToPrimitiveCopy()) {
+                continue;
+              }
+              if (field->getType().isNonTrivialToPrimitiveCopy() !=
+                  clang::QualType::PCK_PtrAuth) {
+                return false;
+              }
+            }
+            return true;
+          };
+
+      bool isNonTrivialPtrAuth = false;
       // FIXME: We should actually support strong ARC references and similar in
       // C structs. That'll require some SIL and IRGen work, though.
       if (decl->isNonTrivialToPrimitiveCopy() ||
           decl->isNonTrivialToPrimitiveDestroy()) {
-        // Note that there is a third predicate related to these,
-        // isNonTrivialToPrimitiveDefaultInitialize. That one's not important
-        // for us because Swift never "trivially default-initializes" a struct
-        // (i.e. uses whatever bits were lying around as an initial value).
-
-        // FIXME: It would be nice to instead import the declaration but mark
-        // it as unavailable, but then it might get used as a type for an
-        // imported function and the developer would be able to use it without
-        // referencing the name, which would sidestep our availability
-        // diagnostics.
-        Impl.addImportDiagnostic(
-            decl, Diagnostic(
-                      diag::record_non_trivial_copy_destroy,
-                      Impl.SwiftContext.AllocateCopy(decl->getNameAsString())),
-            decl->getLocation());
-        return nullptr;
+        isNonTrivialPtrAuth = isNonTrivialDueToAddressDiversifiedPtrAuth(decl);
+        if (!isNonTrivialPtrAuth) {
+          // Note that there is a third predicate related to these,
+          // isNonTrivialToPrimitiveDefaultInitialize. That one's not important
+          // for us because Swift never "trivially default-initializes" a struct
+          // (i.e. uses whatever bits were lying around as an initial value).
+
+          // FIXME: It would be nice to instead import the declaration but mark
+          // it as unavailable, but then it might get used as a type for an
+          // imported function and the developer would be able to use it without
+          // referencing the name, which would sidestep our availability
+          // diagnostics.
+          Impl.addImportDiagnostic(
+              decl,
+              Diagnostic(
+                  diag::record_non_trivial_copy_destroy,
+                  Impl.SwiftContext.AllocateCopy(decl->getNameAsString())),
+              decl->getLocation());
+          return nullptr;
+        }
       }
 
       // Import the name.
@@ -2346,8 +2365,12 @@ namespace {
         }
       }
 
-      if (auto structResult = dyn_cast<StructDecl>(result))
+      if (auto structResult = dyn_cast<StructDecl>(result)) {
         structResult->setHasUnreferenceableStorage(hasUnreferenceableStorage);
+        if (isNonTrivialPtrAuth) {
+          structResult->setHasNonTrivialPtrAuth(true);
+        }
+      }
 
       if (cxxRecordDecl) {
         if (auto structResult = dyn_cast<StructDecl>(result))

lib/IRGen/Callee.h

@@ -78,8 +78,9 @@ namespace irgen {
                                 llvm::Value *storageAddress,
                                 const PointerAuthEntity &entity);
 
-    static PointerAuthInfo emit(IRGenModule &IGM,
-                                clang::PointerAuthQualifier pointerAuthQual);
+    static PointerAuthInfo emit(IRGenFunction &IGF,
+                                clang::PointerAuthQualifier pointerAuthQual,
+                                llvm::Value *storageAddress);
 
     static PointerAuthInfo forFunctionPointer(IRGenModule &IGM,
                                               CanSILFunctionType fnType);

lib/IRGen/GenPointerAuth.cpp

@@ -239,12 +239,30 @@ PointerAuthInfo PointerAuthInfo::emit(IRGenFunction &IGF,
 }
 
 PointerAuthInfo
-PointerAuthInfo::emit(IRGenModule &IGM,
-                      clang::PointerAuthQualifier pointerAuthQual) {
-  return PointerAuthInfo(
-      pointerAuthQual.getKey(),
-      llvm::ConstantInt::get(IGM.Int64Ty,
-                             pointerAuthQual.getExtraDiscriminator()));
+PointerAuthInfo::emit(IRGenFunction &IGF,
+                      clang::PointerAuthQualifier pointerAuthQual,
+                      llvm::Value *storageAddress) {
+  unsigned key = pointerAuthQual.getKey();
+
+  // Produce the 'other' discriminator.
+  auto otherDiscriminator = pointerAuthQual.getExtraDiscriminator();
+  llvm::Value *discriminator =
+      llvm::ConstantInt::get(IGF.IGM.Int64Ty, otherDiscriminator);
+
+  // Factor in the address.
+  if (pointerAuthQual.isAddressDiscriminated()) {
+    assert(storageAddress &&
+           "no storage address for address-discriminated schema");
+
+    if (otherDiscriminator != 0) {
+      discriminator = emitPointerAuthBlend(IGF, storageAddress, discriminator);
+    } else {
+      discriminator =
+          IGF.Builder.CreatePtrToInt(storageAddress, IGF.IGM.Int64Ty);
+    }
+  }
+
+  return PointerAuthInfo(key, discriminator);
 }
 
 llvm::ConstantInt *

lib/IRGen/GenStruct.cpp

@@ -1045,6 +1045,8 @@ class ClangRecordLowering {
       return AddressOnlyCXXClangRecordTypeInfo::create(
           FieldInfos, llvmType, TotalStride, TotalAlignment, ClangDecl);
     }
+    // TODO: New AddressOnlyPtrAuthTypeInfo to handle address diversified field
+    // function ptrs in C structs.
     return LoadableClangRecordTypeInfo::create(
         FieldInfos, NextExplosionIndex, llvmType, TotalStride,
         std::move(SpareBits), TotalAlignment, ClangDecl);

lib/IRGen/IRGenSIL.cpp

@@ -5811,7 +5811,8 @@ void IRGenSILFunction::visitBeginAccessInst(BeginAccessInst *access) {
       auto *signedFptr = Builder.CreateLoad(pointerToIntPtr, Int64PtrTy,
                                             IGM.getPointerAlignment());
       auto *resignedFptr = emitPointerAuthResign(
-          *this, signedFptr, PointerAuthInfo::emit(IGM, pointerAuthQual),
+          *this, signedFptr,
+          PointerAuthInfo::emit(*this, pointerAuthQual, pointerToSignedFptr),
           PointerAuthInfo::emit(*this,
                                 IGM.getOptions().PointerAuth.FunctionPointers,
                                 pointerToSignedFptr, PointerAuthEntity()));
@@ -5926,6 +5927,8 @@ void IRGenSILFunction::visitEndAccessInst(EndAccessInst *i) {
     auto pointerAuthQual = cast<StructElementAddrInst>(access->getOperand())
                                ->getField()
                                ->getPointerAuthQualifier();
+    auto *pointerToSignedFptr =
+        getLoweredAddress(access->getOperand()).getAddress();
     auto tempAddress = getLoweredAddress(access);
     auto *tempAddressToIntPtr =
         Builder.CreateBitCast(tempAddress.getAddress(), Int64PtrPtrTy);
@@ -5936,10 +5939,8 @@ void IRGenSILFunction::visitEndAccessInst(EndAccessInst *i) {
         PointerAuthInfo::emit(*this,
                               IGM.getOptions().PointerAuth.FunctionPointers,
                               tempAddress.getAddress(), PointerAuthEntity()),
-        PointerAuthInfo::emit(IGM, pointerAuthQual));
+        PointerAuthInfo::emit(*this, pointerAuthQual, pointerToSignedFptr));
 
-    auto *pointerToSignedFptr =
-        getLoweredAddress(access->getOperand()).getAddress();
     auto *pointerToIntPtr =
         Builder.CreateBitCast(pointerToSignedFptr, Int64PtrPtrTy);
     Builder.CreateStore(signedFptr, pointerToIntPtr, IGM.getPointerAlignment());

lib/SIL/IR/TypeLowering.cpp

@@ -2271,6 +2271,10 @@ namespace {
 
       // Classify the type according to its stored properties.
       for (auto field : D->getStoredProperties()) {
+        auto pointerAuthQual = field->getPointerAuthQualifier();
+        if (pointerAuthQual && pointerAuthQual.isAddressDiscriminated()) {
+          properties.setAddressOnly();
+        }
         auto substFieldType =
           field->getInterfaceType().subst(subMap)
                ->getCanonicalType();

lib/SILGen/SILGenLValue.cpp

@@ -854,7 +854,7 @@ namespace {
       auto Res = SGF.B.createStructElementAddr(loc, base.getValue(),
                                                Field, SubstFieldType);
 
-      if (!Field->getPointerAuthQualifier() ||
+      if (!Field->getPointerAuthQualifier().isPresent() ||
           !SGF.getOptions().EnableImportPtrauthFieldFunctionPointers) {
         return ManagedValue::forLValue(Res);
       }

test/IRGen/Inputs/ptrauth_field_fptr_import.h

@@ -2,9 +2,16 @@
 #define TEST_C_FUNCTION
 
 int returnInt() { return 111; }
+
 struct SecureStruct {
-  int (*__ptrauth(2, 0, 88)(secure_func_ptr))();
+  int (*__ptrauth(1, 0, 88)(secure_func_ptr))();
+};
+
+struct AddressDiscriminatedSecureStruct {
+  int (*__ptrauth(1, 1, 88)(secure_func_ptr))();
 };
 
 struct SecureStruct *ptr_to_secure_struct;
+struct AddressDiscriminatedSecureStruct
+    *ptr_to_addr_discriminated_secure_struct;
 #endif

test/IRGen/ptrauth_field_fptr_import.swift

@@ -1,4 +1,5 @@
 // RUN: %swift-frontend %s -enable-import-ptrauth-field-function-pointers -emit-ir -target arm64e-apple-ios13.0 -I %S/Inputs/ -validate-tbd-against-ir=none 2>&1 | %FileCheck %s
+
 // REQUIRES: CPU=arm64e
 // REQUIRES: OS=ios
 
@@ -16,7 +17,7 @@ import PointerAuth
 // CHECK:   [[CAST2:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %.secure_func_ptr to i64*
 // CHECK:   [[PTR:%.*]] = load i64*, i64** [[CAST2]], align 8
 // CHECK:   [[SIGNEDINT:%.*]] = ptrtoint i64* [[PTR]] to i64
-// CHECK:   [[DEFAULTSIGNVAL:%.*]] = call i64 @llvm.ptrauth.resign(i64 [[SIGNEDINT]], i32 2, i64 88, i32 0, i64 0)
+// CHECK:   [[DEFAULTSIGNVAL:%.*]] = call i64 @llvm.ptrauth.resign(i64 [[SIGNEDINT]], i32 1, i64 88, i32 0, i64 0)
 // CHECK:   [[AUTHPTR:%.*]] = inttoptr i64 [[DEFAULTSIGNVAL]] to i64*
 // CHECK:   [[TMPCAST1:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %ptrauth.temp to i64**
 // CHECK:   store i64* [[AUTHPTR]], i64** [[TMPCAST1]], align 8
@@ -37,13 +38,66 @@ func test_field_fn_read() -> Int32 {
 // CHECK:   [[CAST3:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %ptrauth.temp to i64**
 // CHECK:   [[LD:%.*]] = load i64*, i64** [[CAST3]], align 8
 // CHECK:   [[CAST4:%.*]] = ptrtoint i64* [[LD]] to i64
-// CHECK:   [[SIGN:%.*]] = call i64 @llvm.ptrauth.resign(i64 [[CAST4]], i32 0, i64 0, i32 2, i64 88)
+// CHECK:   [[SIGN:%.*]] = call i64 @llvm.ptrauth.resign(i64 [[CAST4]], i32 0, i64 0, i32 1, i64 88)
 // CHECK:   [[CAST5:%.*]] = inttoptr i64 [[SIGN]] to i64*
 // CHECK:   [[CAST6:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %.secure_func_ptr to i64**
 // CHECK:   store i64* [[CAST5]], i64** [[CAST6]], align 8
 func test_field_fn_ptr_modify() {
   ptr_to_secure_struct!.pointee.secure_func_ptr = returnInt
 }
 
+// CHECK: define hidden swiftcc i32 @"$s25ptrauth_field_fptr_import024test_addr_discriminated_B8_fn_reads5Int32VyF"() #0 {
+// CHECK: [[LD:%.*]] = load i64, i64* bitcast (%struct.AddressDiscriminatedSecureStruct** @ptr_to_addr_discriminated_secure_struct to i64*), align 8
+// CHECK: 5:                                                ; preds = %entry
+// CHECK:   [[CAST0:%.*]] = inttoptr i64 [[LD]] to i8*
+// CHECK:   br label %12
+// CHECK: 12:
+// CHECK:   [[AddressDiscriminatedSecureStruct:%.*]] = phi i8* [ [[CAST0]], %5 ]
+// CHECK:   [[CAST1:%.*]] = bitcast i8* [[AddressDiscriminatedSecureStruct]] to %TSo32AddressDiscriminatedSecureStructV*
+// CHECK:   %.secure_func_ptr = getelementptr inbounds %TSo32AddressDiscriminatedSecureStructV, %TSo32AddressDiscriminatedSecureStructV* %14, i32 0, i32 0
+// CHECK:   [[CAST2:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %.secure_func_ptr to i64**
+// CHECK:   [[PTR:%.*]] = load i64*, i64** [[CAST2]], align 8
+// CHECK:   [[ADDR:%.*]] = ptrtoint %Ts5Int32VIetCd_Sg* %.secure_func_ptr to i64
+// CHECK:   [[BLEND:%.*]] = call i64 @llvm.ptrauth.blend(i64 [[ADDR]], i64 88)
+// CHECK:   [[CAST3:%.*]] = ptrtoint i64* [[PTR]] to i64
+// CHECK:   [[DEFAULTSIGNVAL:%.*]] = call i64 @llvm.ptrauth.resign(i64 [[CAST3]], i32 1, i64 [[BLEND]], i32 0, i64 0)
+// CHECK:   [[AUTHPTR:%.*]] = inttoptr i64 [[DEFAULTSIGNVAL]] to i64*
+// CHECK:   [[TMPCAST1:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %ptrauth.temp to i64**
+// CHECK:   store i64* [[AUTHPTR]], i64** [[TMPCAST1]], align 8
+// CHECK:   [[TMPCAST2:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %ptrauth.temp to i64*
+// CHECK:   [[FUNCPTR:%.*]] = load i64, i64* [[TMPCAST2]], align 8
+func test_addr_discriminated_field_fn_read() -> Int32 {
+  let fn = ptr_to_addr_discriminated_secure_struct!.pointee.secure_func_ptr!
+  return fn()
+}
+
+// CHECK-LABEL: define hidden swiftcc void @"$s25ptrauth_field_fptr_import024test_addr_discriminated_B14_fn_ptr_modifyyyF"() #0 {
+// CHECK: 11:                                               ; preds = %4
+// CHECK:   [[AddressDiscriminatedSecureStruct:%.*]] = phi i8* [ %5, %4 ]
+// CHECK:   [[CAST1:%.*]] = bitcast i8* [[AddressDiscriminatedSecureStruct]] to %TSo32AddressDiscriminatedSecureStructV*
+// CHECK:   %.secure_func_ptr = getelementptr inbounds %TSo32AddressDiscriminatedSecureStructV, %TSo32AddressDiscriminatedSecureStructV* [[CAST1]], i32 0, i32 0
+// CHECK:   [[CAST2:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %ptrauth.temp to i64*
+// CHECK:   store i64 ptrtoint ({ i8*, i32, i64, i64 }* @returnInt.ptrauth to i64), i64* [[CAST2]], align 8
+// CHECK:   [[CAST3:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %ptrauth.temp to i64**
+// CHECK:   [[LD:%.*]] = load i64*, i64** [[CAST3]], align 8
+// CHECK:   [[CAST4:%.*]] = ptrtoint %Ts5Int32VIetCd_Sg* %.secure_func_ptr to i64
+// CHECK:   [[BLEND:%.*]] = call i64 @llvm.ptrauth.blend(i64 [[CAST4]], i64 88)
+// CHECK:   [[CAST5:%.*]] = ptrtoint i64* [[LD]] to i64
+// CHECK:   [[SIGN:%.*]] = call i64 @llvm.ptrauth.resign(i64 [[CAST5]], i32 0, i64 0, i32 1, i64 [[BLEND]])
+// CHECK:   [[CAST5:%.*]] = inttoptr i64 [[SIGN]] to i64*
+// CHECK:   [[CAST6:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %.secure_func_ptr to i64**
+// CHECK:   store i64* [[CAST5]], i64** [[CAST6]], align 8
+func test_addr_discriminated_field_fn_ptr_modify() {
+  ptr_to_addr_discriminated_secure_struct!.pointee.secure_func_ptr = returnInt
+}
+
+// TODO: Unimplemented non trivial pointer auth copy function
+// func test_addr_discriminated_copy() -> Int32 {
+//   let struct_with_signed_val = ptr_to_addr_discriminated_secure_struct.pointee
+//   return struct_with_signed_val.secure_func_ptr()
+// }
+
 print(test_field_fn_read())
-print(test_field_fn_ptr_modify())
+test_field_fn_ptr_modify()
+print(test_addr_discriminated_field_fn_read())
+test_addr_discriminated_field_fn_ptr_modify()