Fix SslEngineByteBufferInputStream to be prepared that read() may still result in BUFFER_UNDERFLOW

Author: dimas

src/main/java/com/rabbitmq/client/impl/nio/NioHelper.java

@@ -47,4 +47,11 @@ static int retryRead(ReadableByteChannel channel, ByteBuffer buffer) throws IOEx
         return read;
     }
 
+    static int readWithRetry(ReadableByteChannel channel, ByteBuffer buffer) throws IOException {
+        int bytesRead = NioHelper.read(channel, buffer);
+        if (bytesRead <= 0) {
+            bytesRead = NioHelper.retryRead(channel, buffer);
+        }
+        return bytesRead;
+    }
 }

src/main/java/com/rabbitmq/client/impl/nio/SslEngineByteBufferInputStream.java

@@ -43,57 +43,40 @@ public SslEngineByteBufferInputStream(SSLEngine sslEngine, ByteBuffer plainIn, B
 
     @Override
     public int read() throws IOException {
+
         if (plainIn.hasRemaining()) {
             return readFromBuffer(plainIn);
-        } else {
-            plainIn.clear();
+        }
+
+        plainIn.clear();
+
+        while (true) {
+
             SSLEngineResult result = sslEngine.unwrap(cipherIn, plainIn);
 
             switch (result.getStatus()) {
-            case OK:
-                plainIn.flip();
-                if (plainIn.hasRemaining()) {
-                    return readFromBuffer(plainIn);
-                }
-                break;
-            case BUFFER_OVERFLOW:
-                throw new SSLException("buffer overflow in read");
-            case BUFFER_UNDERFLOW:
-                if (cipherIn.hasRemaining()) {
+                case OK:
+                    plainIn.flip();
+                    if (plainIn.hasRemaining()) {
+                        return readFromBuffer(plainIn);
+                    }
+                    break;
+                case BUFFER_OVERFLOW:
+                    throw new SSLException("buffer overflow in read");
+                case BUFFER_UNDERFLOW:
                     cipherIn.compact();
-                } else {
-                    cipherIn.clear();
-                }
-
-                int bytesRead = NioHelper.read(channel, cipherIn);
-                // see https://github.com/rabbitmq/rabbitmq-java-client/issues/307
-                if (bytesRead <= 0) {
-                    bytesRead = NioHelper.retryRead(channel, cipherIn);
-                    if(bytesRead <= 0) {
+                    int bytesRead = NioHelper.readWithRetry(channel, cipherIn);
+                    if (bytesRead <= 0) {
                         throw new IllegalStateException("Should be reading something from the network");
                     }
-                }
-                cipherIn.flip();
-
-                plainIn.clear();
-                result = sslEngine.unwrap(cipherIn, plainIn);
-
-                if (result.getStatus() != SSLEngineResult.Status.OK) {
-                    throw new SSLException("Unexpected result: " + result);
-                }
-                plainIn.flip();
-                if (plainIn.hasRemaining()) {
-                    return readFromBuffer(plainIn);
-                }
-                break;
-            case CLOSED:
-                throw new SSLException("closed in read");
-            default:
-                throw new IllegalStateException("Invalid SSL status: " + result.getStatus());
+                    cipherIn.flip();
+                    break;
+                case CLOSED:
+                    throw new SSLException("closed in read");
+                default:
+                    throw new IllegalStateException("Invalid SSL status: " + result.getStatus());
             }
         }
-
-        return -1;
     }
 
     private int readFromBuffer(ByteBuffer buffer) {

src/test/java/com/rabbitmq/client/test/ssl/NioTlsUnverifiedConnection.java

@@ -96,15 +96,28 @@ public void configure(SSLEngine sslEngine) throws IOException {
     }
 
     @Test public void messageSize() throws Exception {
-        int [] sizes = new int [] {100, 1000, 10 * 1000, 1 * 1000 * 1000, 5 * 1000 * 1000};
-        for(int size : sizes) {
-            CountDownLatch latch = new CountDownLatch(1);
-            connection = basicGetBasicConsume(connection, "tls.nio.queue", latch, size);
-            boolean messagesReceived = latch.await(5, TimeUnit.SECONDS);
-            assertTrue("Message has not been received", messagesReceived);
+        int[] sizes = new int[]{100, 1000, 10 * 1000, 1 * 1000 * 1000, 5 * 1000 * 1000};
+        for (int size : sizes) {
+            sendAndVerifyMessage(size);
         }
     }
 
+    // The purpose of this test is to put some stress on client TLS layer (SslEngineByteBufferInputStream to be specific)
+    // in an attempt to trigger condition described in https://github.com/rabbitmq/rabbitmq-java-client/issues/317
+    // Unfortunately it is not guaranteed to be reproducible
+    @Test public void largeMessagesTlsTraffic() throws Exception {
+        for (int i = 0; i < 50; i++) {
+            sendAndVerifyMessage(76390);
+        }
+    }
+
+    private void sendAndVerifyMessage(int size) throws Exception {
+        CountDownLatch latch = new CountDownLatch(1);
+        connection = basicGetBasicConsume(connection, "tls.nio.queue", latch, size);
+        boolean messagesReceived = latch.await(5, TimeUnit.SECONDS);
+        assertTrue("Message has not been received", messagesReceived);
+    }
+
     private Connection basicGetBasicConsume(Connection connection, String queue, final CountDownLatch latch, int msgSize)
         throws IOException, TimeoutException {
         Channel channel = connection.createChannel();