Merge pull request #35755 from nate-chandler/concurrency/irgen/rdar72357371-3

[IRGen] Sign task resume context in swift_suspend_point.

Author: DougGregor

include/swift/AST/IRGenOptions.h

@@ -147,6 +147,9 @@ struct PointerAuthOptions : clang::PointerAuthOptions {
   /// The resume function stored in AsyncTask.
   PointerAuthSchema TaskResumeFunction;
 
+  /// The async context stored in AsyncTask.
+  PointerAuthSchema TaskResumeContext;
+
   /// The swift async context entry in the extended frame info.
   PointerAuthSchema AsyncContextExtendedFrameEntry;
 };

lib/IRGen/GenFunc.cpp

@@ -2559,6 +2559,11 @@ llvm::Function *IRGenFunction::createAsyncSuspendFn() {
   auto *resumeAddr = Builder.CreateStructGEP(task, 4);
   Builder.CreateStore(resumeFunction, Address(resumeAddr, ptrAlign));
   auto *contextAddr = Builder.CreateStructGEP(task, 5);
+  if (auto schema = IGM.getOptions().PointerAuth.TaskResumeContext) {
+    auto authInfo = PointerAuthInfo::emit(suspendIGF, schema, contextAddr,
+                                          PointerAuthEntity());
+    context = emitPointerAuthSign(suspendIGF, context, authInfo);
+  }
   Builder.CreateStore(context, Address(contextAddr, ptrAlign));
   auto *suspendCall = Builder.CreateCall(
       IGM.getTaskSwitchFuncFn(),

lib/IRGen/IRGen.cpp

@@ -703,6 +703,10 @@ static void setPointerAuthOptions(PointerAuthOptions &opts,
       PointerAuthSchema(codeKey, /*address*/ true, Discrimination::Constant,
                         SpecialPointerAuthDiscriminators::TaskResumeFunction);
 
+  opts.TaskResumeContext =
+      PointerAuthSchema(dataKey, /*address*/ true, Discrimination::Constant,
+                        SpecialPointerAuthDiscriminators::TaskResumeContext);
+
   opts.AsyncContextExtendedFrameEntry = PointerAuthSchema(
       dataKey, /*address*/ true, Discrimination::Constant,
       SpecialPointerAuthDiscriminators::SwiftAsyncContextExtendedFrameEntry);

lib/IRGen/IRGenFunction.cpp

@@ -25,6 +25,7 @@
 
 #include "Callee.h"
 #include "Explosion.h"
+#include "GenPointerAuth.h"
 #include "IRGenDebugInfo.h"
 #include "IRGenFunction.h"
 #include "IRGenModule.h"
@@ -563,7 +564,13 @@ void IRGenFunction::emitGetAsyncContinuation(SILType resumeTy,
   // TODO: add lifetime with matching lifetime in await_async_continuation
   auto contResumeAddr =
       Builder.CreateStructGEP(continuationContext.getAddress(), 0);
-  Builder.CreateStore(getAsyncContext(),
+  llvm::Value *asyncContextValue = getAsyncContext();
+  if (auto schema = IGM.getOptions().PointerAuth.AsyncContextParent) {
+    auto authInfo = PointerAuthInfo::emit(*this, schema, contResumeAddr,
+                                          PointerAuthEntity());
+    asyncContextValue = emitPointerAuthSign(*this, asyncContextValue, authInfo);
+  }
+  Builder.CreateStore(asyncContextValue,
                       Address(contResumeAddr, pointerAlignment));
   auto contErrResultAddr =
       Builder.CreateStructGEP(continuationContext.getAddress(), 2);
@@ -607,15 +614,27 @@ void IRGenFunction::emitGetAsyncContinuation(SILType resumeTy,
   assert(AsyncCoroutineCurrentResume == nullptr &&
          "Don't support nested get_async_continuation");
   AsyncCoroutineCurrentResume = coroResume;
-  Builder.CreateStore(
-      Builder.CreateBitOrPointerCast(coroResume, IGM.FunctionPtrTy),
-      Address(currTaskResumeTaskAddr, pointerAlignment));
+  llvm::Value *coroResumeValue =
+      Builder.CreateBitOrPointerCast(coroResume, IGM.FunctionPtrTy);
+  if (auto schema = IGM.getOptions().PointerAuth.TaskResumeFunction) {
+    auto authInfo = PointerAuthInfo::emit(*this, schema, currTaskResumeTaskAddr,
+                                          PointerAuthEntity());
+    coroResumeValue = emitPointerAuthSign(*this, coroResumeValue, authInfo);
+  }
+  Builder.CreateStore(coroResumeValue,
+                      Address(currTaskResumeTaskAddr, pointerAlignment));
   // currTask->ResumeContext = &continuation_context;
   auto currTaskResumeCtxtAddr = Builder.CreateStructGEP(currTask, 5);
-  Builder.CreateStore(
-      Builder.CreateBitOrPointerCast(continuationContext.getAddress(),
-                                     IGM.SwiftContextPtrTy),
-      Address(currTaskResumeCtxtAddr, pointerAlignment));
+  llvm::Value *continuationContextValue = Builder.CreateBitOrPointerCast(
+      continuationContext.getAddress(), IGM.SwiftContextPtrTy);
+  if (auto schema = IGM.getOptions().PointerAuth.TaskResumeContext) {
+    auto authInfo = PointerAuthInfo::emit(*this, schema, currTaskResumeCtxtAddr,
+                                          PointerAuthEntity());
+    continuationContextValue =
+        emitPointerAuthSign(*this, continuationContextValue, authInfo);
+  }
+  Builder.CreateStore(continuationContextValue,
+                      Address(currTaskResumeCtxtAddr, pointerAlignment));
 
   // Publish all the writes.
   // continuation_context.awaitSynchronization =(atomic release) nullptr;

test/Concurrency/Runtime/actor_counters.swift

@@ -6,8 +6,6 @@
 
 // Remove with rdar://problem/72439642
 // UNSUPPORTED: asan
-// Remove with rdar://problem/72357371
-// UNSUPPORTED: CPU=arm64e
 
 #if canImport(Darwin)
 import Darwin

test/Concurrency/Runtime/async_taskgroup_add_handle_completion.swift

@@ -1,8 +1,8 @@
 // RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck %s --dump-input=always
 // REQUIRES: executable_test
 // REQUIRES: concurrency
-// REQUIRES: OS=macosx
-// REQUIRES: CPU=x86_64
+// XFAIL: windows
+// XFAIL: linux
 
 import Dispatch
 import Darwin

test/Concurrency/Runtime/async_taskgroup_is_empty.swift

@@ -1,8 +1,8 @@
 // RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck %s --dump-input always
 // REQUIRES: executable_test
 // REQUIRES: concurrency
-// REQUIRES: OS=macosx
-// REQUIRES: CPU=x86_64
+// XFAIL: windows
+// XFAIL: linux
 
 import Dispatch
 

test/Concurrency/Runtime/async_taskgroup_next_not_invoked_cancelAll.swift

@@ -1,8 +1,8 @@
 // RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck %s --dump-input=always
 // REQUIRES: executable_test
 // REQUIRES: concurrency
-// REQUIRES: OS=macosx
-// REQUIRES: CPU=x86_64
+// XFAIL: windows
+// XFAIL: linux
 
 import Dispatch
 

test/Concurrency/Runtime/async_taskgroup_next_not_invoked_without_cancelAll.swift

@@ -1,8 +1,8 @@
 // RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck %s --dump-input=always
 // REQUIRES: executable_test
 // REQUIRES: concurrency
-// REQUIRES: OS=macosx
-// REQUIRES: CPU=x86_64
+// XFAIL: windows
+// XFAIL: linux
 
 import func Foundation.sleep
 

test/Concurrency/Runtime/async_taskgroup_throw_recover.swift

@@ -1,8 +1,8 @@
 // RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck %s --dump-input always
 // REQUIRES: executable_test
 // REQUIRES: concurrency
-// REQUIRES: OS=macosx
-// REQUIRES: CPU=x86_64
+// XFAIL: windows
+// XFAIL: linux
 
 import Dispatch
 

test/Concurrency/Runtime/async_taskgroup_throw_rethrow.swift

@@ -1,8 +1,6 @@
 // RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck %s --dump-input always
 // REQUIRES: executable_test
 // REQUIRES: concurrency
-// REQUIRES: OS=macosx
-// REQUIRES: CPU=x86_64
 
 // REQUIRES: rdar73154198
 

test/Concurrency/Runtime/mainactor.swift

@@ -5,8 +5,6 @@
 
 // REQUIRES: OS=macosx || OS=ios
 // FIXME: should not require Darwin to run this test once we have async main!
-// Remove with rdar://problem/72357371
-// UNSUPPORTED: CPU=arm64e
 
 // for exit(:Int)
 #if canImport(Darwin)

test/Concurrency/Runtime/objc_async.swift

@@ -7,9 +7,6 @@
 // REQUIRES: concurrency
 // REQUIRES: objc_interop
 
-// Remove with rdar://problem/72357371
-// UNSUPPORTED: CPU=arm64e
-
 
 @main struct Main {
   static func main() async {

test/IRGen/async/get_async_continuation.sil

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -enable-experimental-concurrency -enable-objc-interop  -primary-file %s -emit-ir -sil-verify-all -disable-llvm-optzns -disable-swift-specific-llvm-optzns | %FileCheck %s
+// RUN: %target-swift-frontend -enable-experimental-concurrency -enable-objc-interop  -primary-file %s -emit-ir -sil-verify-all -disable-llvm-optzns -disable-swift-specific-llvm-optzns | %FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-%target-cpu
 // RUN: %target-swift-frontend -enable-experimental-concurrency -enable-objc-interop  -primary-file %s -emit-ir -sil-verify-all
 
 // REQUIRES: concurrency
@@ -27,7 +27,15 @@ bb0:
 //   Initialize the async  continuation context.
 // CHECK: [[context_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 0
 // CHECK: [[ctxt:%.*]] = load %swift.context*, %swift.context** [[ctxt_addr]]
-// CHECK: store %swift.context* [[ctxt]], %swift.context** [[context_addr]]
+
+// CHECK-arm64e: [[ctxt_addr_int:%[0-9]+]] = ptrtoint %swift.context** [[context_addr]] to i64
+// CHECK-arm64e: [[ptrauth_blend:%[0-9]+]]   = call i64 @llvm.ptrauth.blend.i64(i64 [[ctxt_addr_int]], i64 48546)
+// CHECK-arm64e: [[ctxt_int:%[0-9]+]]      = ptrtoint %swift.context* [[ctxt]] to i64
+// CHECK-arm64e: [[signed_int:%[0-9]+]]      = call i64 @llvm.ptrauth.sign.i64(i64 [[ctxt_int]], i32 2, i64 [[ptrauth_blend]])
+// CHECK-arm64e: [[signed_ctxt:%[0-9]+]]   = inttoptr i64 [[signed_int]] to %swift.context*
+// CHECK-arm64e: store %swift.context* [[signed_ctxt]], %swift.context** [[context_addr]]
+// CHECK-x86_64: store %swift.context* [[ctxt]], %swift.context** [[context_addr]]
+
 // CHECK: [[error_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 2
 // CHECK: store %swift.error* null, %swift.error** [[error_addr]]
 // CHECK: [[result_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 3
@@ -37,12 +45,28 @@ bb0:
 // CHECK:  [[exe:%.*]] = load %swift.executor*, %swift.executor** [[exe_addr]]
 // CHECK:  store %swift.executor* [[exe]], %swift.executor** [[exectuor_addr]]
 //   Initialize the async task with the continuation function and async continuation context.
-// CHECK: [[task_continuation_fn_addr:%.*]] = getelementptr inbounds %swift.task, %swift.task* [[tsk]], i32 0, i32 4
+// CHECK: [[continuation_fn_addr:%.*]] = getelementptr inbounds %swift.task, %swift.task* [[tsk]], i32 0, i32 4
 // CHECK: [[continuation_fn:%.*]] = call i8* @llvm.coro.async.resume()
-// CHECK: store i8* [[continuation_fn]], i8** [[task_continuation_fn_addr]]
+
+// CHECK-arm64e: [[continuation_fn_addr_int:%[0-9]+]] = ptrtoint i8** [[continuation_fn_addr]] to i64
+// CHECK-arm64e: [[ptrauth_blend:%[0-9]+]]   = call i64 @llvm.ptrauth.blend.i64(i64 [[continuation_fn_addr_int]], i64 11330)
+// CHECK-arm64e: [[continuation_fn_int:%[0-9]+]]      = ptrtoint i8* [[continuation_fn]] to i64
+// CHECK-arm64e: [[signed_int:%[0-9]+]]      = call i64 @llvm.ptrauth.sign.i64(i64 [[continuation_fn_int]], i32 0, i64 [[ptrauth_blend]])
+// CHECK-arm64e: [[signed_continuation_fn:%[0-9]+]]   = inttoptr i64 [[signed_int]] to i8*
+// CHECK-arm64e: store i8* [[signed_continuation_fn]], i8** [[continuation_fn_addr]]
+// CHECK-x86_64: store i8* [[continuation_fn]], i8** [[continuation_fn_addr]]
+
 // CHECK: [[task_resume_context_addr:%.*]] = getelementptr inbounds %swift.task, %swift.task* [[tsk]], i32 0, i32 5
-// CHECK: [[cont_context2:%.*]] = bitcast %swift.async_continuation_context* [[cont_context]] to %swift.context*
-// CHECK: store %swift.context* [[cont_context2]], %swift.context** [[task_resume_context_addr]]
+// CHECK: [[task_resume_context:%.*]] = bitcast %swift.async_continuation_context* [[cont_context]] to %swift.context*
+
+// CHECK-arm64e: [[task_resume_context_addr_int:%[0-9]+]] = ptrtoint %swift.context** [[task_resume_context_addr]] to i64
+// CHECK-arm64e: [[ptrauth_blend:%[0-9]+]]   = call i64 @llvm.ptrauth.blend.i64(i64 [[task_resume_context_addr_int]], i64 30010)
+// CHECK-arm64e: [[task_resume_context_int:%[0-9]+]]      = ptrtoint %swift.context* [[task_resume_context]] to i64
+// CHECK-arm64e: [[signed_int:%[0-9]+]]      = call i64 @llvm.ptrauth.sign.i64(i64 [[task_resume_context_int]], i32 2, i64 [[ptrauth_blend]])
+// CHECK-arm64e: [[signed_task_resume_context:%[0-9]+]]   = inttoptr i64 [[signed_int]] to %swift.context*
+// CHECK-arm64e: store %swift.context* [[signed_task_resume_context]], %swift.context** [[task_resume_context_addr]]
+// CHECK-x86_64: store %swift.context* [[task_resume_context]], %swift.context** [[task_resume_context_addr]]
+
 //  Initialize the synchronization variable.
 // CHECK: [[synchronization_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 1
 // CHECK: store atomic {{(i64|i32)}} 0, {{(i64|i32)}}* [[synchronization_addr]] release