Merge pull request #1999 from square/jw/consumers-close

Force consumers to call WS#close on IO error.

okhttp-ws-tests/src/test/java/com/squareup/okhttp/internal/ws/RealWebSocketTest.java

@@ -40,7 +40,6 @@ import static com.squareup.okhttp.ws.WebSocket.TEXT;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -219,37 +218,6 @@ public final class RealWebSocketTest {
     }
   }
 
-  @Test public void clientWritingThrowsSendsClientClose() throws IOException {
-    final IOException brokenException = new IOException("Broken!");
-    RequestBody brokenBody = new RequestBody() {
-      @Override public MediaType contentType() {
-        return TEXT;
-      }
-
-      @Override public void writeTo(BufferedSink sink) throws IOException {
-        throw brokenException;
-      }
-    };
-
-    try {
-      client.sendMessage(brokenBody);
-      fail();
-    } catch (IOException e) {
-      assertSame(brokenException, e);
-    }
-
-    // A failed write prevents further use of the WebSocket instance.
-    try {
-      client.sendPing(new Buffer().writeUtf8("Ping!"));
-      fail();
-    } catch (IllegalStateException e) {
-      assertEquals("closed", e.getMessage());
-    }
-
-    server.readMessage();
-    serverListener.assertClose(1001, "");
-  }
-
   @Test public void socketClosedDuringPingKillsWebSocket() throws IOException {
     client2Server.close();
 
@@ -260,11 +228,17 @@ public final class RealWebSocketTest {
     }
 
     // A failed write prevents further use of the WebSocket instance.
+    try {
+      client.sendMessage(RequestBody.create(TEXT, "Hello!"));
+      fail();
+    } catch (IllegalStateException e) {
+      assertEquals("must call close()", e.getMessage());
+    }
     try {
       client.sendPing(new Buffer().writeUtf8("Ping!"));
       fail();
     } catch (IllegalStateException e) {
-      assertEquals("closed", e.getMessage());
+      assertEquals("must call close()", e.getMessage());
     }
   }
 
@@ -279,28 +253,16 @@ public final class RealWebSocketTest {
 
     // A failed write prevents further use of the WebSocket instance.
     try {
-      client.sendPing(new Buffer().writeUtf8("Ping!"));
+      client.sendMessage(RequestBody.create(TEXT, "Hello!"));
       fail();
     } catch (IllegalStateException e) {
-      assertEquals("closed", e.getMessage());
+      assertEquals("must call close()", e.getMessage());
     }
-  }
-
-  @Test public void socketClosedDuringCloseKillsWebSocket() throws IOException {
-    client2Server.close();
-
-    try {
-      client.close(1000, "I'm done.");
-      fail();
-    } catch (IOException ignored) {
-    }
-
-    // A failed write prevents further use of the WebSocket instance.
     try {
       client.sendPing(new Buffer().writeUtf8("Ping!"));
       fail();
     } catch (IllegalStateException e) {
-      assertEquals("closed", e.getMessage());
+      assertEquals("must call close()", e.getMessage());
     }
   }
 
@@ -388,7 +350,7 @@ public final class RealWebSocketTest {
   @Test public void serverCloseClosesConnection() throws IOException {
     server.close(1000, "Hello!");
 
-    client.readMessage(); // Read server close, sense client close, close connection.
+    client.readMessage(); // Read server close, send client close, close connection.
     assertTrue(clientConnectionClosed);
     clientListener.assertClose(1000, "Hello!");
 
@@ -426,8 +388,8 @@ public final class RealWebSocketTest {
     server2client.raw().write(ByteString.decodeHex("0a00")); // Invalid non-final ping frame.
 
     client.readMessage(); // Detects error, send close, close connection.
-    clientListener.assertFailure(ProtocolException.class, "Control frames must be final.");
     assertTrue(clientConnectionClosed);
+    clientListener.assertFailure(ProtocolException.class, "Control frames must be final.");
 
     server.readMessage();
     serverListener.assertClose(1002, "");
@@ -439,11 +401,24 @@ public final class RealWebSocketTest {
     server2client.raw().write(ByteString.decodeHex("0a00")); // Invalid non-final ping frame.
 
     client.readMessage(); // Detects error, closes connection immediately since close already sent.
-    clientListener.assertFailure(ProtocolException.class, "Control frames must be final.");
     assertTrue(clientConnectionClosed);
+    clientListener.assertFailure(ProtocolException.class, "Control frames must be final.");
 
     server.readMessage();
     serverListener.assertClose(1000, "Hello!");
+
+    serverListener.assertExhausted(); // Client should not have sent second close.
+  }
+
+  @Test public void closeThrowingClosesConnection() {
+    client2Server.close();
+
+    try {
+      client.close(1000, null);
+      fail();
+    } catch (IOException ignored) {
+    }
+    assertTrue(clientConnectionClosed);
   }
 
   @Test public void closeMessageAndConnectionCloseThrowingDoesNotMaskOriginal() throws IOException {

okhttp-ws/src/main/java/com/squareup/okhttp/internal/ws/RealWebSocket.java

@@ -36,7 +36,6 @@ import static com.squareup.okhttp.internal.ws.WebSocketProtocol.OPCODE_TEXT;
 import static com.squareup.okhttp.internal.ws.WebSocketReader.FrameCallback;
 
 public abstract class RealWebSocket implements WebSocket {
-  private static final int CLOSE_GOING_AWAY = 1001;
   private static final int CLOSE_PROTOCOL_EXCEPTION = 1002;
 
   private final WebSocketWriter writer;
@@ -45,6 +44,8 @@ public abstract class RealWebSocket implements WebSocket {
 
   /** True after calling {@link #close(int, String)}. No writes are allowed afterward. */
   private volatile boolean writerSentClose;
+  /** True after {@link IOException}. {@link #close(int, String)} becomes only valid call. */
+  private boolean writerWantsClose;
   /** True after a close frame was read by the reader. No frames will follow it. */
   private boolean readerSentClose;
 
@@ -104,6 +105,7 @@ public abstract class RealWebSocket implements WebSocket {
   @Override public void sendMessage(RequestBody message) throws IOException {
     if (message == null) throw new NullPointerException("message == null");
     if (writerSentClose) throw new IllegalStateException("closed");
+    if (writerWantsClose) throw new IllegalStateException("must call close()");
 
     MediaType contentType = message.contentType();
     if (contentType == null) {
@@ -128,23 +130,19 @@ public abstract class RealWebSocket implements WebSocket {
       message.writeTo(sink);
       sink.close();
     } catch (IOException e) {
-      try {
-        close(CLOSE_GOING_AWAY, null);
-      } catch (IOException ignored) {
-      }
+      writerWantsClose = true;
       throw e;
     }
   }
 
   @Override public void sendPing(Buffer payload) throws IOException {
     if (writerSentClose) throw new IllegalStateException("closed");
+    if (writerWantsClose) throw new IllegalStateException("must call close()");
+
     try {
       writer.writePing(payload);
     } catch (IOException e) {
-      try {
-        close(CLOSE_GOING_AWAY, null);
-      } catch (IOException ignored) {
-      }
+      writerWantsClose = true;
       throw e;
     }
   }
@@ -152,13 +150,12 @@ public abstract class RealWebSocket implements WebSocket {
   /** Send an unsolicited pong with the specified payload. */
   public void sendPong(Buffer payload) throws IOException {
     if (writerSentClose) throw new IllegalStateException("closed");
+    if (writerWantsClose) throw new IllegalStateException("must call close()");
+
     try {
       writer.writePong(payload);
     } catch (IOException e) {
-      try {
-        close(CLOSE_GOING_AWAY, null);
-      } catch (IOException ignored) {
-      }
+      writerWantsClose = true;
       throw e;
     }
   }

okhttp-ws/src/main/java/com/squareup/okhttp/ws/WebSocket.java

@@ -33,8 +33,8 @@ public interface WebSocket {
    * <p>The {@linkplain RequestBody#contentType() content type} of {@code message} should be either
    * {@link #TEXT} or {@link #BINARY}.
    *
-   * @throws IOException if unable to write the message. If thrown, an attempt to send a close
-   * frame with code 1001 will be made and this instance will no longer be valid to use.
+   * @throws IOException if unable to write the message. Clients must call {@link #close} when this
+   * happens to ensure resources are cleaned up.
    * @throws IllegalStateException if not connected, already closed, or another writer is active.
    */
   void sendMessage(RequestBody message) throws IOException;
@@ -42,8 +42,8 @@ public interface WebSocket {
   /**
    * Send a ping to the server with optional payload.
    *
-   * @throws IOException if unable to write the ping. If thrown, an attempt to send a close frame
-   * with code 1001 will be made and this instance will no longer be valid to use.
+   * @throws IOException if unable to write the ping.  Clients must call {@link #close} when this
+   * happens to ensure resources are cleaned up.
    * @throws IllegalStateException if already closed.
    */
   void sendPing(Buffer payload) throws IOException;
@@ -57,8 +57,7 @@ public interface WebSocket {
    * It is an error to call this method before calling close on an active writer. Calling this
    * method more than once has no effect.
    *
-   * @throws IOException if unable to write the close message. If thrown this instance will no
-   * longer be valid to use.
+   * @throws IOException if unable to write the close message. Resources will still be freed.
    * @throws IllegalStateException if already closed.
    */
   void close(int code, String reason) throws IOException;