From d7a09309104f1e59e550760c08753d660acac545 Mon Sep 17 00:00:00 2001 From: Jake Wharton Date: Sat, 14 Nov 2015 18:49:41 -0500 Subject: [PATCH] Force consumers to call WS#close on IO error. --- .../okhttp/internal/ws/RealWebSocketTest.java | 77 +++++++------------ .../okhttp/internal/ws/RealWebSocket.java | 23 +++--- .../com/squareup/okhttp/ws/WebSocket.java | 11 ++- 3 files changed, 41 insertions(+), 70 deletions(-) diff --git a/okhttp-ws-tests/src/test/java/com/squareup/okhttp/internal/ws/RealWebSocketTest.java b/okhttp-ws-tests/src/test/java/com/squareup/okhttp/internal/ws/RealWebSocketTest.java index a70435fbf..8da32b3a6 100644 --- a/okhttp-ws-tests/src/test/java/com/squareup/okhttp/internal/ws/RealWebSocketTest.java +++ b/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 { diff --git a/okhttp-ws/src/main/java/com/squareup/okhttp/internal/ws/RealWebSocket.java b/okhttp-ws/src/main/java/com/squareup/okhttp/internal/ws/RealWebSocket.java index dd5e96c2d..ea55b5a85 100644 --- a/okhttp-ws/src/main/java/com/squareup/okhttp/internal/ws/RealWebSocket.java +++ b/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; } } diff --git a/okhttp-ws/src/main/java/com/squareup/okhttp/ws/WebSocket.java b/okhttp-ws/src/main/java/com/squareup/okhttp/ws/WebSocket.java index eb1df7394..a3eebe7d1 100644 --- a/okhttp-ws/src/main/java/com/squareup/okhttp/ws/WebSocket.java +++ b/okhttp-ws/src/main/java/com/squareup/okhttp/ws/WebSocket.java @@ -33,8 +33,8 @@ public interface WebSocket { *

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;