From 89eaf410dd8fe3845fa6fb0de3469b55da0205cb Mon Sep 17 00:00:00 2001
From: doyledavidson <ddavidson@automatedlogic.com>
Date: Mon, 9 Sep 2019 12:11:47 -0400
Subject: [PATCH] suggested fix for issue 665 "data read at end of SSL
 handshake is discarded"

---
 .../org/java_websocket/SSLSocketChannel2.java | 32 ++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/main/java/org/java_websocket/SSLSocketChannel2.java b/src/main/java/org/java_websocket/SSLSocketChannel2.java
index 75a56c68..896faa40 100644
--- a/src/main/java/org/java_websocket/SSLSocketChannel2.java
+++ b/src/main/java/org/java_websocket/SSLSocketChannel2.java
@@ -215,6 +215,7 @@ protected void consumeDelegatedTasks() {
     }
 
     protected void createBuffers( SSLSession session ) {
+        saveCryptedData(); // save any remaining data in inCrypt
         int netBufferMax = session.getPacketBufferSize();
         int appBufferMax = Math.max(session.getApplicationBufferSize(), netBufferMax);
 
@@ -269,6 +270,7 @@ public int write( ByteBuffer src ) throws IOException {
      * @return the number of bytes read.
      **/
     public int read(ByteBuffer dst) throws IOException {
+        tryRestoreCryptedData();
         while (true) {
             if (!dst.hasRemaining())
                 return 0;
@@ -329,6 +331,7 @@ private int readRemaining( ByteBuffer dst ) throws SSLException {
         }
         if( !inData.hasRemaining() )
             inData.clear();
+        tryRestoreCryptedData();
         // test if some bytes left from last read (e.g. BUFFER_UNDERFLOW)
         if( inCrypt.hasRemaining() ) {
             unwrap();
@@ -396,7 +399,7 @@ public void writeMore() throws IOException {
 
     @Override
     public boolean isNeedRead() {
-        return inData.hasRemaining() || ( inCrypt.hasRemaining() && readEngineResult.getStatus() != Status.BUFFER_UNDERFLOW && readEngineResult.getStatus() != Status.CLOSED );
+        return saveCryptData != null || inData.hasRemaining() || ( inCrypt.hasRemaining() && readEngineResult.getStatus() != Status.BUFFER_UNDERFLOW && readEngineResult.getStatus() != Status.CLOSED );
     }
 
     @Override
@@ -430,4 +433,31 @@ public boolean isBlocking() {
     public SSLEngine getSSLEngine() {
         return sslEngine;
     }
+
+
+    // to avoid complexities with inCrypt, extra unwrapped data after SSL handshake will be saved off in a byte array
+    // and the inserted back on first read
+    private byte[] saveCryptData = null;
+    private void saveCryptedData()
+    {
+        // did we find any extra data?
+        if (inCrypt != null && inCrypt.remaining() > 0)
+        {
+            int saveCryptSize = inCrypt.remaining();
+            saveCryptData = new byte[saveCryptSize];
+            inCrypt.get(saveCryptData);
+        }
+    }
+
+    private void tryRestoreCryptedData()
+    {
+        // was there any extra data, then put into inCrypt and clean up
+        if ( saveCryptData != null )
+        {
+            inCrypt.clear();
+            inCrypt.put( saveCryptData );
+            inCrypt.flip();
+            saveCryptData = null;
+        }
+    }
 }
\ No newline at end of file