From 9f5f0034851047f238ec4028ffce84bb0d8a9508 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Quenaudon?= <benoit@quenaudon.com>
Date: Wed, 4 Sep 2024 14:11:51 +0100
Subject: [PATCH] Don't throw if reading trailers fail

---
 .../kotlin/com/squareup/wire/internal/grpc.kt | 10 ++--
 .../java/com/squareup/wire/GrpcClientTest.kt  | 53 +++++++++++++++++++
 2 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/wire-grpc-client/src/jvmMain/kotlin/com/squareup/wire/internal/grpc.kt b/wire-grpc-client/src/jvmMain/kotlin/com/squareup/wire/internal/grpc.kt
index 4aedd37be2..83de498321 100644
--- a/wire-grpc-client/src/jvmMain/kotlin/com/squareup/wire/internal/grpc.kt
+++ b/wire-grpc-client/src/jvmMain/kotlin/com/squareup/wire/internal/grpc.kt
@@ -202,10 +202,12 @@ private fun Response.checkGrpcResponse() {
 internal fun Response.grpcResponseToException(suppressed: IOException? = null): IOException? {
   var trailers = headersOf()
   var transportException = suppressed
-  try {
-    trailers = trailers()
-  } catch (e: IOException) {
-    if (transportException == null) transportException = e
+  if (suppressed == null) {
+    try {
+      trailers = trailers()
+    } catch (e: IOException) {
+      transportException = e
+    }
   }
 
   val grpcStatus = trailers["grpc-status"] ?: header("grpc-status")
diff --git a/wire-grpc-tests/src/test/java/com/squareup/wire/GrpcClientTest.kt b/wire-grpc-tests/src/test/java/com/squareup/wire/GrpcClientTest.kt
index 727304e5ee..c30effb008 100644
--- a/wire-grpc-tests/src/test/java/com/squareup/wire/GrpcClientTest.kt
+++ b/wire-grpc-tests/src/test/java/com/squareup/wire/GrpcClientTest.kt
@@ -53,7 +53,9 @@ import okhttp3.ResponseBody
 import okhttp3.ResponseBody.Companion.toResponseBody
 import okio.Buffer
 import okio.ByteString
+import okio.ForwardingSource
 import okio.IOException
+import okio.buffer
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.After
 import org.junit.Assert.assertThrows
@@ -94,6 +96,10 @@ class GrpcClientTest {
     override fun intercept(chain: Chain) = chain.proceed(chain.request())
   }
 
+  private var networkInterceptor: Interceptor = object : Interceptor {
+    override fun intercept(chain: Chain) = chain.proceed(chain.request())
+  }
+
   @Before
   fun setUp() {
     okhttpClient = OkHttpClient.Builder()
@@ -101,6 +107,9 @@ class GrpcClientTest {
         callReference.set(chain.call())
         interceptor.intercept(chain)
       }
+      .addNetworkInterceptor { chain ->
+        networkInterceptor.intercept(chain)
+      }
       .protocols(listOf(H2_PRIOR_KNOWLEDGE))
       .callTimeout(okHttpClientTimeout)
       .build()
@@ -555,6 +564,50 @@ class GrpcClientTest {
     }
   }
 
+  /**
+   * Force an IOException in the gRPC response body stream, and confirm that Wire fails with the
+   * right exception (IOException). We've had bugs where we incorrectly attempt to read trailers
+   * while handling an exception.
+   */
+  @Test
+  fun duplexSuspend_responseBodyThrows() {
+    mockService.enqueue(ReceiveCall("/routeguide.RouteGuide/RouteChat"))
+    mockService.enqueueSendNote(message = "marco")
+
+    // This interceptor throws an exception at the start of the stream instead of returning -1L.
+    networkInterceptor = object : Interceptor {
+      override fun intercept(chain: Chain): Response {
+        val response = chain.proceed(chain.request())
+        return response.newBuilder()
+          .body(object : ResponseBody() {
+            override fun contentLength() = response.body.contentLength()
+            override fun contentType() = response.body.contentType()
+            override fun source() = object : ForwardingSource(response.body.source()) {
+              override fun read(sink: Buffer, byteCount: Long) = throw IOException("boom!")
+            }.buffer()
+          })
+          .build()
+      }
+    }
+
+    runBlocking {
+      val (_, responseChannel) = routeGuideService.RouteChat().executeIn(this)
+
+      try {
+        responseChannel.receive()
+        fail()
+      } catch (expected: IOException) {
+        assertThat(expected)
+          .hasMessage("gRPC transport failure (HTTP status=200, grpc-status=null, grpc-message=null)")
+          .cause() // Synthetic exception added by coroutines in StackTraceRecovery.kt.
+          .cause()
+          .hasMessage("boom!")
+      }
+
+      assertThat(responseChannel.isClosedForReceive).isTrue()
+    }
+  }
+
   @Test
   fun duplexSuspend_responseChannelThrowsWhenCallCanceledByClient() {
     mockService.enqueue(ReceiveCall("/routeguide.RouteGuide/RouteChat"))