From 58928f537078b7fe5d4ef9582a9b3cf5b254adde Mon Sep 17 00:00:00 2001
From: Dave Roberge <david.roberge12@gmail.com>
Date: Tue, 15 Aug 2017 08:46:39 -0400
Subject: [PATCH] Recover from a previous allocated connection with
 noNewStreams.

This is an edge case that can occur with HTTP/2. Since multiple requests
use the same connection, it's possible for one request to flag the
connection as bad during a follow-up request.

https://github.com/square/okhttp/issues/3521
---
 .../internal/http2/HttpOverHttp2Test.java     | 60 ++++++++++++++++++-
 .../internal/connection/StreamAllocation.java | 27 ++++++++-
 2 files changed, 84 insertions(+), 3 deletions(-)

diff --git a/okhttp-tests/src/test/java/okhttp3/internal/http2/HttpOverHttp2Test.java b/okhttp-tests/src/test/java/okhttp3/internal/http2/HttpOverHttp2Test.java
index 320c7d7c8e80..c9db4d03bbf1 100644
--- a/okhttp-tests/src/test/java/okhttp3/internal/http2/HttpOverHttp2Test.java
+++ b/okhttp-tests/src/test/java/okhttp3/internal/http2/HttpOverHttp2Test.java
@@ -29,7 +29,6 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.SynchronousQueue;
-
 import javax.net.ssl.HostnameVerifier;
 import okhttp3.Cache;
 import okhttp3.Call;
@@ -46,6 +45,7 @@
 import okhttp3.Request;
 import okhttp3.RequestBody;
 import okhttp3.Response;
+import okhttp3.Route;
 import okhttp3.TestUtil;
 import okhttp3.internal.DoubleInetAddressDns;
 import okhttp3.internal.RecordingOkAuthenticator;
@@ -736,6 +736,64 @@ private void noRecoveryFromErrorWithRetryDisabled(ErrorCode errorCode) throws Ex
     }
   }
 
+  @Test public void recoverFromConnectionNoNewStreamsOnFollowUp() throws InterruptedException {
+    server.enqueue(new MockResponse()
+        .setResponseCode(401));
+    server.enqueue(new MockResponse()
+        .setSocketPolicy(SocketPolicy.RESET_STREAM_AT_START)
+        .setHttp2ErrorCode(ErrorCode.CANCEL.httpCode));
+    server.enqueue(new MockResponse()
+        .setBody("DEF"));
+    server.enqueue(new MockResponse()
+        .setBody("ABC"));
+
+    final CountDownLatch latch = new CountDownLatch(1);
+    final BlockingQueue<String> responses = new SynchronousQueue<>();
+    okhttp3.Authenticator authenticator = new okhttp3.Authenticator() {
+      @Override public Request authenticate(Route route, Response response) throws IOException {
+        responses.offer(response.body().string());
+        try {
+          latch.await();
+        } catch (InterruptedException e) {
+          throw new AssertionError();
+        }
+        return response.request();
+      }
+    };
+
+    OkHttpClient blockingAuthClient = client.newBuilder()
+        .authenticator(authenticator)
+        .build();
+
+    Callback callback = new Callback() {
+      @Override public void onFailure(Call call, IOException e) {
+        fail();
+      }
+
+      @Override public void onResponse(Call call, Response response) throws IOException {
+        responses.offer(response.body().string());
+      }
+    };
+
+    // Make the first request waiting until we get our auth challenge.
+    Request request = new Request.Builder()
+        .url(server.url("/"))
+        .build();
+    blockingAuthClient.newCall(request).enqueue(callback);
+    String response1 = responses.take();
+    assertEquals("", response1);
+
+    // Now make the second request which will cause the HTTP/2 connection to be flagged as bad.
+    client.newCall(request).enqueue(callback);
+    String response2 = responses.take();
+    assertEquals("DEF", response2);
+
+    // Let the first request proceed.
+    latch.countDown();
+    String response3 = responses.take();
+    assertEquals("ABC", response3);
+  }
+
   @Test public void nonAsciiResponseHeader() throws Exception {
     server.enqueue(new MockResponse()
         .addHeaderLenient("Alpha", "α")
diff --git a/okhttp/src/main/java/okhttp3/internal/connection/StreamAllocation.java b/okhttp/src/main/java/okhttp3/internal/connection/StreamAllocation.java
index 0920121c27ac..770b876350f3 100644
--- a/okhttp/src/main/java/okhttp3/internal/connection/StreamAllocation.java
+++ b/okhttp/src/main/java/okhttp3/internal/connection/StreamAllocation.java
@@ -166,8 +166,8 @@ private RealConnection findConnection(int connectTimeout, int readTimeout, int w
       if (canceled) throw new IOException("Canceled");
 
       // Attempt to use an already-allocated connection.
-      RealConnection allocatedConnection = this.connection;
-      if (allocatedConnection != null && !allocatedConnection.noNewStreams) {
+      RealConnection allocatedConnection = previousAllocatedConnection();
+      if (allocatedConnection != null) {
         return allocatedConnection;
       }
 
@@ -257,6 +257,29 @@ private RealConnection findConnection(int connectTimeout, int readTimeout, int w
     return result;
   }
 
+  /**
+   * Returns the previous allocated connection or null if there is no previous allocated connection.
+   * We will have a previous allocated connection in the case of follow-up requests. With HTTP/2
+   * multiple requests share the same connection so it's possible that our connection is flagged to
+   * disallow new streams during a follow-up request.
+   */
+  private RealConnection previousAllocatedConnection() {
+    assert (Thread.holdsLock(connectionPool));
+
+    RealConnection allocatedConnection = this.connection;
+    if (allocatedConnection == null) return null;
+
+    if (allocatedConnection.noNewStreams) {
+      // Our connection was flagged to disallow new streams. This could happen in HTTP/2 when
+      // multiple requests use the same connection. Discard it.
+      deallocate(false, true, true);
+      return null;
+    }
+
+    // We're good!
+    return allocatedConnection;
+  }
+
   public void streamFinished(boolean noNewStreams, HttpCodec codec) {
     Socket socket;
     Connection releasedConnection;