Merge pull request #41539 from rjmccall/uaf-task-tracing

Fix a use-after-free with task tracing

Author: rjmccall

stdlib/public/Concurrency/Actor.cpp

@@ -1540,7 +1540,7 @@ static void swift_job_runImpl(Job *job, ExecutorRef executor) {
   SWIFT_TASK_DEBUG_LOG("%s(%p)", __func__, job);
   runJobInEstablishedExecutorContext(job);
 
-  concurrency::trace::job_run_end(job, &executor, traceHandle);
+  concurrency::trace::job_run_end(&executor, traceHandle);
   trackingInfo.leave();
 
   // Give up the current executor if this is a switching context

stdlib/public/Concurrency/TaskPrivate.h

@@ -280,6 +280,12 @@ class alignas(2 * sizeof(void*)) ActiveTaskStatus {
     /// actor. This bit is cleared when a starts running on a thread, suspends
     /// or is completed.
     IsEnqueued = 0x1000,
+
+#ifndef NDEBUG
+    /// Task has been completed.  This is purely used to enable an assertion
+    /// that the task is completed when we destroy it.
+    IsComplete = 0x2000,
+#endif
   };
 
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION && SWIFT_POINTER_IS_4_BYTES
@@ -393,6 +399,20 @@ class alignas(2 * sizeof(void*)) ActiveTaskStatus {
 #endif
   }
 
+#ifndef NDEBUG
+  bool isComplete() const {
+    return Flags & IsComplete;
+  }
+
+  ActiveTaskStatus withComplete() const {
+#if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
+    return ActiveTaskStatus(Record, Flags | IsComplete, ExecutionLock);
+#else
+    return ActiveTaskStatus(Record, Flags | IsComplete);
+#endif
+  }
+#endif
+
   /// Is there a lock on the linked list of status records?
   bool isStatusRecordLocked() const { return Flags & IsStatusRecordLocked; }
   ActiveTaskStatus withLockingRecord(TaskStatusRecord *lockRecord) const {
@@ -564,6 +584,9 @@ struct AsyncTask::PrivateStorage {
       auto newStatus = oldStatus.withRunning(false);
       newStatus = newStatus.withoutStoredPriorityEscalation();
       newStatus = newStatus.withoutEnqueued();
+#ifndef NDEBUG
+      newStatus = newStatus.withComplete();
+#endif
 
       // This can fail since the task can still get concurrently cancelled or
       // escalated.
@@ -650,6 +673,7 @@ retry:;
   while (true) {
     // We can get here from being suspended or being enqueued
     assert(!oldStatus.isRunning());
+    assert(!oldStatus.isComplete());
 
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
     // Task's priority is greater than the thread's - do a self escalation

stdlib/public/Concurrency/Tracing.h

@@ -75,11 +75,22 @@ void job_enqueue_global_with_delay(unsigned long long delay, Job *job);
 
 void job_enqueue_main_executor(Job *job);
 
-// This returns a handle that must be passed to the corresponding call to
-// task_run_end.
-uint64_t job_run_begin(Job *job, ExecutorRef *executor);
-
-void job_run_end(Job *job, ExecutorRef *executor, uint64_t beginHandle);
+struct job_run_info {
+  /// The ID of the task that started running.
+  uint64_t taskId;
+
+  /// The signpost ID for this task execution, or OS_SIGNPOST_ID_INVALID
+  /// if the job was not a task.
+  uint64_t handle;
+};
+
+// This returns information that must be passed to the corresponding
+// call to task_run_end.  Any information we want to log must be
+// extracted from the job when we start to run it because execution
+// will invalidate the job.
+job_run_info job_run_begin(Job *job, ExecutorRef *executor);
+
+void job_run_end(ExecutorRef *executor, job_run_info info);
 
 } // namespace trace
 } // namespace concurrency

stdlib/public/Concurrency/TracingSignpost.h

@@ -244,29 +244,33 @@ inline void job_enqueue_main_executor(Job *job) {
   }
 }
 
-inline uint64_t job_run_begin(Job *job, ExecutorRef *executor) {
+inline job_run_info job_run_begin(Job *job, ExecutorRef *executor) {
+  auto invalidInfo = []{
+    return job_run_info{ 0, OS_SIGNPOST_ID_INVALID };
+  };
+
   if (AsyncTask *task = dyn_cast<AsyncTask>(job)) {
-    ENSURE_LOGS(0);
-    auto id = os_signpost_id_generate(TaskLog);
+    ENSURE_LOGS(invalidInfo());
+    auto handle = os_signpost_id_generate(TaskLog);
+    auto taskId = task->getTaskId();
     os_signpost_interval_begin(
-        TaskLog, id, SWIFT_LOG_JOB_RUN_NAME,
+        TaskLog, handle, SWIFT_LOG_JOB_RUN_NAME,
         "task=%" PRIx64
         " executorIdentity=%p executorImplementation=0x%" PRIxPTR,
-        task->getTaskId(), executor->getIdentity(),
-        executor->getRawImplementation());
-    return id;
+        taskId, executor->getIdentity(), executor->getRawImplementation());
+    return { taskId, handle };
   }
-  return 0;
+  return invalidInfo();
 }
 
-inline void job_run_end(Job *job, ExecutorRef *executor, uint64_t beginHandle) {
-  if (AsyncTask *task = dyn_cast<AsyncTask>(job)) {
+inline void job_run_end(ExecutorRef *executor, job_run_info info) {
+  if (info.handle != OS_SIGNPOST_ID_INVALID) {
     ENSURE_LOGS();
     os_signpost_interval_end(
-        TaskLog, beginHandle, SWIFT_LOG_JOB_RUN_NAME,
+        TaskLog, info.handle, SWIFT_LOG_JOB_RUN_NAME,
         "task=%" PRIx64
         " executorIdentity=%p executorImplementation=0x%" PRIxPTR,
-        task->getTaskId(), executor->getIdentity(),
+        info.taskId, executor->getIdentity(),
         executor->getRawImplementation());
   }
 }

stdlib/public/Concurrency/TracingStubs.h

@@ -60,9 +60,9 @@ inline void job_enqueue_global_with_delay(unsigned long long delay, Job *job) {}
 
 inline void job_enqueue_main_executor(Job *job) {}
 
-inline uint64_t job_run_begin(Job *job, ExecutorRef *executor) { return 0; }
+inline job_run_info job_run_begin(Job *job, ExecutorRef *executor) { return {}; }
 
-inline void job_run_end(Job *job, ExecutorRef *executor, uint64_t beginHandle) {
+inline void job_run_end(ExecutorRef *executor, job_run_info info) {
 }
 
 } // namespace trace