From 4ba627ca0c43c3797183947e49bbcf76546bbb28 Mon Sep 17 00:00:00 2001
From: yandryakov <yandryakov@solanteq.com>
Date: Thu, 29 Jul 2021 16:07:02 +0300
Subject: [PATCH] support underflow handling without thread sleep

(cherry picked from commit ba696986acad8aaee7365aa56e90ccef08bce4bc)
---
 .../nio/SocketChannelFrameHandlerState.java   | 24 ++++++++++++++---
 .../impl/nio/SslEngineFrameBuilder.java       | 27 ++++++++++++-------
 2 files changed, 37 insertions(+), 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
index 50f08a59f..87c1d0892 100644
--- 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
index c2f192387..34295fb36 100644
--- 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;
         }
     }