[BugFix] Fix read listener not support ssl (backport #65291) (#65357)

Signed-off-by: stdpain <[email protected]>
Co-authored-by: stdpain <[email protected]>

Author: mergify[bot]

fe/fe-core/src/main/java/com/starrocks/mysql/MysqlChannel.java

@@ -37,6 +37,7 @@
 import com.starrocks.common.util.NetUtils;
 import com.starrocks.mysql.nio.MySQLReadListener;
 import com.starrocks.mysql.ssl.SSLChannel;
+import com.starrocks.mysql.ssl.SSLDecoder;
 import com.starrocks.qe.ConnectContext;
 import com.starrocks.qe.ConnectProcessor;
 import org.apache.logging.log4j.LogManager;
@@ -143,6 +144,13 @@ public void setSSLChannel(SSLChannel sslChannel) {
         this.sslChannel = sslChannel;
     }
 
+    public SSLDecoder getSSLDecoder() {
+        if (this.sslChannel == null) {
+            return null;
+        }
+        return this.sslChannel.createDecoder();
+    }
+
     protected int readAll(ByteBuffer dstBuf) throws IOException {
         if (sslChannel != null) {
             return sslChannel.readAll(dstBuf);

fe/fe-core/src/main/java/com/starrocks/mysql/nio/MySQLReadListener.java

@@ -17,6 +17,7 @@
 import com.starrocks.common.Config;
 import com.starrocks.mysql.MysqlPackageDecoder;
 import com.starrocks.mysql.RequestPackage;
+import com.starrocks.mysql.ssl.SSLDecoder;
 import com.starrocks.qe.ConnectContext;
 import com.starrocks.qe.ConnectProcessor;
 import com.starrocks.rpc.RpcException;
@@ -31,14 +32,16 @@ public class MySQLReadListener implements ChannelListener<ConduitStreamSourceCha
     private static final Logger LOG = LogManager.getLogger(MySQLReadListener.class);
     private final ConnectContext ctx;
     private final ConnectProcessor connectProcessor;
-    private final MysqlPackageDecoder decoder = new MysqlPackageDecoder();
+    private final MysqlPackageDecoder packageDecoder = new MysqlPackageDecoder();
 
     protected static final int DEFAULT_BUFFER_SIZE = 16 * 1024;
     private final ByteBuffer readBuffer = ByteBuffer.allocate(DEFAULT_BUFFER_SIZE);
+    private final SSLDecoder sslDecoder;
 
     public MySQLReadListener(ConnectContext connectContext, ConnectProcessor connectProcessor) {
         this.ctx = connectContext;
         this.connectProcessor = connectProcessor;
+        this.sslDecoder = this.ctx.getMysqlChannel().getSSLDecoder();
     }
 
     @Override
@@ -66,11 +69,18 @@ public void handleEvent(ConduitStreamSourceChannel channel) {
                 }
 
                 readBuffer.flip();
-                decoder.consume(readBuffer);
+
+                if (sslDecoder != null) {
+                    sslDecoder.feed(readBuffer);
+                    packageDecoder.consume(sslDecoder.decode());
+                } else {
+                    packageDecoder.consume(readBuffer);
+                }
+
                 readBuffer.compact();
 
                 RequestPackage pkg;
-                while ((pkg = decoder.poll()) != null) {
+                while ((pkg = packageDecoder.poll()) != null) {
                     final RequestPackage req = pkg;
                     channel.getWorker().execute(() -> {
                         handleRequest(req);

fe/fe-core/src/main/java/com/starrocks/mysql/ssl/SSLChannel.java

@@ -24,4 +24,6 @@ public interface SSLChannel {
     int readAll(ByteBuffer buffer) throws IOException;
 
     void write(ByteBuffer buffer) throws IOException;
+
+    SSLDecoder createDecoder();
 }

fe/fe-core/src/main/java/com/starrocks/mysql/ssl/SSLChannelImp.java

@@ -433,4 +433,8 @@ protected ByteBuffer handleBufferOverflow(SSLEngine engine, ByteBuffer buffer) {
         replaceBuffer.put(buffer);
         return replaceBuffer;
     }
+
+    public SSLDecoder createDecoder() {
+        return new SSLDecoder(sslEngine, peerNetData, peerAppData);
+    }
 }

fe/fe-core/src/main/java/com/starrocks/mysql/ssl/SSLDecoder.java

@@ -0,0 +1,103 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.starrocks.mysql.ssl;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import javax.net.ssl.SSLEngine;
+import javax.net.ssl.SSLEngineResult;
+
+public class SSLDecoder {
+    private static final Logger LOG = LogManager.getLogger(SSLDecoder.class);
+
+    private final SSLEngine sslEngine;
+
+    // Encrypted data buffer (network)
+    private ByteBuffer netBuffer;
+    // Decrypted plaintext buffer
+    private ByteBuffer appBuffer;
+
+    public SSLDecoder(SSLEngine sslEngine, ByteBuffer netBuffer, ByteBuffer appBuffer) {
+        this.sslEngine = sslEngine;
+        this.netBuffer = netBuffer;
+        this.appBuffer = appBuffer;
+        this.netBuffer.compact();
+    }
+
+    public void feed(ByteBuffer encryptedInput) throws IOException {
+        // Check if netBuffer has enough remaining capacity
+        if (encryptedInput.remaining() > netBuffer.remaining()) {
+            // Expand netBuffer if needed
+            ByteBuffer newBuf = ByteBuffer.allocate(
+                    Math.max(netBuffer.capacity() * 2, netBuffer.position() + encryptedInput.remaining()));
+            netBuffer.flip();
+            newBuf.put(netBuffer);
+            netBuffer = newBuf;
+        }
+        netBuffer.put(encryptedInput);
+    }
+
+    public ByteBuffer decode() throws IOException {
+        if (!appBuffer.hasRemaining()) {
+            appBuffer.clear();
+        }
+        netBuffer.flip();
+
+        while (netBuffer.hasRemaining()) {
+            SSLEngineResult result = sslEngine.unwrap(netBuffer, appBuffer);
+            switch (result.getStatus()) {
+                case OK:
+                    // No progress (no bytes consumed or produced), break the loop to avoid busy-wait
+                    if (result.bytesConsumed() == 0 && result.bytesProduced() == 0) {
+                        netBuffer.compact();
+                        appBuffer.flip();
+                        return appBuffer;
+                    }
+                    break;
+                case BUFFER_OVERFLOW:
+                    // appBuffer is too small, expand it and retry
+                    expandAppBuffer();
+                    continue;
+                case BUFFER_UNDERFLOW:
+                    // Not enough encrypted data, preserve remaining data in netBuffer
+                    // and exit. More data will be fed in the next read event.
+                    netBuffer.compact();
+                    appBuffer.flip();
+                    return appBuffer;
+                case CLOSED:
+                    throw new IOException("SSL engine closed");
+                default:
+                    throw new IllegalStateException("Unexpected SSL status: " + result.getStatus());
+            }
+        }
+
+        // All encrypted data consumed, flip appBuffer for reading
+        netBuffer.compact();
+        appBuffer.flip();
+        return appBuffer;
+    }
+
+    private void expandAppBuffer() {
+        int newCapacity = appBuffer.capacity() * 2;
+        ByteBuffer newBuf = ByteBuffer.allocateDirect(newCapacity);
+        appBuffer.flip();
+        newBuf.put(appBuffer);
+        appBuffer = newBuf;
+    }
+
+}