support underflow handling without thread sleep

(cherry picked from commit ba69698)
rabbitmq · Jul 30, 2021 · 4ba627c · 4ba627c
1 parent 02ccdc2
commit 4ba627c
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 14 deletions.
diff --git a/src/main/java/com/rabbitmq/client/impl/nio/SocketChannelFrameHandlerState.java b/src/main/java/com/rabbitmq/client/impl/nio/SocketChannelFrameHandlerState.java
@@ -26,6 +26,8 @@
 import java.nio.ByteBuffer;
 import java.nio.channels.SelectionKey;
 import java.nio.channels.SocketChannel;
+import java.util.concurrent.atomic.AtomicBoolean;
+
 
 /**
  *
@@ -70,6 +72,8 @@ public class SocketChannelFrameHandlerState {
 
     final FrameBuilder frameBuilder;
 
+    private final AtomicBoolean isUnderflowHandlingEnabled = new AtomicBoolean(false);
+
     public SocketChannelFrameHandlerState(SocketChannel channel, NioLoopContext nioLoopsState, NioParams nioParams, SSLEngine sslEngine) {
         this.channel = channel;
         this.readSelectorState = nioLoopsState.readSelectorState;
@@ -105,7 +109,7 @@ public SocketChannelFrameHandlerState(SocketChannel channel, NioLoopContext nioL
             this.outputStream = new DataOutputStream(
                 new SslEngineByteBufferOutputStream(sslEngine, plainOut, cipherOut, channel)
             );
-            this.frameBuilder = new SslEngineFrameBuilder(sslEngine, plainIn, cipherIn, channel);
+            this.frameBuilder = new SslEngineFrameBuilder(sslEngine, plainIn, cipherIn, channel, isUnderflowHandlingEnabled);
         }
 
     }
@@ -176,11 +180,14 @@ void endWriteSequence() {
 
     void prepareForReadSequence() throws IOException {
         if(ssl) {
-            cipherIn.clear();
-            plainIn.clear();
+            if (!isUnderflowHandlingEnabled.get()) {
+                cipherIn.clear();
+                cipherIn.flip();
+            }
 
-            cipherIn.flip();
+            plainIn.clear();
             plainIn.flip();
+
         } else {
             NioHelper.read(channel, plainIn);
             plainIn.flip();
@@ -189,6 +196,15 @@ void prepareForReadSequence() throws IOException {
 
     boolean continueReading() throws IOException {
         if(ssl) {
+            if (isUnderflowHandlingEnabled.get()) {
+                int bytesRead = NioHelper.read(channel, cipherIn);
+                if (bytesRead == 0) {
+                    return false;
+                } else {
+                    cipherIn.flip();
+                    return true;
+                }
+            }
             if (!plainIn.hasRemaining() && !cipherIn.hasRemaining()) {
                 // need to try to read something
                 cipherIn.clear();

diff --git a/src/main/java/com/rabbitmq/client/impl/nio/SslEngineFrameBuilder.java b/src/main/java/com/rabbitmq/client/impl/nio/SslEngineFrameBuilder.java
@@ -21,6 +21,8 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.channels.ReadableByteChannel;
+import java.util.concurrent.atomic.AtomicBoolean;
+
 
 /**
  * Sub-class of {@link FrameBuilder} that unwraps crypted data from the network.
@@ -32,20 +34,25 @@ public class SslEngineFrameBuilder extends FrameBuilder {
 
     private final ByteBuffer cipherBuffer;
 
-    public SslEngineFrameBuilder(SSLEngine sslEngine, ByteBuffer plainIn, ByteBuffer cipherIn, ReadableByteChannel channel) {
+    private final AtomicBoolean isUnderflowHandlingEnabled;
+
+    public SslEngineFrameBuilder(SSLEngine sslEngine, ByteBuffer plainIn, ByteBuffer cipherIn, ReadableByteChannel channel, final AtomicBoolean isUnderflowHandlingEnabled) {
         super(channel, plainIn);
         this.sslEngine = sslEngine;
         this.cipherBuffer = cipherIn;
+        this.isUnderflowHandlingEnabled = isUnderflowHandlingEnabled;
     }
 
     @Override
     protected boolean somethingToRead() throws IOException {
-        if (applicationBuffer.hasRemaining()) {
+        if (applicationBuffer.hasRemaining() && !isUnderflowHandlingEnabled.get()) {
             return true;
         } else {
             applicationBuffer.clear();
 
-            while (true) {
+            boolean underflowHandling = false;
+
+            try {
                 SSLEngineResult result = sslEngine.unwrap(cipherBuffer, applicationBuffer);
                 switch (result.getStatus()) {
                     case OK:
@@ -59,18 +66,18 @@ protected boolean somethingToRead() throws IOException {
                         throw new SSLException("buffer overflow in read");
                     case BUFFER_UNDERFLOW:
                         cipherBuffer.compact();
-                        int read = NioHelper.read(channel, cipherBuffer);
-                        if (read == 0) {
-                            return false;
-                        }
-                        cipherBuffer.flip();
-                        break;
+                        underflowHandling = true;
+                        return false;
                     case CLOSED:
                         throw new SSLException("closed in read");
                     default:
                         throw new IllegalStateException("Invalid SSL status: " + result.getStatus());
-                    }
+                }
+            } finally {
+                isUnderflowHandlingEnabled.set(underflowHandling);
             }
+
+            return false;
         }
     }