Follow-up: KTOR-9497 Preventing a fatal crash in DarwinSession on close (#5533)

* Throw CancellationException when session closed
* Port session synchronization to Darwin legacy

Author: osipxd

ktor-client/ktor-client-darwin-legacy/darwin/src/io/ktor/client/engine/darwin/internal/legacy/DarwinLegacySession.kt

@@ -8,11 +8,16 @@ import io.ktor.client.engine.darwin.*
 import io.ktor.client.request.*
 import io.ktor.utils.io.*
 import io.ktor.utils.io.core.*
-import kotlinx.atomicfu.*
-import kotlinx.cinterop.*
+import kotlinx.atomicfu.atomic
+import kotlinx.atomicfu.locks.SynchronizedObject
+import kotlinx.atomicfu.locks.synchronized
+import kotlinx.cinterop.UnsafeNumber
 import kotlinx.coroutines.job
-import platform.Foundation.*
-import kotlin.coroutines.*
+import platform.Foundation.NSOperationQueue
+import platform.Foundation.NSURLSession
+import platform.Foundation.NSURLSessionConfiguration
+import platform.Foundation.NSURLSessionTaskStateRunning
+import kotlin.coroutines.CoroutineContext
 
 @OptIn(UnsafeNumber::class)
 @Suppress("DEPRECATION")
@@ -21,6 +26,7 @@ internal class DarwinLegacySession(
     requestQueue: NSOperationQueue?
 ) : Closeable {
     private val closed = atomic(false)
+    private val sessionLock = SynchronizedObject()
 
     private val sessionAndDelegate = config.sessionAndDelegate ?: createSession(config, requestQueue)
     private val session = sessionAndDelegate.first
@@ -30,7 +36,7 @@ internal class DarwinLegacySession(
     internal suspend fun execute(request: HttpRequestData, callContext: CoroutineContext): HttpResponseData {
         val nativeRequest = request.toNSUrlRequest()
             .apply(config.requestConfig)
-        val task = session.dataTaskWithRequest(nativeRequest)
+        val task = withSession { dataTaskWithRequest(nativeRequest) }
         val response = delegate.read(request, callContext, task)
 
         callContext.job.invokeOnCompletion { cause ->
@@ -49,9 +55,24 @@ internal class DarwinLegacySession(
         }
     }
 
+    private inline fun <T> withSession(block: NSURLSession.() -> T): T {
+        cancelIfClosed()
+        return synchronized(sessionLock) {
+            cancelIfClosed()
+            block(session)
+        }
+    }
+
     override fun close() {
-        if (!closed.compareAndSet(expect = false, update = true)) return
-        session.finishTasksAndInvalidate()
+        if (closed.value) return
+        synchronized(sessionLock) {
+            if (!closed.compareAndSet(expect = false, update = true)) return
+            session.finishTasksAndInvalidate()
+        }
+    }
+
+    private fun cancelIfClosed() {
+        if (closed.value) throw CancellationException("Darwin session is closed")
     }
 }
 

ktor-client/ktor-client-darwin-legacy/darwin/test/DarwinLegacyEngineTest.kt

@@ -9,8 +9,11 @@ import io.ktor.client.request.*
 import io.ktor.client.statement.*
 import io.ktor.client.test.base.*
 import io.ktor.http.*
+import io.ktor.utils.io.*
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.withTimeout
 import platform.Foundation.NSHTTPCookieStorage.Companion.sharedHTTPCookieStorage
 import platform.Foundation.NSOperationQueue
@@ -20,6 +23,7 @@ import platform.Foundation.setValue
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertTrue
+import kotlin.test.fail
 import kotlin.time.Duration.Companion.seconds
 
 @Suppress("DEPRECATION")
@@ -138,6 +142,38 @@ class DarwinLegacyEngineTest : ClientEngineTest<DarwinLegacyClientEngineConfig>(
         }
     }
 
+    @Test
+    fun testExecuteAfterSessionClose() = runTest {
+        val config = DarwinLegacyClientEngineConfig()
+        val session = DarwinLegacySession(config, null)
+
+        session.close()
+        launch {
+            val request = request { url(TEST_SERVER) }.build()
+            session.execute(request, coroutineContext)
+            fail("Execution expected to be cancelled")
+        }
+    }
+
+    @OptIn(InternalAPI::class)
+    @Test
+    fun testWebSocketExecuteAfterSessionClose() = runTest {
+        val config = DarwinLegacyClientEngineConfig()
+        val session = DarwinLegacySession(config, null)
+
+        session.close()
+        launch {
+            val request = request {
+                url(TEST_WEBSOCKET_SERVER)
+                body = object : ClientUpgradeContent() {
+                    override fun verify(headers: Headers) = Unit
+                }
+            }.build()
+            session.execute(request, coroutineContext)
+            fail("Execution expected to be cancelled")
+        }
+    }
+
     private fun stringToNSUrlString(value: String): String {
         return Url(value).toNSUrl().absoluteString!!
     }

ktor-client/ktor-client-darwin/darwin/src/io/ktor/client/engine/darwin/internal/DarwinSession.kt

@@ -8,13 +8,17 @@ import io.ktor.client.engine.darwin.*
 import io.ktor.client.request.*
 import io.ktor.utils.io.*
 import io.ktor.utils.io.core.*
-import kotlinx.atomicfu.*
+import kotlinx.atomicfu.atomic
 import kotlinx.atomicfu.locks.SynchronizedObject
 import kotlinx.atomicfu.locks.synchronized
-import kotlinx.cinterop.*
-import kotlinx.coroutines.*
-import platform.Foundation.*
-import kotlin.coroutines.*
+import kotlinx.cinterop.ExperimentalForeignApi
+import kotlinx.cinterop.UnsafeNumber
+import kotlinx.coroutines.job
+import platform.Foundation.NSOperationQueue
+import platform.Foundation.NSURLSession
+import platform.Foundation.NSURLSessionConfiguration
+import platform.Foundation.NSURLSessionTaskStateRunning
+import kotlin.coroutines.CoroutineContext
 
 @OptIn(UnsafeNumber::class)
 internal class DarwinSession(
@@ -75,7 +79,7 @@ internal class DarwinSession(
     }
 
     private fun cancelIfClosed() {
-        check(!closed.value) { "Darwin session is closed" }
+        if (closed.value) throw CancellationException("Darwin session is closed")
     }
 }
 

ktor-client/ktor-client-darwin/darwin/test/DarwinEngineTest.kt

@@ -16,10 +16,8 @@ import io.ktor.utils.io.*
 import io.ktor.websocket.*
 import kotlinx.cinterop.ExperimentalForeignApi
 import kotlinx.cinterop.UnsafeNumber
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.TimeoutCancellationException
-import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.withTimeout
+import kotlinx.coroutines.*
+import kotlinx.coroutines.test.runTest
 import platform.Foundation.NSHTTPCookieStorage.Companion.sharedHTTPCookieStorage
 import platform.Foundation.NSOperationQueue
 import platform.Foundation.NSURLSession
@@ -277,45 +275,35 @@ class DarwinEngineTest : ClientEngineTest<DarwinClientEngineConfig>(Darwin) {
     }
 
     @Test
-    fun testExecuteAfterSessionCloseThrowsCatchableException() = runBlocking {
+    fun testExecuteAfterSessionClose() = runTest {
         val config = DarwinClientEngineConfig()
         val session = DarwinSession(config, null)
-        session.close()
-
-        val request = HttpRequestBuilder().apply {
-            url(TEST_SERVER)
-        }.build()
 
-        // Executing on a closed session must throw a catchable Kotlin exception,
-        // not crash with SIGABRT from an NSException on the invalidated NSURLSession
-        assertFailsWith<IllegalStateException> {
+        session.close()
+        launch {
+            val request = request { url(TEST_SERVER) }.build()
             session.execute(request, coroutineContext)
+            fail("Execution expected to be cancelled")
         }
-        return@runBlocking
     }
 
     @OptIn(InternalAPI::class)
     @Test
-    fun testWebSocketExecuteAfterSessionCloseThrowsCatchableException() = runBlocking {
+    fun testWebSocketExecuteAfterSessionClose() = runTest {
         val config = DarwinClientEngineConfig()
         val session = DarwinSession(config, null)
-        session.close()
 
-        // ClientUpgradeContent body is needed to trigger the webSocketTaskWithRequest branch
-        // in DarwinSession.execute(), since isUpgradeRequest() checks `body is ClientUpgradeContent`
-        val request = HttpRequestBuilder().apply {
-            url(TEST_WEBSOCKET_SERVER)
-            body = object : ClientUpgradeContent() {
-                override fun verify(headers: Headers) = Unit
-            }
-        }.build()
-
-        // WebSocket execute on a closed session must throw a catchable Kotlin exception,
-        // not crash with SIGABRT from an NSException on the invalidated NSURLSession
-        assertFailsWith<IllegalStateException> {
+        session.close()
+        launch {
+            val request = request {
+                url(TEST_WEBSOCKET_SERVER)
+                body = object : ClientUpgradeContent() {
+                    override fun verify(headers: Headers) = Unit
+                }
+            }.build()
             session.execute(request, coroutineContext)
+            fail("Execution expected to be cancelled")
         }
-        return@runBlocking
     }
 
     private fun stringToNSUrlString(value: String): String {