From a1b66acb32d0ebea91e115b8d7f1b42f1dd2507c Mon Sep 17 00:00:00 2001
From: Dave Roberge <david.roberge12@gmail.com>
Date: Sun, 13 Mar 2016 11:42:15 -0400
Subject: [PATCH] Ensure network interceptors always have access to the
 underlying connection.

---
 .../test/java/okhttp3/InterceptorTest.java    | 26 +++++++++++++++++++
 .../okhttp3/internal/http/HttpEngine.java     | 11 +++++---
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/okhttp-tests/src/test/java/okhttp3/InterceptorTest.java b/okhttp-tests/src/test/java/okhttp3/InterceptorTest.java
index dd1f3a977..287f12a31 100644
--- a/okhttp-tests/src/test/java/okhttp3/InterceptorTest.java
+++ b/okhttp-tests/src/test/java/okhttp3/InterceptorTest.java
@@ -26,6 +26,7 @@ import java.util.concurrent.TimeUnit;
 import okhttp3.mockwebserver.MockResponse;
 import okhttp3.mockwebserver.MockWebServer;
 import okhttp3.mockwebserver.RecordedRequest;
+import okhttp3.mockwebserver.SocketPolicy;
 import okio.Buffer;
 import okio.BufferedSink;
 import okio.ForwardingSink;
@@ -630,6 +631,31 @@ public final class InterceptorTest {
     }
   }
 
+  @Test public void networkInterceptorReturnsConnectionOnEmptyBody() throws Exception {
+    server.enqueue(new MockResponse()
+        .setSocketPolicy(SocketPolicy.DISCONNECT_AT_END)
+        .addHeader("Connection", "Close"));
+
+    Interceptor interceptor = new Interceptor() {
+      @Override public Response intercept(Chain chain) throws IOException {
+        Response response = chain.proceed(chain.request());
+        assertNotNull(chain.connection());
+        return response;
+      }
+    };
+
+    client = client.newBuilder()
+        .addNetworkInterceptor(interceptor)
+        .build();
+
+    Request request = new Request.Builder()
+        .url(server.url("/"))
+        .build();
+
+    Response response = client.newCall(request).execute();
+    response.body().close();
+  }
+
   private RequestBody uppercase(final RequestBody original) {
     return new RequestBody() {
       @Override public MediaType contentType() {
diff --git a/okhttp/src/main/java/okhttp3/internal/http/HttpEngine.java b/okhttp/src/main/java/okhttp3/internal/http/HttpEngine.java
index b7c9f6432..b0ca5dbb3 100644
--- a/okhttp/src/main/java/okhttp3/internal/http/HttpEngine.java
+++ b/okhttp/src/main/java/okhttp3/internal/http/HttpEngine.java
@@ -560,7 +560,8 @@ public final class HttpEngine {
       httpStream.writeRequestHeaders(networkRequest);
       networkResponse = readNetworkResponse();
     } else if (!callerWritesRequestBody) {
-      networkResponse = new NetworkInterceptorChain(0, networkRequest).proceed(networkRequest);
+      networkResponse = new NetworkInterceptorChain(0, networkRequest,
+          streamAllocation.connection()).proceed(networkRequest);
     } else {
       // Emit the request body's buffer so that everything is in requestBodyOut.
       if (bufferedRequestBody != null && bufferedRequestBody.buffer().size() > 0) {
@@ -638,15 +639,17 @@ public final class HttpEngine {
   class NetworkInterceptorChain implements Interceptor.Chain {
     private final int index;
     private final Request request;
+    private final Connection connection;
     private int calls;
 
-    NetworkInterceptorChain(int index, Request request) {
+    NetworkInterceptorChain(int index, Request request, Connection connection) {
       this.index = index;
       this.request = request;
+      this.connection = connection;
     }
 
     @Override public Connection connection() {
-      return streamAllocation.connection();
+      return connection;
     }
 
     @Override public Request request() {
@@ -676,7 +679,7 @@ public final class HttpEngine {
 
       if (index < client.networkInterceptors().size()) {
         // There's another interceptor in the chain. Call that.
-        NetworkInterceptorChain chain = new NetworkInterceptorChain(index + 1, request);
+        NetworkInterceptorChain chain = new NetworkInterceptorChain(index + 1, request, connection);
         Interceptor interceptor = client.networkInterceptors().get(index);
         Response interceptedResponse = interceptor.intercept(chain);