Ensure network interceptors always have access to the underlying conn…

…ection.
yang12sanpi · Mar 13, 2016 · a1b66ac · a1b66ac
1 parent fb3c390
commit a1b66ac
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 4 deletions.
diff --git a/okhttp-tests/src/test/java/okhttp3/InterceptorTest.java b/okhttp-tests/src/test/java/okhttp3/InterceptorTest.java
@@ -26,6 +26,7 @@
 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 @@ private void interceptorThrowsRuntimeExceptionAsynchronous(boolean network) thro
     }
   }
 
+  @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
@@ -560,7 +560,8 @@ public void readResponse() throws IOException {
       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 void readResponse() throws IOException {
   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 @@ class NetworkInterceptorChain implements Interceptor.Chain {
 
       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);