From cb60265ec192a5c152d0ce6d4e4ce65633945073 Mon Sep 17 00:00:00 2001
From: Riza Suminto <riza.suminto@cloudera.com>
Date: Wed, 5 Apr 2023 13:31:44 -0700
Subject: THRIFT-5696: Allow custom TConfiguration for TByteBuffer.java

---
 .../org/apache/thrift/transport/TByteBuffer.java   | 18 ++++++++++++---
 .../apache/thrift/transport/TestTByteBuffer.java   | 27 ++++++++++++++++++++++
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/lib/java/src/main/java/org/apache/thrift/transport/TByteBuffer.java b/lib/java/src/main/java/org/apache/thrift/transport/TByteBuffer.java
index fa296e7c4..226f02a04 100644
--- a/lib/java/src/main/java/org/apache/thrift/transport/TByteBuffer.java
+++ b/lib/java/src/main/java/org/apache/thrift/transport/TByteBuffer.java
@@ -10,17 +10,29 @@ public final class TByteBuffer extends TEndpointTransport {
   private final ByteBuffer byteBuffer;
 
   /**
-   * Creates a new TByteBuffer wrapping a given NIO ByteBuffer.
+   * Creates a new TByteBuffer wrapping a given NIO ByteBuffer and custom TConfiguration.
    *
+   * @param configuration the custom TConfiguration.
    * @param byteBuffer the NIO ByteBuffer to wrap.
    * @throws TTransportException on error.
    */
-  public TByteBuffer(ByteBuffer byteBuffer) throws TTransportException {
-    super(new TConfiguration());
+  public TByteBuffer(TConfiguration configuration, ByteBuffer byteBuffer)
+      throws TTransportException {
+    super(configuration);
     this.byteBuffer = byteBuffer;
     updateKnownMessageSize(byteBuffer.capacity());
   }
 
+  /**
+   * Creates a new TByteBuffer wrapping a given NIO ByteBuffer.
+   *
+   * @param byteBuffer the NIO ByteBuffer to wrap.
+   * @throws TTransportException on error.
+   */
+  public TByteBuffer(ByteBuffer byteBuffer) throws TTransportException {
+    this(new TConfiguration(), byteBuffer);
+  }
+
   @Override
   public boolean isOpen() {
     return true;
diff --git a/lib/java/src/test/java/org/apache/thrift/transport/TestTByteBuffer.java b/lib/java/src/test/java/org/apache/thrift/transport/TestTByteBuffer.java
index 748de121e..26ffc5d58 100644
--- a/lib/java/src/test/java/org/apache/thrift/transport/TestTByteBuffer.java
+++ b/lib/java/src/test/java/org/apache/thrift/transport/TestTByteBuffer.java
@@ -1,10 +1,12 @@
 package org.apache.thrift.transport;
 
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
+import org.apache.thrift.TConfiguration;
 import org.junit.jupiter.api.Test;
 
 public class TestTByteBuffer {
@@ -39,4 +41,29 @@ public class TestTByteBuffer {
             () -> byteBuffer.write("Hello World".getBytes(StandardCharsets.UTF_8)));
     assertEquals("Not enough room in output buffer", e.getMessage());
   }
+
+  @Test
+  public void testSmallTConfiguration() throws Exception {
+    // Test that TByteBuffer init fail with small max message size.
+    final TConfiguration configSmall =
+        new TConfiguration(
+            4, TConfiguration.DEFAULT_MAX_FRAME_SIZE, TConfiguration.DEFAULT_RECURSION_DEPTH);
+    TTransportException e =
+        assertThrows(
+            TTransportException.class,
+            () -> new TByteBuffer(configSmall, ByteBuffer.allocate(100)));
+    assertEquals("MaxMessageSize reached", e.getMessage());
+  }
+
+  @Test
+  public void testLargeTConfiguration() throws Exception {
+    // Test that TByteBuffer init pass with large max message size beyond
+    // TConfiguration.DEFAULT_MAX_MESSAGE_SIZE.
+    int maxSize = 101 * 1024 * 1024;
+    int bufferSize = (100 * 1024 + 512) * 1024;
+    final TConfiguration configLarge =
+        new TConfiguration(
+            maxSize, TConfiguration.DEFAULT_MAX_FRAME_SIZE, TConfiguration.DEFAULT_RECURSION_DEPTH);
+    assertDoesNotThrow(() -> new TByteBuffer(configLarge, ByteBuffer.allocate(bufferSize)));
+  }
 }
-- 
cgit v1.2.1