commit cb60265ec192a5c152d0ce6d4e4ce65633945073
author Riza Suminto <riza.suminto@cloudera.com> Wed Apr 05 13:31:44 2023 -0700
committer Yuxuan 'fishy' Wang <fishywang@gmail.com> Thu Apr 06 11:46:34 2023 -0700
tree ad8cdb140f3dc2f94a161e902ea2894ca9aa9ae8
parent 12ab079ecfe0c87e984fc7498a43a614d8e8a240

THRIFT-5696: Allow custom TConfiguration for TByteBuffer.java

lib/java/src/main/java/org/apache/thrift/transport/TByteBuffer.java

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 fa296e7..226f02a 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,15 +10,27 @@
   private final ByteBuffer 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(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 {
-    super(new TConfiguration());
-    this.byteBuffer = byteBuffer;
-    updateKnownMessageSize(byteBuffer.capacity());
+    this(new TConfiguration(), byteBuffer);
   }
 
   @Override

lib/java/src/test/java/org/apache/thrift/transport/TestTByteBuffer.java

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 748de12..26ffc5d 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 @@
             () -> 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)));
+  }
 }