From 997881197ac3d43bbc5b5e451c13010cb74a23e2 Mon Sep 17 00:00:00 2001
From: Yuri Schimke <yuri@schimke.ee>
Date: Sun, 10 Apr 2022 12:49:25 +0100
Subject: [PATCH] Fix Conscrypt NPE workaround (#7219)

(cherry picked from commit 6ba23dcd710dac2589d8740086092cda988ec7a4)
---
 .../okhttp/android/envoy/EnvoyInterceptor.kt  | 30 +++++++++++++-----
 .../okhttp/android/envoy/EnvoyMobileTest.kt   | 31 +++++++++++++++++--
 2 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/android-test/src/androidTest/java/okhttp/android/envoy/EnvoyInterceptor.kt b/android-test/src/androidTest/java/okhttp/android/envoy/EnvoyInterceptor.kt
index 0cf058dce971..2b18929a8e75 100644
--- a/android-test/src/androidTest/java/okhttp/android/envoy/EnvoyInterceptor.kt
+++ b/android-test/src/androidTest/java/okhttp/android/envoy/EnvoyInterceptor.kt
@@ -22,7 +22,6 @@ import io.envoyproxy.envoymobile.RequestHeadersBuilder
 import io.envoyproxy.envoymobile.RequestMethod
 import io.envoyproxy.envoymobile.StreamPrototype
 import java.io.IOException
-import java.lang.IllegalArgumentException
 import java.nio.ByteBuffer
 import java.util.concurrent.Executors
 import kotlin.coroutines.cancellation.CancellationException
@@ -31,11 +30,15 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.suspendCancellableCoroutine
-import okhttp3.*
+import okhttp3.Headers
+import okhttp3.Interceptor
+import okhttp3.MediaType
 import okhttp3.MediaType.Companion.toMediaTypeOrNull
+import okhttp3.Protocol
+import okhttp3.Request
+import okhttp3.Response
 import okhttp3.ResponseBody.Companion.asResponseBody
 import okio.Buffer
-import okio.Okio
 import okio.Pipe
 import okio.buffer
 
@@ -52,7 +55,6 @@ suspend fun makeRequest(engine: Engine, request: Request) =
   suspendCancellableCoroutine<Response> { continuation ->
     val responseBuilder = Response.Builder()
       .request(request)
-      .protocol(if (request.isHttps) Protocol.QUIC else Protocol.HTTP_1_1)
       .sentRequestAtMillis(System.currentTimeMillis())
 
     val bodyPipe = Pipe(1024L * 1024L)
@@ -69,7 +71,16 @@ suspend fun makeRequest(engine: Engine, request: Request) =
 
         contentType = headers["content-type"]?.toMediaTypeOrNull()
 
+        // TODO check this logic
+        val alpn = headers["x-envoy-upstream-alpn"]
+        val protocol = if (alpn != null) {
+          Protocol.get(alpn)
+        } else {
+          if (request.isHttps) Protocol.QUIC else Protocol.HTTP_1_1
+        }
+
         responseBuilder
+          .protocol(protocol)
           .code(responseHeaders.httpStatus ?: 0)
           .message(responseHeaders.httpStatus.toString())
           .receivedResponseAtMillis(System.currentTimeMillis())
@@ -85,12 +96,13 @@ suspend fun makeRequest(engine: Engine, request: Request) =
         println("Dropping trailers " + responseTrailers.toHeaders())
       }
       .setOnResponseData { data, endStream, streamIntel ->
-        bodySink.write(data)
-
         if (endStream) {
           bodySink.close()
         }
       }
+      .setOnComplete { streamIntel, finalStreamIntel ->
+        bodySink.close()
+      }
       .setOnError { error, streamIntel, finalStreamIntel ->
         // TODO how to signal error correctly?
         bodySource.close()
@@ -114,6 +126,8 @@ suspend fun makeRequest(engine: Engine, request: Request) =
       request.url.host,
       request.url.encodedPath
     ).apply {
+      // TODO addH2RawDomains for Protocol.H2C?
+
       request.headers.toMultimap().forEach { (name, values) ->
         values.forEach { value ->
           add(name, value)
@@ -143,7 +157,9 @@ suspend fun makeRequest(engine: Engine, request: Request) =
         .sendHeaders(requestHeaders.build(), endStream = true)
     }
 
-    continuation.invokeOnCancellation { stream.cancel() }
+    continuation.invokeOnCancellation {
+      stream.cancel()
+    }
   }
 
 private fun io.envoyproxy.envoymobile.Headers.toHeaders(): Headers {
diff --git a/android-test/src/androidTest/java/okhttp/android/envoy/EnvoyMobileTest.kt b/android-test/src/androidTest/java/okhttp/android/envoy/EnvoyMobileTest.kt
index 3520ea9fb568..7ca4fe4d21e7 100644
--- a/android-test/src/androidTest/java/okhttp/android/envoy/EnvoyMobileTest.kt
+++ b/android-test/src/androidTest/java/okhttp/android/envoy/EnvoyMobileTest.kt
@@ -39,6 +39,7 @@ import okhttp3.HttpUrl
 import okhttp3.HttpUrl.Companion.toHttpUrl
 import okhttp3.MediaType.Companion.toMediaType
 import okhttp3.OkHttpClient
+import okhttp3.Protocol
 import okhttp3.Request
 import okhttp3.RequestBody.Companion.toRequestBody
 import okhttp3.Response
@@ -67,9 +68,9 @@ class EnvoyMobileTest {
     val application = ApplicationProvider.getApplicationContext<Application>()
 
     engine = AndroidEngineBuilder(application, baseConfiguration = Standard())
-      .addLogLevel(LogLevel.TRACE)
-      .setOnEngineRunning { println("Envoy async internal setup completed") }
+      .addLogLevel(LogLevel.INFO)
       .setLogger { println(it) }
+      // .enableHappyEyeballs(true)
       .build()
 
     client = OkHttpClient.Builder()
@@ -95,6 +96,31 @@ class EnvoyMobileTest {
     response.use {
       printResponse(response)
     }
+
+    assertEquals(Protocol.QUIC, response.protocol)
+  }
+
+  @Test
+  fun get2() = runTest {
+    val client = OkHttpClient.Builder()
+      .addInterceptor(EnvoyInterceptor(engine))
+      .build()
+
+    val getRequest = Request(url = aiortc + "get")
+
+    val response = client.newCall(getRequest).executeAsync()
+
+    response.use {
+      printResponse(response)
+    }
+
+    val response2 = client.newCall(getRequest).executeAsync()
+
+    response.use {
+      printResponse(response2)
+    }
+
+    assertEquals(Protocol.QUIC, response.protocol)
   }
 
   @Test
@@ -117,6 +143,7 @@ class EnvoyMobileTest {
   }
 
   @Test
+  @Disabled
   fun cancel() = runTest {
     val client = OkHttpClient.Builder()
       .addInterceptor(EnvoyInterceptor(engine))