Add option to send and received close code and reason

Author: gusavila92

src/main/java/dev/gustavoavila/websocketclient/WebSocketClient.java

@@ -2,8 +2,8 @@
 
 import dev.gustavoavila.websocketclient.common.Utils;
 import dev.gustavoavila.websocketclient.exceptions.IllegalSchemeException;
+import dev.gustavoavila.websocketclient.exceptions.InvalidReceivedFrameException;
 import dev.gustavoavila.websocketclient.exceptions.InvalidServerHandshakeException;
-import dev.gustavoavila.websocketclient.exceptions.UnknownOpcodeException;
 import dev.gustavoavila.websocketclient.model.Payload;
 
 import java.io.BufferedInputStream;
@@ -17,11 +17,7 @@
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.security.SecureRandom;
-import java.util.HashMap;
-import java.util.LinkedList;
-import java.util.Map;
-import java.util.Random;
-import java.util.Queue;
+import java.util.*;
 
 import javax.net.SocketFactory;
 import javax.net.ssl.SSLSocketFactory;
@@ -32,6 +28,8 @@
  * @author Gustavo Avila
  */
 public abstract class WebSocketClient {
+    public static final int CLOSE_CODE_NORMAL = 1000;
+
     /**
      * Max number of response handshake bytes to read before raising an exception
      */
@@ -137,6 +135,8 @@ public abstract class WebSocketClient {
      */
     private SSLSocketFactory sslSocketFactory;
 
+    private volatile Timer closeTimer;
+
     /**
      * Initialize all the variables
      *
@@ -199,7 +199,7 @@ public WebSocketClient(URI uri) {
     /**
      * Called when a close code has been received
      */
-    public abstract void onCloseReceived();
+    public abstract void onCloseReceived(int reason, String description);
 
     /**
      * Adds a new header to the set of headers that will be send into the
@@ -449,14 +449,31 @@ private void notifyOnException(Exception e) {
     /**
      * If the close method wasn't called, call onCloseReceived() method.
      */
-    private void notifyOnCloseReceived() {
+    private void notifyOnCloseReceived(int reason, String description) {
         synchronized (globalLock) {
             if (isRunning) {
-                onCloseReceived();
+                onCloseReceived(reason, description);
             }
         }
     }
 
+    private void forceClose() {
+        new Thread(new Runnable() {
+            @Override
+            public void run() {
+                synchronized (globalLock) {
+                    isRunning = false;
+
+                    if (reconnectionThread != null) {
+                        reconnectionThread.interrupt();
+                    }
+
+                    webSocketConnection.closeInternal();
+                }
+            }
+        }).start();
+    }
+
     /**
      * Sends a text message If the WebSocket is not connected yet, message will
      * be send the next time the connection is opened
@@ -465,7 +482,7 @@ private void notifyOnCloseReceived() {
      */
     public void send(String message) {
         byte[] data = message.getBytes(Charset.forName("UTF-8"));
-        final Payload payload = new Payload(OPCODE_TEXT, data);
+        final Payload payload = new Payload(OPCODE_TEXT, data, false);
 
         new Thread(new Runnable() {
             @Override
@@ -483,7 +500,7 @@ public void run() {
      * @param data Binary data that will be send to the WebSocket server
      */
     public void send(byte[] data) {
-        final Payload payload = new Payload(OPCODE_BINARY, data);
+        final Payload payload = new Payload(OPCODE_BINARY, data, false);
 
         new Thread(new Runnable() {
             @Override
@@ -503,7 +520,7 @@ public void sendPing(byte[] data) {
             throw new IllegalArgumentException("Control frame payload cannot be greater than 125 bytes");
         }
 
-        final Payload payload = new Payload(OPCODE_PING, data);
+        final Payload payload = new Payload(OPCODE_PING, data, false);
         new Thread(new Runnable() {
             @Override
             public void run() {
@@ -522,7 +539,7 @@ public void sendPong(byte[] data) {
             throw new IllegalArgumentException("Control frame payload cannot be greater than 125 bytes");
         }
 
-        final Payload payload = new Payload(OPCODE_PONG, data);
+        final Payload payload = new Payload(OPCODE_PONG, data, false);
         new Thread(new Runnable() {
             @Override
             public void run() {
@@ -534,21 +551,40 @@ public void run() {
     /**
      * Closes the WebSocket connection
      */
-    public void close() {
-        new Thread(new Runnable() {
-            @Override
-            public void run() {
-                synchronized (globalLock) {
-                    isRunning = false;
+    public void close(final int timeout, int code, String reason) {
+        if (timeout == 0) {
+            forceClose();
+        } else if (code < 0 || code >= 5000) {
+            throw new IllegalArgumentException("Close frame code must be greater or equal than zero and less than 5000");
+        } else {
+            byte[] internalReason = new byte[0];
+            if (reason != null) {
+                internalReason = reason.getBytes(Charset.forName("UTF-8"));
+                if (internalReason.length > 123) {
+                    throw new IllegalArgumentException("Close frame reason is too large");
+                }
+            }
 
-                    if (reconnectionThread != null) {
-                        reconnectionThread.interrupt();
-                    }
+            byte[] codeLength = Utils.to2ByteArray(code);
+            byte[] data = Arrays.copyOf(codeLength, 2 + internalReason.length);
+            System.arraycopy(internalReason, 0, data, codeLength.length, internalReason.length);
 
-                    webSocketConnection.closeInternal();
+            final Payload payload = new Payload(OPCODE_CLOSE, data, false);
+            new Thread(new Runnable() {
+                @Override
+                public void run() {
+                    webSocketConnection.sendInternal(payload);
                 }
-            }
-        }).start();
+            }).start();
+
+            closeTimer = new Timer();
+            closeTimer.schedule(new TimerTask() {
+                @Override
+                public void run() {
+                    forceClose();
+                }
+            }, timeout);
+        }
     }
 
     /**
@@ -569,10 +605,15 @@ private class WebSocketConnection {
          */
         private volatile boolean isClosed;
 
+        /**
+         * Flag that indicates that a graceful close is in process
+         */
+        private volatile boolean isClosing;
+
         /**
          * Data waiting to be read from the writer thread
          */
-        private final LinkedList<Payload> outBuffer;
+        private final Queue<Payload> queue;
 
         /**
          * This will act as a lock for synchronized statements
@@ -606,7 +647,8 @@ private class WebSocketConnection {
         private WebSocketConnection() {
             this.pendingMessages = false;
             this.isClosed = false;
-            this.outBuffer = new LinkedList<Payload>();
+            this.isClosing = false;
+            this.queue = new LinkedList<Payload>();
             this.internalLock = new Object();
 
             this.writerThread = new Thread(new Runnable() {
@@ -627,13 +669,17 @@ public void run() {
                             if (socket.isClosed()) {
                                 return;
                             } else {
-                                while (outBuffer.size() > 0) {
-                                    Payload payload = outBuffer.removeFirst();
+                                while (queue.size() > 0) {
+                                    Payload payload = queue.poll();
                                     int opcode = payload.getOpcode();
                                     byte[] data = payload.getData();
 
                                     try {
                                         send(opcode, data);
+
+                                        if (payload.isCloseEcho()) {
+                                            closeInternalInsecure();
+                                        }
                                     } catch (IOException e) {
                                         // Reader thread will notify this
                                         // exception
@@ -1125,9 +1171,29 @@ private void read() throws IOException {
                         notifyOnBinaryReceived(data);
                         break;
                     case OPCODE_CLOSE:
-                        closeInternal();
-                        notifyOnCloseReceived();
-                        return;
+                        if (data.length > 125) {
+                            closeInternal();
+                            Exception e = new InvalidReceivedFrameException("Close frame payload is too big");
+                            notifyOnException(e);
+                            return;
+                        } else {
+                            int code = getCloseCode(data);
+                            String reason = getCloseReason(data);
+                            notifyOnCloseReceived(code, reason);
+                        }
+
+                        synchronized (internalLock) {
+                            if (isClosing) {
+                                // This is the echo of a client initiated close so the connection can be closed immediately
+                                closeInternalInsecure();
+                                return;
+                            } else {
+                                // This is a server initiated close so an echo must be sent
+                                Payload payload = new Payload(OPCODE_CLOSE, data, true);
+                                sendInternalInsecure(payload);
+                                break;
+                            }
+                        }
                     case OPCODE_PING:
                         notifyOnPingReceived(data);
                         sendPong(data);
@@ -1137,28 +1203,38 @@ private void read() throws IOException {
                         break;
                     default:
                         closeInternal();
-                        Exception e = new UnknownOpcodeException("Unknown opcode: 0x" + Integer.toHexString(opcode));
+                        Exception e = new InvalidReceivedFrameException("Unknown opcode: 0x" + Integer.toHexString(opcode));
                         notifyOnException(e);
                         return;
                 }
             }
 
-            // If there are not more data to be read,
-            // and if the connection didn't receive a close frame,
-            // an IOException must be thrown because the connection didn't close
-            // gracefully
-            throw new IOException("Unexpected end of stream");
+            synchronized (internalLock) {
+                // There is no need to notify an exception if the connection is closing
+                if (!isClosing) {
+                    // An IOException must be thrown because the connection didn't close gracefully
+                    throw new IOException("Unexpected end of stream");
+                }
+            }
         }
 
         /**
-         * Puts the payload into the out buffer and notifies the writer thread
-         * that new data is available
+         * Puts the payload into the out queue and notifies the writer thread that new data is available
          *
-         * @param payload Payload to be send to the WebSocket server
+         * @param payload Payload to be sent to the WebSocket server
          */
         private void sendInternal(Payload payload) {
             synchronized (internalLock) {
-                outBuffer.addLast(payload);
+                sendInternalInsecure(payload);
+            }
+        }
+
+        private void sendInternalInsecure(Payload payload) {
+            if (!isClosing) {
+                if (payload.getOpcode() == OPCODE_CLOSE) {
+                    isClosing = true;
+                }
+                queue.offer(payload);
                 pendingMessages = true;
                 internalLock.notify();
             }
@@ -1169,20 +1245,45 @@ private void sendInternal(Payload payload) {
          * thread and the reconnection thread that they must finish
          */
         private void closeInternal() {
+            synchronized (internalLock) {
+                closeInternalInsecure();
+            }
+        }
+
+        private void closeInternalInsecure() {
             try {
-                synchronized (internalLock) {
-                    if (!isClosed) {
-                        isClosed = true;
-                        if (socket != null) {
-                            socket.close();
-                            pendingMessages = true;
-                            internalLock.notify();
-                        }
+                if (!isClosed) {
+                    isClosed = true;
+                    if (socket != null) {
+                        socket.close();
+                        pendingMessages = true;
+                        internalLock.notify();
                     }
                 }
+
+                if (closeTimer != null) {
+                    closeTimer.cancel();
+                    closeTimer = null;
+                }
             } catch (IOException e) {
                 // This should never happen
             }
         }
+
+        private int getCloseCode(byte[] data) {
+            if (data.length > 1) {
+                byte[] baseCode = Arrays.copyOfRange(data, 0, 2);
+                return Utils.fromByteArray(new byte[]{0, 0, baseCode[0], baseCode[1]});
+            }
+            return -1;
+        }
+
+        private String getCloseReason(byte[] data) {
+            if (data.length > 2) {
+                byte[] baseReason = Arrays.copyOfRange(data, 2, data.length);
+                return new String(baseReason, Charset.forName("UTF-8"));
+            }
+            return null;
+        }
     }
 }

src/main/java/dev/gustavoavila/websocketclient/exceptions/InvalidReceivedFrameException.java

@@ -0,0 +1,13 @@
+package dev.gustavoavila.websocketclient.exceptions;
+
+/**
+ * Used to indicate a protocol problem with a received frame
+ *
+ * @author Gustavo Avila
+ *
+ */
+public class InvalidReceivedFrameException extends RuntimeException {
+    public InvalidReceivedFrameException(String message) {
+        super(message);
+    }
+}