From 62f2f823410661128d1d9549b9b727f824637f6c Mon Sep 17 00:00:00 2001
From: Jesse Wilson <jwilson@squareup.com>
Date: Tue, 6 Nov 2018 22:36:54 +1100
Subject: [PATCH] Fix connection leaks on failed web socket upgrades.

Closes: https://github.com/square/okhttp/issues/4258
---
 .../internal/ws/WebSocketHttpTest.java        | 32 +++++++++++++++++++
 .../http/RetryAndFollowUpInterceptor.java     |  4 +--
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/okhttp-tests/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java b/okhttp-tests/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java
index 4ef795d33471..1e17baa9c2ab 100644
--- a/okhttp-tests/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java
+++ b/okhttp-tests/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java
@@ -17,8 +17,10 @@
 
 import java.io.EOFException;
 import java.io.IOException;
+import java.net.HttpURLConnection;
 import java.net.ProtocolException;
 import java.net.SocketTimeoutException;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Random;
 import java.util.concurrent.CountDownLatch;
@@ -27,6 +29,7 @@
 import java.util.logging.Logger;
 import okhttp3.Interceptor;
 import okhttp3.OkHttpClient;
+import okhttp3.Protocol;
 import okhttp3.RecordingEventListener;
 import okhttp3.RecordingHostnameVerifier;
 import okhttp3.Request;
@@ -695,6 +698,35 @@ public final class WebSocketHttpTest {
     clientListener.assertTextMessage("Hello, WebSockets!");
   }
 
+  /**
+   * We had a bug where web socket connections were leaked if the HTTP connection upgrade was not
+   * successful. This test confirms that connections are released back to the connection pool!
+   * https://github.com/square/okhttp/issues/4258
+   */
+  @Test public void webSocketConnectionIsReleased() throws Exception {
+    // This test assumes HTTP/1.1 pooling semantics.
+    client = client.newBuilder()
+        .protocols(Arrays.asList(Protocol.HTTP_1_1))
+        .build();
+
+    webServer.enqueue(new MockResponse()
+        .setResponseCode(HttpURLConnection.HTTP_NOT_FOUND)
+        .setBody("not found!"));
+    webServer.enqueue(new MockResponse());
+
+    newWebSocket();
+    clientListener.assertFailure();
+
+    Request regularRequest = new Request.Builder()
+        .url(webServer.url("/"))
+        .build();
+    Response response = client.newCall(regularRequest).execute();
+    response.close();
+
+    assertEquals(0, webServer.takeRequest().getSequenceNumber());
+    assertEquals(1, webServer.takeRequest().getSequenceNumber());
+  }
+
   private MockResponse upgradeResponse(RecordedRequest request) {
     String key = request.getHeader("Sec-WebSocket-Key");
     return new MockResponse()
diff --git a/okhttp/src/main/java/okhttp3/internal/http/RetryAndFollowUpInterceptor.java b/okhttp/src/main/java/okhttp3/internal/http/RetryAndFollowUpInterceptor.java
index ffb75058aa7b..d682de8bf04f 100644
--- a/okhttp/src/main/java/okhttp3/internal/http/RetryAndFollowUpInterceptor.java
+++ b/okhttp/src/main/java/okhttp3/internal/http/RetryAndFollowUpInterceptor.java
@@ -164,9 +164,7 @@ public StreamAllocation streamAllocation() {
       }
 
       if (followUp == null) {
-        if (!forWebSocket) {
-          streamAllocation.release();
-        }
+        streamAllocation.release();
         return response;
       }