Fix RequestBody events on upgraded connections (#8970)

* Add SocketSink{Start,End} events

* Add SocketSource{Start,End} events

* Emit SocketSink/Source events

* Ensure same connection

* Chore

* Chore

* There and back again

* Update okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/connection/Exchange.kt

* Update okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/connection/Exchange.kt

---------

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

Author: gesellix

okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/connection/Exchange.kt

@@ -46,6 +46,10 @@ class Exchange(
   internal var isDuplex: Boolean = false
     private set
 
+  /** True if the request body should not be used, but the socket, instead. */
+  internal var isSocket: Boolean = false
+    private set
+
   /** True if there was an exception on the connection to the peer. */
   internal var hasFailure: Boolean = false
     private set
@@ -143,11 +147,10 @@ class Exchange(
   fun peekTrailers(): Headers? = codec.peekTrailers()
 
   fun upgradeToSocket(): Socket {
+    isSocket = true
     call.timeoutEarlyExit()
     (codec.carrier as RealConnection).useAsSocket()
 
-    eventListener.requestBodyStart(call)
-
     return object : Socket {
       override fun cancel() {
         this@Exchange.cancel()
@@ -233,6 +236,7 @@ class Exchange(
   ) : ForwardingSink(delegate) {
     private var completed = false
     private var bytesReceived = 0L
+    private var invokeStartEvent = isSocket
     private var closed = false
 
     @Throws(IOException::class)
@@ -247,6 +251,10 @@ class Exchange(
         )
       }
       try {
+        if (invokeStartEvent) {
+          invokeStartEvent = false
+          eventListener.requestBodyStart(call)
+        }
         super.write(source, byteCount)
         this.bytesReceived += byteCount
       } catch (e: IOException) {

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

@@ -44,39 +44,41 @@ object CallServerInterceptor : Interceptor {
     try {
       exchange.writeRequestHeaders(request)
 
-      if (hasRequestBody) {
-        // If there's a "Expect: 100-continue" header on the request, wait for a "HTTP/1.1 100
-        // Continue" response before transmitting the request body. If we don't get that, return
-        // what we did get (such as a 4xx response) without ever transmitting the request body.
-        if ("100-continue".equals(request.header("Expect"), ignoreCase = true)) {
-          exchange.flushRequest()
-          responseBuilder = exchange.readResponseHeaders(expectContinue = true)
-          exchange.responseHeadersStart()
-          invokeStartEvent = false
-        }
-        if (responseBuilder == null) {
-          if (requestBody.isDuplex()) {
-            // Prepare a duplex body so that the application can send a request body later.
+      if (!isUpgradeRequest) {
+        if (hasRequestBody) {
+          // If there's a "Expect: 100-continue" header on the request, wait for a "HTTP/1.1 100
+          // Continue" response before transmitting the request body. If we don't get that, return
+          // what we did get (such as a 4xx response) without ever transmitting the request body.
+          if ("100-continue".equals(request.header("Expect"), ignoreCase = true)) {
             exchange.flushRequest()
-            val bufferedRequestBody = exchange.createRequestBody(request, true).buffer()
-            requestBody.writeTo(bufferedRequestBody)
+            responseBuilder = exchange.readResponseHeaders(expectContinue = true)
+            exchange.responseHeadersStart()
+            invokeStartEvent = false
+          }
+          if (responseBuilder == null) {
+            if (requestBody.isDuplex()) {
+              // Prepare a duplex body so that the application can send a request body later.
+              exchange.flushRequest()
+              val bufferedRequestBody = exchange.createRequestBody(request, true).buffer()
+              requestBody.writeTo(bufferedRequestBody)
+            } else {
+              // Write the request body if the "Expect: 100-continue" expectation was met.
+              val bufferedRequestBody = exchange.createRequestBody(request, false).buffer()
+              requestBody.writeTo(bufferedRequestBody)
+              bufferedRequestBody.close()
+            }
           } else {
-            // Write the request body if the "Expect: 100-continue" expectation was met.
-            val bufferedRequestBody = exchange.createRequestBody(request, false).buffer()
-            requestBody.writeTo(bufferedRequestBody)
-            bufferedRequestBody.close()
+            exchange.noRequestBody()
+            if (!exchange.connection.isMultiplexed) {
+              // If the "Expect: 100-continue" expectation wasn't met, prevent the HTTP/1 connection
+              // from being reused. Otherwise we're still obligated to transmit the request body to
+              // leave the connection in a consistent state.
+              exchange.noNewExchangesOnConnection()
+            }
           }
         } else {
           exchange.noRequestBody()
-          if (!exchange.connection.isMultiplexed) {
-            // If the "Expect: 100-continue" expectation wasn't met, prevent the HTTP/1 connection
-            // from being reused. Otherwise we're still obligated to transmit the request body to
-            // leave the connection in a consistent state.
-            exchange.noNewExchangesOnConnection()
-          }
         }
-      } else if (!isUpgradeRequest) {
-        exchange.noRequestBody()
       }
 
       if (requestBody == null || !requestBody.isDuplex()) {

okhttp/src/jvmTest/kotlin/okhttp3/internal/http/HttpUpgradesTest.kt

@@ -208,8 +208,8 @@ class HttpUpgradesTest {
       "RequestHeadersEnd",
       "ResponseHeadersStart",
       "ResponseHeadersEnd",
-      "RequestBodyStart",
       "FollowUpDecision",
+      "RequestBodyStart",
       "ResponseBodyStart",
       "ResponseBodyEnd",
       "RequestBodyEnd",
@@ -233,17 +233,15 @@ class HttpUpgradesTest {
       "ConnectionAcquired",
       "RequestHeadersStart",
       "RequestHeadersEnd",
-      "RequestBodyStart",
-      "RequestBodyEnd",
       "ResponseHeadersStart",
       "ResponseHeadersEnd",
-      "RequestBodyStart",
       "FollowUpDecision",
+      "RequestBodyStart",
       "ResponseBodyStart",
       "ResponseBodyEnd",
+      "RequestBodyEnd",
       "ConnectionReleased",
       "CallEnd",
-      "RequestBodyEnd",
     )
   }