From 4005f6925126d7fd1a9d36e90a2da0f3573a1b50 Mon Sep 17 00:00:00 2001
From: Vsevolod Tolstopyatov <qwwdfsad@gmail.com>
Date: Mon, 7 Jun 2021 18:15:18 +0300
Subject: [PATCH] Verify integrity of the recovered exception's message on both
 code paths

Fixes #2749
---
 .../jvm/src/internal/StackTraceRecovery.kt    | 13 +++++---
 .../StackTraceRecoveryCustomExceptionsTest.kt | 30 +++++++++++++++++++
 .../test/exceptions/StackTraceRecoveryTest.kt | 14 ---------
 3 files changed, 39 insertions(+), 18 deletions(-)

diff --git a/kotlinx-coroutines-core/jvm/src/internal/StackTraceRecovery.kt b/kotlinx-coroutines-core/jvm/src/internal/StackTraceRecovery.kt
index 48e8790cd1..4e17c08201 100644
--- a/kotlinx-coroutines-core/jvm/src/internal/StackTraceRecovery.kt
+++ b/kotlinx-coroutines-core/jvm/src/internal/StackTraceRecovery.kt
@@ -29,7 +29,7 @@ private val stackTraceRecoveryClassName = runCatching {
 internal actual fun <E : Throwable> recoverStackTrace(exception: E): E {
     if (!RECOVER_STACK_TRACES) return exception
     // No unwrapping on continuation-less path: exception is not reported multiple times via slow paths
-    val copy = tryCopyException(exception) ?: return exception
+    val copy = tryCopyAndVerify(exception) ?: return exception
     return copy.sanitizeStackTrace()
 }
 
@@ -66,9 +66,7 @@ private fun <E : Throwable> recoverFromStackFrame(exception: E, continuation: Co
     val (cause, recoveredStacktrace) = exception.causeAndStacktrace()
 
     // Try to create an exception of the same type and get stacktrace from continuation
-    val newException = tryCopyException(cause) ?: return exception
-    // Verify that the new exception has the same message as the original one (bail out if not, see #1631)
-    if (newException.message != cause.message) return exception
+    val newException = tryCopyAndVerify(cause) ?: return exception
     // Update stacktrace
     val stacktrace = createStackTrace(continuation)
     if (stacktrace.isEmpty()) return exception
@@ -80,6 +78,13 @@ private fun <E : Throwable> recoverFromStackFrame(exception: E, continuation: Co
     return createFinalException(cause, newException, stacktrace)
 }
 
+private fun <E : Throwable> tryCopyAndVerify(exception: E): E? {
+    val newException = tryCopyException(exception) ?: return null
+    // Verify that the new exception has the same message as the original one (bail out if not, see #1631)
+    if (newException.message != exception.message) return null
+    return newException
+}
+
 /*
  * Here we partially copy original exception stackTrace to make current one much prettier.
  * E.g. for
diff --git a/kotlinx-coroutines-core/jvm/test/exceptions/StackTraceRecoveryCustomExceptionsTest.kt b/kotlinx-coroutines-core/jvm/test/exceptions/StackTraceRecoveryCustomExceptionsTest.kt
index 70336659e8..3cac3ebabe 100644
--- a/kotlinx-coroutines-core/jvm/test/exceptions/StackTraceRecoveryCustomExceptionsTest.kt
+++ b/kotlinx-coroutines-core/jvm/test/exceptions/StackTraceRecoveryCustomExceptionsTest.kt
@@ -5,6 +5,7 @@
 package kotlinx.coroutines.exceptions
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.*
 import org.junit.Test
 import kotlin.test.*
 
@@ -71,4 +72,33 @@ class StackTraceRecoveryCustomExceptionsTest : TestBase() {
             assertEquals("custom", cause.message)
         }
     }
+
+    class WrongMessageException(token: String) : RuntimeException("Token $token")
+
+    @Test
+    fun testWrongMessageException() = runTest {
+        val result = runCatching {
+            coroutineScope<Unit> {
+                throw WrongMessageException("OK")
+            }
+        }
+        val ex = result.exceptionOrNull() ?: error("Expected to fail")
+        assertTrue(ex is WrongMessageException)
+        assertEquals("Token OK", ex.message)
+    }
+
+    @Test
+    fun testWrongMessageExceptionInChannel() = runTest {
+        // Separate code path
+        val result = produce<Unit>(SupervisorJob() + Dispatchers.Unconfined) {
+            throw WrongMessageException("OK")
+        }
+        val ex = runCatching {
+            for (unit in result) {
+                // Iterator has a special code path
+            }
+        }.exceptionOrNull() ?: error("Expected to fail")
+        assertTrue(ex is WrongMessageException)
+        assertEquals("Token OK", ex.message)
+    }
 }
diff --git a/kotlinx-coroutines-core/jvm/test/exceptions/StackTraceRecoveryTest.kt b/kotlinx-coroutines-core/jvm/test/exceptions/StackTraceRecoveryTest.kt
index 8dc106bc22..0a8b6530e2 100644
--- a/kotlinx-coroutines-core/jvm/test/exceptions/StackTraceRecoveryTest.kt
+++ b/kotlinx-coroutines-core/jvm/test/exceptions/StackTraceRecoveryTest.kt
@@ -261,18 +261,4 @@ class StackTraceRecoveryTest : TestBase() {
         }
         yield() // nop to make sure it is not a tail call
     }
-
-    @Test
-    fun testWrongMessageException() = runTest {
-        val result = runCatching {
-            coroutineScope<Unit> {
-                throw WrongMessageException("OK")
-            }
-        }
-        val ex = result.exceptionOrNull() ?: error("Expected to fail")
-        assertTrue(ex is WrongMessageException)
-        assertEquals("Token OK", ex.message)
-    }
-
-    public class WrongMessageException(token: String) : RuntimeException("Token $token")
 }