[Concurrency] Fix unsafe continuation validation when a continued task has been destroyed.

SWIFT_DEBUG_VALIDATE_UNCHECKED_CONTINUATIONS works by tracking the context pointers of active continuations, and verifying that a resumed context is in the set of active continuations. However, the resume calls are passed the task pointer, not the context pointer. The context pointer is recovered from the task. If the task has been destroyed, the context pointer is invalid. This can result in a weird error message or it can crash if the context pointer is used before checking it against the active continuations.

Instead, track tasks that are suspended pending an unchecked continuation. If the task is destroyed, we'll still be passed the dangling pointer and check that pointer against the tracking info. We must be sure to check that before trying to use anything inside it.

rdar://131858544

Author: mikeash

stdlib/public/Concurrency/Task.cpp

@@ -1397,7 +1397,7 @@ enum class State : uint8_t { Uninitialized, On, Off };
 static std::atomic<State> CurrentState;
 
 static LazyMutex ActiveContinuationsLock;
-static Lazy<std::unordered_set<ContinuationAsyncContext *>> ActiveContinuations;
+static Lazy<std::unordered_set<AsyncTask *>> ActiveContinuations;
 
 static bool isEnabled() {
   auto state = CurrentState.load(std::memory_order_relaxed);
@@ -1410,39 +1410,39 @@ static bool isEnabled() {
   return state == State::On;
 }
 
-static void init(ContinuationAsyncContext *context) {
+static void init(AsyncTask *task) {
   if (!isEnabled())
     return;
 
   LazyMutex::ScopedLock guard(ActiveContinuationsLock);
-  auto result = ActiveContinuations.get().insert(context);
+  auto result = ActiveContinuations.get().insert(task);
   auto inserted = std::get<1>(result);
   if (!inserted)
     swift_Concurrency_fatalError(
         0,
-        "Initializing continuation context %p that was already initialized.\n",
-        context);
+        "Initializing continuation for task %p that was already initialized.\n",
+        task);
 }
 
-static void willResume(ContinuationAsyncContext *context) {
+static void willResume(AsyncTask *task) {
   if (!isEnabled())
     return;
 
   LazyMutex::ScopedLock guard(ActiveContinuationsLock);
-  auto removed = ActiveContinuations.get().erase(context);
+  auto removed = ActiveContinuations.get().erase(task);
   if (!removed)
-    swift_Concurrency_fatalError(0,
-                      "Resuming continuation context %p that was not awaited "
-                      "(may have already been resumed).\n",
-                      context);
+    swift_Concurrency_fatalError(
+        0,
+        "Resuming continuation for task %p that is not awaited "
+        "(may have already been resumed).\n",
+        task);
 }
 
 } // namespace continuationChecking
 
 SWIFT_CC(swift)
 static AsyncTask *swift_continuation_initImpl(ContinuationAsyncContext *context,
                                               AsyncContinuationFlags flags) {
-  continuationChecking::init(context);
   context->Flags = ContinuationAsyncContext::FlagsType();
   if (flags.canThrow()) context->Flags.setCanThrow(true);
   if (flags.isExecutorSwitchForced())
@@ -1479,6 +1479,7 @@ static AsyncTask *swift_continuation_initImpl(ContinuationAsyncContext *context,
   task->ResumeContext = context;
   task->ResumeTask = context->ResumeParent;
 
+  continuationChecking::init(task);
   concurrency::trace::task_continuation_init(task, context);
 
   return task;
@@ -1593,8 +1594,6 @@ static void swift_continuation_awaitImpl(ContinuationAsyncContext *context) {
 
 static void resumeTaskAfterContinuation(AsyncTask *task,
                                         ContinuationAsyncContext *context) {
-  continuationChecking::willResume(context);
-
   auto &sync = context->AwaitSynchronization;
 
   auto status = sync.load(std::memory_order_acquire);
@@ -1644,13 +1643,15 @@ static void resumeTaskAfterContinuation(AsyncTask *task,
 
 SWIFT_CC(swift)
 static void swift_continuation_resumeImpl(AsyncTask *task) {
+  continuationChecking::willResume(task);
   auto context = static_cast<ContinuationAsyncContext*>(task->ResumeContext);
   concurrency::trace::task_continuation_resume(context, false);
   resumeTaskAfterContinuation(task, context);
 }
 
 SWIFT_CC(swift)
 static void swift_continuation_throwingResumeImpl(AsyncTask *task) {
+  continuationChecking::willResume(task);
   auto context = static_cast<ContinuationAsyncContext*>(task->ResumeContext);
   concurrency::trace::task_continuation_resume(context, false);
   resumeTaskAfterContinuation(task, context);
@@ -1660,6 +1661,7 @@ static void swift_continuation_throwingResumeImpl(AsyncTask *task) {
 SWIFT_CC(swift)
 static void swift_continuation_throwingResumeWithErrorImpl(AsyncTask *task,
                                                 /* +1 */ SwiftError *error) {
+  continuationChecking::willResume(task);
   auto context = static_cast<ContinuationAsyncContext*>(task->ResumeContext);
   concurrency::trace::task_continuation_resume(context, true);
   context->ErrorResult = error;