From 0bca4eb84a0cd784472e809ca68ee334d0592d16 Mon Sep 17 00:00:00 2001
From: Tom Lee <github@tomlee.co>
Date: Thu, 19 Nov 2015 03:03:30 -0800
Subject: [PATCH] THRIFT-3432 Add the TByteBuffer transport for Java

---
 .../apache/thrift/transport/TByteBuffer.java  | 87 +++++++++++++++++++
 .../thrift/transport/TestTByteBuffer.java     | 36 ++++++++
 2 files changed, 123 insertions(+)
 create mode 100644 lib/java/src/org/apache/thrift/transport/TByteBuffer.java
 create mode 100644 lib/java/test/org/apache/thrift/transport/TestTByteBuffer.java

diff --git a/lib/java/src/org/apache/thrift/transport/TByteBuffer.java b/lib/java/src/org/apache/thrift/transport/TByteBuffer.java
new file mode 100644
index 00000000000..a09f33da81b
--- /dev/null
+++ b/lib/java/src/org/apache/thrift/transport/TByteBuffer.java
@@ -0,0 +1,87 @@
+package org.apache.thrift.transport;
+
+import java.nio.BufferOverflowException;
+import java.nio.BufferUnderflowException;
+import java.nio.ByteBuffer;
+
+/**
+ * ByteBuffer-backed implementation of TTransport.
+ */
+public final class TByteBuffer extends TTransport {
+  private final ByteBuffer byteBuffer;
+
+  /**
+   * Creates a new TByteBuffer wrapping a given NIO ByteBuffer.
+   */
+  public TByteBuffer(ByteBuffer byteBuffer) {
+    this.byteBuffer = byteBuffer;
+  }
+
+  @Override
+  public boolean isOpen() {
+    return true;
+  }
+
+  @Override
+  public void open() {
+  }
+
+  @Override
+  public void close() {
+  }
+
+  @Override
+  public int read(byte[] buf, int off, int len) throws TTransportException {
+    final int n = Math.min(byteBuffer.remaining(), len);
+    if (n > 0) {
+      try {
+        byteBuffer.get(buf, off, len);
+      } catch (BufferUnderflowException e) {
+        throw new TTransportException("Unexpected end of input buffer", e);
+      }
+    }
+    return n;
+  }
+
+  @Override
+  public void write(byte[] buf, int off, int len) throws TTransportException {
+    try {
+      byteBuffer.put(buf, off, len);
+    } catch (BufferOverflowException e) {
+      throw new TTransportException("Not enough room in output buffer", e);
+    }
+  }
+
+  /**
+   * Get the underlying NIO ByteBuffer.
+   */
+  public ByteBuffer getByteBuffer() {
+    return byteBuffer;
+  }
+
+  /**
+   * Convenience method to call clear() on the underlying NIO ByteBuffer.
+   */
+  public TByteBuffer clear() {
+    byteBuffer.clear();
+    return this;
+  }
+
+  /**
+   * Convenience method to call flip() on the underlying NIO ByteBuffer.
+     */
+  public TByteBuffer flip() {
+    byteBuffer.flip();
+    return this;
+  }
+
+  /**
+   * Convenience method to convert the underlying NIO ByteBuffer to a
+   * plain old byte array.
+   */
+  public byte[] toByteArray() {
+    final byte[] data = new byte[byteBuffer.remaining()];
+    byteBuffer.slice().get(data);
+    return data;
+  }
+}
diff --git a/lib/java/test/org/apache/thrift/transport/TestTByteBuffer.java b/lib/java/test/org/apache/thrift/transport/TestTByteBuffer.java
new file mode 100644
index 00000000000..a73075b58cd
--- /dev/null
+++ b/lib/java/test/org/apache/thrift/transport/TestTByteBuffer.java
@@ -0,0 +1,36 @@
+package org.apache.thrift.transport;
+
+import junit.framework.TestCase;
+import org.apache.commons.codec.Charsets;
+import org.apache.thrift.TException;
+
+import java.nio.ByteBuffer;
+
+public class TestTByteBuffer extends TestCase {
+  public void testReadWrite() throws Exception {
+    final TByteBuffer byteBuffer = new TByteBuffer(ByteBuffer.allocate(16));
+    byteBuffer.write("Hello World".getBytes(Charsets.UTF_8));
+    assertEquals("Hello World", new String(byteBuffer.flip().toByteArray(), Charsets.UTF_8));
+  }
+
+  public void testReuseReadWrite() throws Exception {
+    final TByteBuffer byteBuffer = new TByteBuffer(ByteBuffer.allocate(16));
+    byteBuffer.write("Hello World".getBytes(Charsets.UTF_8));
+    assertEquals("Hello World", new String(byteBuffer.flip().toByteArray(), Charsets.UTF_8));
+
+    byteBuffer.clear();
+
+    byteBuffer.write("Goodbye Horses".getBytes(Charsets.UTF_8));
+    assertEquals("Goodbye Horses", new String(byteBuffer.flip().toByteArray(), Charsets.UTF_8));
+  }
+
+  public void testOverflow() throws Exception {
+    final TByteBuffer byteBuffer = new TByteBuffer(ByteBuffer.allocate(4));
+    try {
+      byteBuffer.write("Hello World".getBytes(Charsets.UTF_8));
+      fail("Expected write operation to fail with TTransportException");
+    } catch (TTransportException e) {
+      assertEquals("Not enough room in output buffer", e.getMessage());
+    }
+  }
+}