Merge pull request #943 from da-als/master

Fix for discarded data at end of SSL handshake
TooTallNate · Jan 19, 2020 · 7bfa83d · 7bfa83d
2 parents 65300d4 + bc3669f
commit 7bfa83d
Showing 1 changed file with 31 additions and 1 deletion.
diff --git 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;
+        }
+    }
 }