Wrap IOException and include WebSocket

Fixes TooTallNate#900

src/main/java/org/java_websocket/exceptions/WrappedIOException.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright (c) 2010-2019 Nathan Rajlich
+ *
+ * Permission is hereby granted, free of charge, to any person
+ * obtaining a copy of this software and associated documentation
+ * files (the "Software"), to deal in the Software without
+ * restriction, including without limitation the rights to use,
+ * copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following
+ * conditions:
+ *
+ * The above copyright notice and this permission notice shall be
+ * included in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
+ * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+ * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+ * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+ * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+ * OTHER DEALINGS IN THE SOFTWARE.
+ *
+ */
+
+package org.java_websocket.exceptions;
+
+import org.java_websocket.WebSocket;
+
+import java.io.IOException;
+
+/**
+ * Exception to wrap an IOException and include information about the websocket which had the exception
+ * @since 1.4.1
+ */
+public class WrappedIOException extends Throwable {
+
+    /**
+     * The websocket where the IOException happened
+     */
+    private final WebSocket connection;
+
+    /**
+     * The IOException
+     */
+    private final IOException ioException;
+
+    /**
+     * Wrapp an IOException and include the websocket
+     * @param connection the websocket where the IOException happened
+     * @param ioException the IOException
+     */
+    public WrappedIOException(WebSocket connection, IOException ioException) {
+        this.connection = connection;
+        this.ioException = ioException;
+    }
+
+    /**
+     * The websocket where the IOException happened
+     * @return the websocket for the wrapped IOException
+     */
+    public WebSocket getConnection() {
+        return connection;
+    }
+
+    /**
+     * The wrapped IOException
+     * @return IOException which is wrapped
+     */
+    public IOException getIOException() {
+        return ioException;
+    }
+}

src/main/java/org/java_websocket/server/WebSocketServer.java

@@ -48,6 +48,7 @@
 import org.java_websocket.drafts.Draft;
 import org.java_websocket.exceptions.InvalidDataException;
 import org.java_websocket.exceptions.WebsocketNotConnectedException;
+import org.java_websocket.exceptions.WrappedIOException;
 import org.java_websocket.framing.CloseFrame;
 import org.java_websocket.framing.Framedata;
 import org.java_websocket.handshake.ClientHandshake;
@@ -320,7 +321,6 @@ public void run() {
 			int selectTimeout = 0;
 			while ( !selectorthread.isInterrupted() && iShutdownCount != 0) {
 				SelectionKey key = null;
-				WebSocketImpl conn = null;
 				try {
 					if (isclosed.get()) {
 						selectTimeout = 5;
@@ -334,7 +334,6 @@ public void run() {
 
 					while ( i.hasNext() ) {
 						key = i.next();
-						conn = null;
 
 						if( !key.isValid() ) {
 							continue;
@@ -358,10 +357,10 @@ public void run() {
 					// an other thread may cancel the key
 				} catch ( ClosedByInterruptException e ) {
 					return; // do the same stuff as when InterruptedException is thrown
+				} catch ( WrappedIOException ex) {
+					handleIOException( key, ex.getConnection(), ex.getIOException());
 				} catch ( IOException ex ) {
-					if( key != null )
-						key.cancel();
-					handleIOException( key, conn, ex );
+					handleIOException( key, null, ex );
 				} catch ( InterruptedException e ) {
 					// FIXME controlled shutdown (e.g. take care of buffermanagement)
 					Thread.currentThread().interrupt();
@@ -445,7 +444,7 @@ private void doAccept(SelectionKey key, Iterator<SelectionKey> i) throws IOExcep
 	 * @throws InterruptedException thrown by taking a buffer
 	 * @throws IOException if an error happened during read
 	 */
-	private boolean doRead(SelectionKey key, Iterator<SelectionKey> i) throws InterruptedException, IOException {
+	private boolean doRead(SelectionKey key, Iterator<SelectionKey> i) throws InterruptedException, WrappedIOException {
 		WebSocketImpl conn = (WebSocketImpl) key.attachment();
 		ByteBuffer buf = takeBuffer();
 		if(conn.getChannel() == null){
@@ -471,7 +470,7 @@ private boolean doRead(SelectionKey key, Iterator<SelectionKey> i) throws Interr
 			}
 		} catch ( IOException e ) {
 			pushBuffer( buf );
-			throw e;
+			throw new WrappedIOException(conn, e);
 		}
 		return true;
 	}
@@ -481,12 +480,16 @@ private boolean doRead(SelectionKey key, Iterator<SelectionKey> i) throws Interr
 	 * @param key the selectionkey to write on
 	 * @throws IOException if an error happened during batch
 	 */
-	private void doWrite(SelectionKey key) throws IOException {
+	private void doWrite(SelectionKey key) throws WrappedIOException {
 		WebSocketImpl conn = (WebSocketImpl) key.attachment();
-		if( SocketChannelIOHelper.batch( conn, conn.getChannel() ) ) {
-			if( key.isValid() ) {
-				key.interestOps(SelectionKey.OP_READ);
+		try {
+			if (SocketChannelIOHelper.batch(conn, conn.getChannel())) {
+				if (key.isValid()) {
+					key.interestOps(SelectionKey.OP_READ);
+				}
 			}
+		} catch (IOException e) {
+			throw new WrappedIOException(conn, e);
 		}
 	}
 
@@ -598,6 +601,9 @@ private void pushBuffer( ByteBuffer buf ) throws InterruptedException {
 
 	private void handleIOException( SelectionKey key, WebSocket conn, IOException ex ) {
 		// onWebsocketError( conn, ex );// conn may be null here
+		if (key != null) {
+			key.cancel();
+		}
 		if( conn != null ) {
 			conn.closeConnection( CloseFrame.ABNORMAL_CLOSE, ex.getMessage() );
 		} else if( key != null ) {