improved buffer allocation, renamed variables, documented code and fixed potential threading issue in SSLSocketChannel2#processHandshake (TooTallNate#154)

Davidiusdadi · Davidiusdadi · commit d7b546762170 · 2013-02-20T02:15:57.000+01:00
diff --git a/src/main/java/org/java_websocket/SSLSocketChannel2.java b/src/main/java/org/java_websocket/SSLSocketChannel2.java
@@ -44,33 +44,34 @@ public class SSLSocketChannel2 implements ByteChannel, WrappedByteChannel {
 	/** encrypted data incoming */
 	protected ByteBuffer inCrypt;
 
-	protected SocketChannel sc;
-	protected SelectionKey key;
+	/** the underlying channel */
+	protected SocketChannel socketChannel;
+	/** used to set interestOP SelectionKey.OP_WRITE for the underlying channel */
+	protected SelectionKey selectionKey;
 
-	protected SSLEngineResult res;
+
+	protected SSLEngineResult engineResult;
 	protected SSLEngine sslEngine;
-	protected final boolean isblocking;
 
-	private Status status = Status.BUFFER_UNDERFLOW;
+
+	private Status engineStatus = Status.BUFFER_UNDERFLOW;
 
 	public SSLSocketChannel2( SocketChannel channel , SSLEngine sslEngine , ExecutorService exec , SelectionKey key ) throws IOException {
-		this.sc = channel;
+		if( channel == null || sslEngine == null || exec == null )
+			throw new IllegalArgumentException( "parameter must not be null" );
 
+		this.socketChannel = channel;
 		this.sslEngine = sslEngine;
 		this.exec = exec;
 
 		tasks = new ArrayList<Future<?>>( 3 );
 		if( key != null ) {
 			key.interestOps( key.interestOps() | SelectionKey.OP_WRITE );
-			this.key = key;
+			this.selectionKey = key;
 		}
-		isblocking = channel.isBlocking();
-
-		sslEngine.setEnableSessionCreation( true );
-		SSLSession session = sslEngine.getSession();
-		createBuffers( session );
-
-		sc.write( wrap( emptybuffer ) );// initializes res
+		createBuffers( sslEngine.getSession() );
+		// kick off handshake
+		socketChannel.write( wrap( emptybuffer ) );// initializes res
 		processHandshake();
 	}
 
@@ -93,6 +94,8 @@ private void consumeFutureUninterruptible( Future<?> f ) {
 	}
 
 	private synchronized void processHandshake() throws IOException {
+		if( engineResult.getHandshakeStatus() == HandshakeStatus.NOT_HANDSHAKING )
+			return; // since this may be called either from a reading or a writing thread and because this method is synchronized it is necessary to double check if we are still handshaking.
 		if( !tasks.isEmpty() ) {
 			Iterator<Future<?>> it = tasks.iterator();
 			while ( it.hasNext() ) {
@@ -107,28 +110,35 @@ private synchronized void processHandshake() throws IOException {
 			}
 		}
 
-		if( res.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NEED_UNWRAP ) {
-			if( !isblocking || status == Status.BUFFER_UNDERFLOW ) {
+		if( engineResult.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NEED_UNWRAP ) {
+			if( !isBlocking() || engineStatus == Status.BUFFER_UNDERFLOW ) {
 				inCrypt.compact();
-				int read = sc.read( inCrypt );
+				int read = socketChannel.read( inCrypt );
 				if( read == -1 ) {
 					throw new IOException( "connection closed unexpectedly by peer" );
 				}
 				inCrypt.flip();
 			}
 			inData.compact();
 			unwrap();
+			if( engineResult.getHandshakeStatus() == HandshakeStatus.FINISHED ) {
+				createBuffers( sslEngine.getSession() );
+				return;
+			}
 		}
 		consumeDelegatedTasks();
-		if( tasks.isEmpty() || res.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NEED_WRAP ) {
-			sc.write( wrap( emptybuffer ) );
+		assert ( engineResult.getHandshakeStatus() != HandshakeStatus.NOT_HANDSHAKING );
+		if( tasks.isEmpty() || engineResult.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NEED_WRAP ) {
+			socketChannel.write( wrap( emptybuffer ) );
+			if( engineResult.getHandshakeStatus() == HandshakeStatus.FINISHED ) {
+				createBuffers( sslEngine.getSession() );
+			}
 		}
-
 	}
 
 	private synchronized ByteBuffer wrap( ByteBuffer b ) throws SSLException {
 		outCrypt.compact();
-		res = sslEngine.wrap( b, outCrypt );
+		engineResult = sslEngine.wrap( b, outCrypt );
 		outCrypt.flip();
 		return outCrypt;
 	}
@@ -137,9 +147,9 @@ private synchronized ByteBuffer unwrap() throws SSLException {
 		int rem;
 		do {
 			rem = inData.remaining();
-			res = sslEngine.unwrap( inCrypt, inData );
-			status = res.getStatus();
-		} while ( status == SSLEngineResult.Status.OK && ( rem != inData.remaining() || res.getHandshakeStatus() == HandshakeStatus.NEED_UNWRAP ) );
+			engineResult = sslEngine.unwrap( inCrypt, inData );
+			engineStatus = engineResult.getStatus();
+		} while ( engineStatus == SSLEngineResult.Status.OK && ( rem != inData.remaining() || engineResult.getHandshakeStatus() == HandshakeStatus.NEED_UNWRAP ) );
 
 		inData.flip();
 		return inData;
@@ -157,11 +167,23 @@ protected void createBuffers( SSLSession session ) {
 		int appBufferMax = session.getApplicationBufferSize();
 		int netBufferMax = session.getPacketBufferSize();
 
-		inData = ByteBuffer.allocate( appBufferMax );
-		outCrypt = ByteBuffer.allocate( netBufferMax );
-		inCrypt = ByteBuffer.allocate( netBufferMax );
+		if( inData == null ) {
+			inData = ByteBuffer.allocate( appBufferMax );
+			outCrypt = ByteBuffer.allocate( netBufferMax );
+			inCrypt = ByteBuffer.allocate( netBufferMax );
+		} else {
+			if( inData.capacity() != appBufferMax )
+				inData = ByteBuffer.allocate( appBufferMax );
+			if( outCrypt.capacity() != netBufferMax )
+				outCrypt = ByteBuffer.allocate( netBufferMax );
+			if( inCrypt.capacity() != netBufferMax )
+				inCrypt = ByteBuffer.allocate( netBufferMax );
+		}
+		inData.rewind();
 		inData.flip();
+		inCrypt.rewind();
 		inCrypt.flip();
+		outCrypt.rewind();
 		outCrypt.flip();
 	}
 
@@ -170,7 +192,7 @@ public int write( ByteBuffer src ) throws IOException {
 			processHandshake();
 			return 0;
 		}
-		int num = sc.write( wrap( src ) );
+		int num = socketChannel.write( wrap( src ) );
 		return num;
 
 	}
@@ -203,8 +225,8 @@ public int read( ByteBuffer dst ) throws IOException {
 		else
 			inCrypt.compact();
 
-		if( ( isblocking && inCrypt.position() == 0 ) || status == Status.BUFFER_UNDERFLOW )
-			if( sc.read( inCrypt ) == -1 ) {
+		if( ( isBlocking() && inCrypt.position() == 0 ) || engineStatus == Status.BUFFER_UNDERFLOW )
+			if( socketChannel.read( inCrypt ) == -1 ) {
 				return -1;
 			}
 		inCrypt.flip();
@@ -230,36 +252,36 @@ private int readRemaining( ByteBuffer dst ) throws SSLException {
 	}
 
 	public boolean isConnected() {
-		return sc.isConnected();
+		return socketChannel.isConnected();
 	}
 
 	public void close() throws IOException {
 		sslEngine.closeOutbound();
 		sslEngine.getSession().invalidate();
-		if( sc.isOpen() )
-			sc.write( wrap( emptybuffer ) );// FIXME what if not all bytes can be written
-		sc.close();
+		if( socketChannel.isOpen() )
+			socketChannel.write( wrap( emptybuffer ) );// FIXME what if not all bytes can be written
+		socketChannel.close();
 	}
 
 	private boolean isHandShakeComplete() {
-		HandshakeStatus status = res.getHandshakeStatus();
+		HandshakeStatus status = engineResult.getHandshakeStatus();
 		return status == SSLEngineResult.HandshakeStatus.FINISHED || status == SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING;
 	}
 
 	public SelectableChannel configureBlocking( boolean b ) throws IOException {
-		return sc.configureBlocking( b );
+		return socketChannel.configureBlocking( b );
 	}
 
 	public boolean connect( SocketAddress remote ) throws IOException {
-		return sc.connect( remote );
+		return socketChannel.connect( remote );
 	}
 
 	public boolean finishConnect() throws IOException {
-		return sc.finishConnect();
+		return socketChannel.finishConnect();
 	}
 
 	public Socket socket() {
-		return sc.socket();
+		return socketChannel.socket();
 	}
 
 	public boolean isInboundDone() {
@@ -268,7 +290,7 @@ public boolean isInboundDone() {
 
 	@Override
 	public boolean isOpen() {
-		return sc.isOpen();
+		return socketChannel.isOpen();
 	}
 
 	@Override
@@ -283,7 +305,7 @@ public void writeMore() throws IOException {
 
 	@Override
 	public boolean isNeedRead() {
-		return inData.hasRemaining() || ( inCrypt.hasRemaining() && res.getStatus() != Status.BUFFER_UNDERFLOW );
+		return inData.hasRemaining() || ( inCrypt.hasRemaining() && engineResult.getStatus() != Status.BUFFER_UNDERFLOW );
 	}
 
 	@Override
@@ -310,7 +332,7 @@ private int transfereTo( ByteBuffer from, ByteBuffer to ) {
 
 	@Override
 	public boolean isBlocking() {
-		return isblocking;
+		return socketChannel.isBlocking();
 	}
 
 }
diff --git a/src/main/java/org/java_websocket/SocketChannelIOHelper.java b/src/main/java/org/java_websocket/SocketChannelIOHelper.java
@@ -19,6 +19,10 @@ public static boolean read( final ByteBuffer buf, WebSocketImpl ws, ByteChannel
 		return read != 0;
 	}
 
+	/**
+	 * @see WrappedByteChannel#readMore(ByteBuffer)
+	 * @return returns whether there is more data left which can be obtained via {@link #readMore(ByteBuffer, WebSocketImpl, WrappedByteChannel)}
+	 **/
 	public static boolean readMore( final ByteBuffer buf, WebSocketImpl ws, WrappedByteChannel channel ) throws IOException {
 		buf.clear();
 		int read = channel.readMore( buf );
diff --git a/src/main/java/org/java_websocket/WrappedByteChannel.java b/src/main/java/org/java_websocket/WrappedByteChannel.java
@@ -10,7 +10,17 @@ public interface WrappedByteChannel extends ByteChannel {
 	public boolean isNeedWrite();
 	public void writeMore() throws IOException;
 
+	/**
+	 * returns whether readMore should be called to fetch data which has been decoded but not yet been returned.
+	 * 
+	 * @see #read(ByteBuffer)
+	 * @see #readMore(ByteBuffer)
+	 **/
 	public boolean isNeedRead();
+	/**
+	 * This function does not read data from the underlying channel at all. It is just a way to fetch data which has already be received or decoded but was but was not yet returned to the user.
+	 * This could be the case when the decoded data did not fit into the buffer the user passed to {@link #read(ByteBuffer)}.
+	 **/
 	public int readMore( ByteBuffer dst ) throws SSLException;
 	public boolean isBlocking();
 }
