From 0c0e0a36582b3d06fd18dd5f262ce6ad543067f3 Mon Sep 17 00:00:00 2001
From: Aleksei Tirman <aleksei.tirman@jetbrains.com>
Date: Thu, 28 Mar 2024 09:57:12 +0200
Subject: [PATCH] =?UTF-8?q?KTOR-6878=20Fix=20stack=20overflow=20when=20Sus?=
 =?UTF-8?q?pendFunctionGun=20trying=20to=20access=E2=80=A6=20(#4007)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* KTOR-6878 Fix stack overflow error when SuspendFunctionGun tries to access continuation context of the same continuation object recursively

(cherry picked from commit 9c78494bfccdb40bbffcba2a7aa55b079703677d)
---
 .../ktor/util/pipeline/SuspendFunctionGun.kt  | 17 ++++++--
 .../io/ktor/util/SuspendFunctionGunTest.kt    | 40 +++++++++++++++++++
 2 files changed, 54 insertions(+), 3 deletions(-)
 create mode 100644 ktor-utils/common/test/io/ktor/util/SuspendFunctionGunTest.kt

diff --git a/ktor-utils/common/src/io/ktor/util/pipeline/SuspendFunctionGun.kt b/ktor-utils/common/src/io/ktor/util/pipeline/SuspendFunctionGun.kt
index 2ac171736e..d8cbadfce6 100644
--- a/ktor-utils/common/src/io/ktor/util/pipeline/SuspendFunctionGun.kt
+++ b/ktor-utils/common/src/io/ktor/util/pipeline/SuspendFunctionGun.kt
@@ -18,7 +18,7 @@ internal class SuspendFunctionGun<TSubject : Any, TContext : Any>(
 
     // this is impossible to inline because of property name clash
     // between PipelineContext.context and Continuation.context
-    private val continuation: Continuation<Unit> = object : Continuation<Unit>, CoroutineStackFrame {
+    internal val continuation: Continuation<Unit> = object : Continuation<Unit>, CoroutineStackFrame {
         override val callerFrame: CoroutineStackFrame? get() = peekContinuation() as? CoroutineStackFrame
 
         var currentIndex: Int = Int.MIN_VALUE
@@ -48,7 +48,18 @@ internal class SuspendFunctionGun<TSubject : Any, TContext : Any>(
         }
 
         override val context: CoroutineContext
-            get() = suspensions[lastSuspensionIndex]?.context ?: error("Not started")
+            get() {
+                val continuation = suspensions[lastSuspensionIndex]
+                if (continuation !== this && continuation != null) return continuation.context
+
+                var index = lastSuspensionIndex - 1
+                while (index >= 0) {
+                    val cont = suspensions[index--]
+                    if (cont !== this && cont != null) return cont.context
+                }
+
+                error("Not started")
+            }
 
         override fun resumeWith(result: Result<Unit>) {
             if (result.isFailure) {
@@ -144,7 +155,7 @@ internal class SuspendFunctionGun<TSubject : Any, TContext : Any>(
         suspensions[lastSuspensionIndex--] = null
     }
 
-    private fun addContinuation(continuation: Continuation<TSubject>) {
+    internal fun addContinuation(continuation: Continuation<TSubject>) {
         suspensions[++lastSuspensionIndex] = continuation
     }
 }
diff --git a/ktor-utils/common/test/io/ktor/util/SuspendFunctionGunTest.kt b/ktor-utils/common/test/io/ktor/util/SuspendFunctionGunTest.kt
new file mode 100644
index 0000000000..062614d60d
--- /dev/null
+++ b/ktor-utils/common/test/io/ktor/util/SuspendFunctionGunTest.kt
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2014-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package io.ktor.util
+
+import io.ktor.test.dispatcher.*
+import io.ktor.util.pipeline.*
+import kotlin.coroutines.*
+import kotlin.test.*
+
+class SuspendFunctionGunTest {
+    @Test
+    fun throwsErrorWhenNoDistinctContinuation() = testSuspend {
+        val gun = SuspendFunctionGun(Unit, Unit, listOf({ _, _, _ -> }))
+        gun.addContinuation(gun.continuation)
+
+        val cause = assertFailsWith<IllegalStateException> { gun.continuation.context }
+        assertEquals("Not started", cause.message)
+    }
+
+    @Test
+    fun returnsLastDistinctContinuationContext() = testSuspend {
+        val gun = SuspendFunctionGun(Unit, Unit, listOf({ _, _, _ -> }, { _, _, _ -> }))
+        val continuation = Continuation<Unit>(EmptyCoroutineContext) {}
+        gun.addContinuation(continuation)
+        gun.addContinuation(gun.continuation)
+
+        assertEquals(gun.continuation.context, continuation.context)
+    }
+
+    @Test
+    fun returnsFirstDistinctContinuationContext() = testSuspend {
+        val gun = SuspendFunctionGun(Unit, Unit, listOf({ _, _, _ -> }))
+        val continuation = Continuation<Unit>(EmptyCoroutineContext) {}
+        gun.addContinuation(continuation)
+
+        assertEquals(gun.continuation.context, continuation.context)
+    }
+}