Don't force a response body read on all trailers (#8904)

We only need to read the response body if its coming
from a real server.

Closes: #8902

Co-authored-by: Jesse Wilson <[email protected]>

Author: swankjesse

okhttp/src/commonJvmAndroid/kotlin/okhttp3/Response.kt

@@ -28,7 +28,6 @@ import okhttp3.internal.connection.Exchange
 import okhttp3.internal.http.HTTP_PERM_REDIRECT
 import okhttp3.internal.http.HTTP_TEMP_REDIRECT
 import okhttp3.internal.http.parseChallenges
-import okhttp3.internal.skipAll
 import okio.Buffer
 
 /**
@@ -191,13 +190,7 @@ class Response internal constructor(
    *     dropped.
    */
   @Throws(IOException::class)
-  fun trailers(): Headers {
-    val source = body.source()
-    if (source.isOpen) {
-      source.skipAll()
-    }
-    return trailersSource.get()
-  }
+  fun trailers(): Headers = trailersSource.get()
 
   /**
    * Peeks up to [byteCount] bytes from the response body and returns them as a new response
@@ -478,7 +471,6 @@ class Response internal constructor(
 
     internal fun initExchange(exchange: Exchange) {
       this.exchange = exchange
-      this.trailersSource = TrailersSource { exchange.trailers() }
     }
 
     open fun build(): Response {

okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http/CallServerInterceptor.kt

@@ -21,6 +21,7 @@ import okhttp3.Interceptor
 import okhttp3.Response
 import okhttp3.internal.connection.Exchange
 import okhttp3.internal.http2.ConnectionShutdownException
+import okhttp3.internal.skipAll
 import okhttp3.internal.stripBody
 import okio.buffer
 
@@ -129,10 +130,17 @@ class CallServerInterceptor(
           // Connection is upgrading, but we need to ensure interceptors see a non-null response body.
           response.stripBody()
         } else {
+          val responseBody = exchange.openResponseBody(response)
           response
             .newBuilder()
-            .body(exchange.openResponseBody(response))
-            .build()
+            .body(responseBody)
+            .trailers {
+              val source = responseBody.source()
+              if (source.isOpen) {
+                source.skipAll()
+              }
+              exchange.trailers()
+            }.build()
         }
       if ("close".equals(response.request.header("Connection"), ignoreCase = true) ||
         "close".equals(response.header("Connection"), ignoreCase = true)

okhttp/src/jvmTest/kotlin/okhttp3/TrailersTest.kt

@@ -35,7 +35,10 @@ import mockwebserver3.SocketEffect.CloseSocket
 import mockwebserver3.SocketEffect.ShutdownConnection
 import mockwebserver3.junit5.StartStop
 import okhttp3.Headers.Companion.headersOf
+import okhttp3.HttpUrl.Companion.toHttpUrl
+import okhttp3.ResponseBody.Companion.toResponseBody
 import okhttp3.testing.PlatformRule
+import okio.BufferedSource
 import okio.IOException
 import okio.Path.Companion.toPath
 import okio.fakefilesystem.FakeFileSystem
@@ -419,6 +422,70 @@ open class TrailersTest {
     }
   }
 
+  @Test
+  fun bufferResponseBodyAndReadTrailersHttp1() {
+    bufferResponseBodyAndReadTrailers(Protocol.HTTP_1_1)
+  }
+
+  @Test
+  fun bufferResponseBodyAndReadTrailersHttp2() {
+    bufferResponseBodyAndReadTrailers(Protocol.HTTP_2)
+  }
+
+  private fun bufferResponseBodyAndReadTrailers(protocol: Protocol) {
+    enableProtocol(protocol)
+
+    server.enqueue(
+      MockResponse
+        .Builder()
+        .trailers(headersOf("t1", "v1"))
+        .body(protocol, "Hello")
+        .build(),
+    )
+
+    val call = client.newCall(Request(server.url("/")))
+    call.execute().use { originalResponse ->
+      val responseBodyData = originalResponse.body.byteString()
+      val responseTrailers = originalResponse.trailers()
+      assertThat(responseTrailers).isEqualTo(headersOf("t1", "v1"))
+
+      val rewrittenResponse =
+        originalResponse
+          .newBuilder()
+          .body(responseBodyData.toResponseBody())
+          .build()
+      assertThat(rewrittenResponse.body.string()).isEqualTo("Hello")
+      assertThat(rewrittenResponse.trailers()).isEqualTo(headersOf("t1", "v1"))
+    }
+  }
+
+  /**
+   * We had a bug where a custom `ResponseBody` interacted poorly with `Response.trailers()`.
+   * Confirm custom trailers can be read without even accessing the response body.
+   */
+  @Test
+  fun customTrailersDoNotUseResponseBody() {
+    val response =
+      Response
+        .Builder()
+        .request(Request(url = "https://example.com".toHttpUrl()))
+        .protocol(Protocol.HTTP_1_1)
+        .code(200)
+        .message("OK")
+        .body(
+          object : ResponseBody() {
+            override fun contentType(): MediaType? = null
+
+            override fun contentLength(): Long = -1L
+
+            override fun source(): BufferedSource = error("unexpected call")
+          },
+        ).trailers { headersOf("t1", "v1") }
+        .build()
+
+    assertThat(response.trailers()).isEqualTo(headersOf("t1", "v1"))
+  }
+
   private fun MockResponse.Builder.body(
     protocol: Protocol,
     body: String,