Retain response body when web socket returns non-101 response code.

This change also fixes an edge case where a server returns a response
that allows a follow up request. Previously, this would cause an
exception because the body was being ignored.

okhttp-ws-tests/src/test/java/okhttp3/ws/WebSocketCallTest.java

@@ -88,17 +88,29 @@ public final class WebSocketCallTest {
     listener.assertTextMessage("Hello, WebSockets!");
   }
 
-  @Test public void okButNotOk() {
-    server.enqueue(new MockResponse().setResponseCode(200));
+  @Test public void non101RetainsBody() throws IOException {
+    server.enqueue(new MockResponse().setResponseCode(200).setBody("Body"));
     awaitWebSocket();
     listener.assertFailure(ProtocolException.class, "Expected HTTP 101 response but was '200 OK'");
+    listener.assertResponse(200, "Body");
   }
 
-  @Test public void notFound() {
+  @Test public void notFound() throws IOException {
     server.enqueue(new MockResponse().setStatus("HTTP/1.1 404 Not Found"));
     awaitWebSocket();
     listener.assertFailure(ProtocolException.class,
         "Expected HTTP 101 response but was '404 Not Found'");
+    listener.assertResponse(404, "");
+  }
+
+  @Test public void clientTimeoutClosesBody() throws IOException {
+    server.enqueue(new MockResponse().setResponseCode(408));
+    WebSocketListener serverListener = new EmptyWebSocketListener();
+    server.enqueue(new MockResponse().withWebSocketUpgrade(serverListener));
+
+    WebSocket webSocket = awaitWebSocket();
+    webSocket.sendPing(new Buffer().writeUtf8("WebSockets are fun!"));
+    listener.assertPong(new Buffer().writeUtf8("WebSockets are fun!"));
   }
 
   @Test public void missingConnectionHeader() {
@@ -236,7 +248,7 @@ private WebSocket awaitWebSocket(Request request) {
       }
 
       @Override public void onFailure(IOException e, Response response) {
-        listener.onFailure(e, null);
+        listener.onFailure(e, response);
         failureRef.set(e);
         latch.countDown();
       }

okhttp-ws-tests/src/test/java/okhttp3/ws/WebSocketRecorder.java

@@ -38,6 +38,7 @@ public interface MessageDelegate {
 
   private final BlockingQueue<Object> events = new LinkedBlockingQueue<>();
   private MessageDelegate delegate;
+  private Response response;
 
   /** Sets a delegate for the next call to {@link #onMessage}. Cleared after invoked. */
   public void setNextMessageDelegate(MessageDelegate delegate) {
@@ -73,6 +74,7 @@ public void setNextMessageDelegate(MessageDelegate delegate) {
 
   @Override public void onFailure(IOException e, Response response) {
     events.add(e);
+    this.response = response;
   }
 
   private Object nextEvent() {
@@ -144,6 +146,12 @@ public void assertExhausted() {
     assertTrue("Remaining events: " + events, events.isEmpty());
   }
 
+  public void assertResponse(int code, String body) throws IOException {
+    assertNotNull(response);
+    assertEquals(code, response.code());
+    assertEquals(body, response.body().string());
+  }
+
   private static class Message {
     public final MediaType mediaType;
     public final Buffer buffer = new Buffer();

okhttp-ws/src/main/java/okhttp3/ws/WebSocketCall.java

@@ -115,7 +115,6 @@ public void cancel() {
 
   private void createWebSocket(Response response, WebSocketListener listener) throws IOException {
     if (response.code() != 101) {
-      Util.closeQuietly(response.body());
       throw new ProtocolException("Expected HTTP 101 response but was '"
           + response.code()
           + " "

okhttp/src/main/java/okhttp3/internal/http/HttpEngine.java

@@ -779,7 +779,7 @@ private Response readNetworkResponse() throws IOException {
         .receivedResponseAtMillis(System.currentTimeMillis())
         .build();
 
-    if (!forWebSocket) {
+    if (!forWebSocket || networkResponse.code() != 101) {
       networkResponse = networkResponse.newBuilder()
           .body(httpStream.openResponseBody(networkResponse))
           .build();