Fix IRGen for pointer auth qualified field function pointers when they are null

meg-gupta · meg-gupta · commit c0204d33ec86 · 2023-03-09T19:16:34.000-08:00
diff --git a/lib/IRGen/IRGenSIL.cpp b/lib/IRGen/IRGenSIL.cpp
@@ -5869,26 +5869,52 @@ void IRGenSILFunction::visitBeginAccessInst(BeginAccessInst *access) {
     auto *Int64PtrPtrTy = Int64PtrTy->getPointerTo();
     if (access->getAccessKind() == SILAccessKind::Read) {
       // When we see a signed read access, generate code to:
-      // authenticate the signed pointer, and store the authenticated value to a
-      // shadow stack location. Set the lowered address of the access to this
-      // stack location.
+      // authenticate the signed pointer if non-null, and store the
+      // authenticated value to a shadow stack location. Set the lowered address
+      // of the access to this stack location.
       auto pointerAuthQual = sea->getField()->getPointerAuthQualifier();
       auto *pointerToSignedFptr = getLoweredAddress(sea).getAddress();
       auto *pointerToIntPtr =
           Builder.CreateBitCast(pointerToSignedFptr, Int64PtrPtrTy);
       auto *signedFptr = Builder.CreateLoad(pointerToIntPtr, Int64PtrTy,
                                             IGM.getPointerAlignment());
-      auto *resignedFptr = emitPointerAuthResign(
-          *this, signedFptr,
-          PointerAuthInfo::emit(*this, pointerAuthQual, pointerToSignedFptr),
-          PointerAuthInfo::emit(*this,
-                                IGM.getOptions().PointerAuth.FunctionPointers,
-                                pointerToSignedFptr, PointerAuthEntity()));
+
+      // Create a stack temporary.
       auto temp = ti.allocateStack(*this, access->getType(), "ptrauth.temp");
       auto *tempAddressToIntPtr =
           Builder.CreateBitCast(temp.getAddressPointer(), Int64PtrPtrTy);
+
+      // Branch based on pointer is null or not.
+      llvm::Value *cond = Builder.CreateICmpNE(
+          signedFptr, llvm::ConstantPointerNull::get(Int64PtrTy));
+      auto *resignNonNull = createBasicBlock("resign-nonnull");
+      auto *resignNull = createBasicBlock("resign-null");
+      auto *resignCont = createBasicBlock("resign-cont");
+      Builder.CreateCondBr(cond, resignNonNull, resignNull);
+
+      // Resign if non-null.
+      Builder.emitBlock(resignNonNull);
+      auto oldAuthInfo =
+          PointerAuthInfo::emit(*this, pointerAuthQual, pointerToSignedFptr);
+      // ClangImporter imports the c function pointer as an optional type.
+      PointerAuthEntity entity(
+          sea->getType().getOptionalObjectType().getAs<SILFunctionType>());
+      auto newAuthInfo = PointerAuthInfo::emit(
+          *this, IGM.getOptions().PointerAuth.FunctionPointers,
+          pointerToSignedFptr, entity);
+      auto *resignedFptr =
+          emitPointerAuthResign(*this, signedFptr, oldAuthInfo, newAuthInfo);
       Builder.CreateStore(resignedFptr, tempAddressToIntPtr,
                           IGM.getPointerAlignment());
+      Builder.CreateBr(resignCont);
+
+      // If null, no need to resign.
+      Builder.emitBlock(resignNull);
+      Builder.CreateStore(signedFptr, tempAddressToIntPtr,
+                          IGM.getPointerAlignment());
+      Builder.CreateBr(resignCont);
+
+      Builder.emitBlock(resignCont);
       setLoweredAddress(access, temp.getAddress());
       return;
     }
@@ -5990,30 +6016,55 @@ void IRGenSILFunction::visitEndAccessInst(EndAccessInst *i) {
       return;
     }
     // When we see a signed modify access, get the lowered address of the
-    // access which is the shadow stack slot, sign the value and write back to
-    // the struct field.
+    // access which is the shadow stack slot, sign the value if non-null and
+    // write back to the struct field.
+    auto *sea = cast<StructElementAddrInst>(access->getOperand());
     auto *Int64PtrTy = llvm::Type::getInt64PtrTy(IGM.getLLVMContext());
     auto *Int64PtrPtrTy = Int64PtrTy->getPointerTo();
     auto pointerAuthQual = cast<StructElementAddrInst>(access->getOperand())
                                ->getField()
                                ->getPointerAuthQualifier();
     auto *pointerToSignedFptr =
         getLoweredAddress(access->getOperand()).getAddress();
+    auto *pointerToIntPtr =
+        Builder.CreateBitCast(pointerToSignedFptr, Int64PtrPtrTy);
     auto tempAddress = getLoweredAddress(access);
     auto *tempAddressToIntPtr =
         Builder.CreateBitCast(tempAddress.getAddress(), Int64PtrPtrTy);
     auto *tempAddressValue = Builder.CreateLoad(tempAddressToIntPtr, Int64PtrTy,
                                                 IGM.getPointerAlignment());
-    auto *signedFptr = emitPointerAuthResign(
-        *this, tempAddressValue,
-        PointerAuthInfo::emit(*this,
-                              IGM.getOptions().PointerAuth.FunctionPointers,
-                              tempAddress.getAddress(), PointerAuthEntity()),
-        PointerAuthInfo::emit(*this, pointerAuthQual, pointerToSignedFptr));
-
-    auto *pointerToIntPtr =
-        Builder.CreateBitCast(pointerToSignedFptr, Int64PtrPtrTy);
+    // Branch based on value is null or not.
+    llvm::Value *cond = Builder.CreateICmpNE(
+        tempAddressValue, llvm::ConstantPointerNull::get(Int64PtrTy));
+    auto *resignNonNull = createBasicBlock("resign-nonnull");
+    auto *resignNull = createBasicBlock("resign-null");
+    auto *resignCont = createBasicBlock("resign-cont");
+
+    Builder.CreateCondBr(cond, resignNonNull, resignNull);
+
+    Builder.emitBlock(resignNonNull);
+
+    // If non-null, resign
+    // ClangImporter imports the c function pointer as an optional type.
+    PointerAuthEntity entity(
+        sea->getType().getOptionalObjectType().getAs<SILFunctionType>());
+    auto oldAuthInfo = PointerAuthInfo::emit(
+        *this, IGM.getOptions().PointerAuth.FunctionPointers,
+        tempAddress.getAddress(), entity);
+    auto newAuthInfo =
+        PointerAuthInfo::emit(*this, pointerAuthQual, pointerToSignedFptr);
+    auto *signedFptr = emitPointerAuthResign(*this, tempAddressValue,
+                                             oldAuthInfo, newAuthInfo);
     Builder.CreateStore(signedFptr, pointerToIntPtr, IGM.getPointerAlignment());
+
+    Builder.CreateBr(resignCont);
+
+    // If null, no need to resign
+    Builder.emitBlock(resignNull);
+    Builder.CreateStore(tempAddressValue, pointerToIntPtr, IGM.getPointerAlignment());
+    Builder.CreateBr(resignCont);
+
+    Builder.emitBlock(resignCont);
     return;
   }
   }
diff --git a/test/IRGen/ptrauth_field_fptr_import.swift b/test/IRGen/ptrauth_field_fptr_import.swift
@@ -15,11 +15,19 @@ import PointerAuth
 // CHECK:   %.secure_func_ptr = getelementptr inbounds %TSo12SecureStructV, %TSo12SecureStructV* %14, i32 0, i32 0
 // CHECK:   [[CAST2:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %.secure_func_ptr to i64*
 // CHECK:   [[PTR:%.*]] = load i64*, i64** [[CAST2]], align 8
+// CHECK:   [[CAST3:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %ptrauth.temp to i64**
+// CHECK:   [[COND:%.*]] = icmp ne i64* %16, null
+// CHECK:   br i1 [[COND]], label %resign-nonnull, label %resign-null
+// CHECK: resign-nonnull:                                   ; preds = %12
 // CHECK:   [[SIGNEDINT:%.*]] = ptrtoint i64* [[PTR]] to i64
 // 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
+// CHECK:   store i64* [[AUTHPTR]], i64** [[CAST3]], align 8
+// CHECK:   br label %resign-cont
+// CHECK: resign-null:                                      ; preds = %12
+// CHECK:   store i64* [[PTR]], i64** [[CAST3]], align 8
+// CHECK:   br label %resign-cont
+// CHECK: resign-cont:                                      ; preds = %resign-null, %resign-nonnull
 // CHECK:   [[TMPCAST2:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %ptrauth.temp to i64*
 // CHECK:   [[FUNCPTR:%.*]] = load i64, i64* [[TMPCAST2]], align 8
 func test_field_fn_read() -> Int32 {
@@ -34,13 +42,16 @@ func test_field_fn_read() -> Int32 {
 // CHECK:   %.secure_func_ptr = getelementptr inbounds %TSo12SecureStructV, %TSo12SecureStructV* [[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:   [[PTR:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %.secure_func_ptr to i64**
 // CHECK:   [[CAST3:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %ptrauth.temp to i64**
 // CHECK:   [[LD:%.*]] = load i64*, i64** [[CAST3]], align 8
+// CHECK:   [[COND:%.*]] = icmp ne i64* [[LD]], null
+// CHECK:   br i1 [[COND]], label %resign-nonnull, label %resign-null
+// CHECK: resign-nonnull:                                   ; preds = %11
 // CHECK:   [[CAST4:%.*]] = ptrtoint i64* [[LD]] to i64
 // 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
+// CHECK:   store i64* [[CAST5]], i64** [[PTR]], align 8
 func test_field_fn_ptr_modify() {
   ptr_to_secure_struct!.pointee.secure_func_ptr = returnInt
 }
@@ -56,13 +67,18 @@ func test_field_fn_ptr_modify() {
 // CHECK:   %.secure_func_ptr = getelementptr inbounds %TSo32AddressDiscriminatedSecureStructV, %TSo32AddressDiscriminatedSecureStructV* [[CAST1]], i32 0, i32 0
 // CHECK:   [[CAST2:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %.secure_func_ptr to i64**
 // CHECK:   [[PTR:%.*]] = load i64*, i64** [[CAST2]], align 8
+// CHECK:   [[CAST3:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %ptrauth.temp to i64**
+// CHECK:   [[COND:%.*]] = icmp ne i64* [[PTR]], null
+// CHECK:   br i1 [[COND]], label %resign-nonnull, label %resign-null
+// CHECK: resign-nonnull:                                   ; preds = %12
 // 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:   [[CAST4:%.*]] = ptrtoint i64* [[PTR]] to i64
+// CHECK:   [[DEFAULTSIGNVAL:%.*]] = call i64 @llvm.ptrauth.resign(i64 [[CAST4]], 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:   store i64* [[AUTHPTR]], i64** [[CAST3]], align 8
+// CHECK:   br label %resign-cont
+// CHECK: resign-cont:                                      ; preds = %resign-null, %resign-nonnull
 // 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 {
@@ -82,15 +98,18 @@ func test_addr_discriminated_field_fn_read() -> Int32 {
 // 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:   [[PTR:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %.secure_func_ptr to i64**
 // CHECK:   [[CAST3:%.*]] = bitcast %Ts5Int32VIetCd_Sg* %ptrauth.temp to i64**
 // CHECK:   [[LD:%.*]] = load i64*, i64** [[CAST3]], align 8
+// CHECK:   [[COND:%.*]] = icmp ne i64* [[LD]], null
+// CHECK:   br i1 [[COND]], label %resign-nonnull, label %resign-null
+// CHECK: resign-nonnull:                                   ; preds = %11
 // 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
+// CHECK:   [[CAST6:%.*]] = inttoptr i64 [[SIGN]] to i64*
+// CHECK:   store i64* [[CAST6]], i64** [[PTR]], align 8
 func test_addr_discriminated_field_fn_ptr_modify() {
   ptr_to_addr_discriminated_secure_struct!.pointee.secure_func_ptr = returnInt
 }
