From 7350c3f90eca4953b3ca891a3fcf6d0d227575ff Mon Sep 17 00:00:00 2001 From: Jonathan Grynspan Date: Wed, 5 Nov 2025 23:31:28 -0500 Subject: [PATCH] Avoid a potential race condition and use-after-free when calling `Test.cancel()`. If a test creates an unstructured, non-detached task that continues running after the test has finished and eventually calls `Test.cancel()`, then it may be able to see a reference to the test's (or test case's) task after it has been destroyed by the Swift runtime. This PR ensures that the infrastructure under `Test.cancel()` clears its reference to the test's task before returning. This then minimizes the risk of observing the task after it has been destroyed. Further work at the Swift runtime level may be required to completely eliminate this race condition, but this change makes it sufficiently narrow that any example I can come up with is contrived. --- Sources/Testing/Test+Cancellation.swift | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/Sources/Testing/Test+Cancellation.swift b/Sources/Testing/Test+Cancellation.swift index 5a4b425c7..1ded9359f 100644 --- a/Sources/Testing/Test+Cancellation.swift +++ b/Sources/Testing/Test+Cancellation.swift @@ -87,10 +87,19 @@ extension TestCancellable { /// the current task, test, or test case is cancelled, it records a /// corresponding cancellation event. func withCancellationHandling(_ body: () async throws -> R) async rethrows -> R { + let taskReference = _TaskReference() var currentTaskReferences = _currentTaskReferences - currentTaskReferences[ObjectIdentifier(Self.self)] = _TaskReference() + currentTaskReferences[ObjectIdentifier(Self.self)] = taskReference return try await $_currentTaskReferences.withValue(currentTaskReferences) { - try await withTaskCancellationHandler { + // Before returning, explicitly clear the stored task. This minimizes + // the potential race condition that can occur if test code creates an + // unstructured task and calls `Test.cancel()` in it after the test body + // has finished. + defer { + _ = taskReference.takeUnsafeCurrentTask() + } + + return try await withTaskCancellationHandler { try await body() } onCancel: { // The current task was cancelled, so cancel the test case or test