mobile: Use direct ByteBuffer to pass data between C++ and Java (#32715)

Signed-off-by: Fredy Wijaya <fredyw@google.com>

mobile/library/java/io/envoyproxy/envoymobile/engine/BUILD

@@ -30,6 +30,7 @@ android_library(
 java_library(
     name = "envoy_base_engine_lib",
     srcs = [
+        "ByteBuffers.java",
         "EnvoyConfiguration.java",
         "EnvoyEngine.java",
         "EnvoyEngineImpl.java",

mobile/library/java/io/envoyproxy/envoymobile/engine/ByteBuffers.java

@@ -0,0 +1,17 @@
+package io.envoyproxy.envoymobile.engine;
+
+import java.nio.ByteBuffer;
+
+public class ByteBuffers {
+  /**
+   * Copies the specified `ByteBuffer` into a new `ByteBuffer`. The `ByteBuffer` created will
+   * be backed by `byte[]`.
+   */
+  public static ByteBuffer copy(ByteBuffer byteBuffer) {
+    byte[] bytes = new byte[byteBuffer.capacity()];
+    byteBuffer.get(bytes);
+    return ByteBuffer.wrap(bytes);
+  }
+
+  private ByteBuffers() {}
+}

mobile/library/java/io/envoyproxy/envoymobile/engine/JvmCallbackContext.java

@@ -1,10 +1,6 @@
 package io.envoyproxy.envoymobile.engine;
 
 import java.nio.ByteBuffer;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 
 import io.envoyproxy.envoymobile.engine.types.EnvoyHTTPCallbacks;
@@ -76,14 +72,15 @@ public Object onResponseTrailers(long trailerCount, long[] streamIntel) {
    * @param streamIntel, internal HTTP stream metrics, context, and other details.
    * @return Object,     not used for response callbacks.
    */
-  public Object onResponseData(byte[] data, boolean endStream, long[] streamIntel) {
+  public Object onResponseData(ByteBuffer data, boolean endStream, long[] streamIntel) {
+    // Create a copy of the `data` because the `data` uses direct `ByteBuffer` and the `data` will
+    // be destroyed after calling this callback.
+    ByteBuffer copiedData = ByteBuffers.copy(data);
     callbacks.getExecutor().execute(new Runnable() {
       public void run() {
-        ByteBuffer dataBuffer = ByteBuffer.wrap(data);
-        callbacks.onData(dataBuffer, endStream, new EnvoyStreamIntelImpl(streamIntel));
+        callbacks.onData(copiedData, endStream, new EnvoyStreamIntelImpl(streamIntel));
       }
     });
-
     return null;
   }
 

mobile/library/java/io/envoyproxy/envoymobile/engine/JvmFilterContext.java

@@ -65,10 +65,12 @@ public Object onRequestHeaders(long headerCount, boolean endStream, long[] strea
    * @param streamIntel, internal HTTP stream metrics, context, and other details.
    * @return Object[],   pair of HTTP filter status and optional modified data.
    */
-  public Object onRequestData(byte[] data, boolean endStream, long[] streamIntel) {
-    ByteBuffer dataBuffer = ByteBuffer.wrap(data);
+  public Object onRequestData(ByteBuffer data, boolean endStream, long[] streamIntel) {
+    // Create a copy of the `data` because the `data` uses direct `ByteBuffer` and the `data` will
+    // be destroyed after calling this callback.
+    ByteBuffer copiedData = ByteBuffers.copy(data);
     return toJniFilterDataStatus(
-        filter.onRequestData(dataBuffer, endStream, new EnvoyStreamIntelImpl(streamIntel)));
+        filter.onRequestData(copiedData, endStream, new EnvoyStreamIntelImpl(streamIntel)));
   }
 
   /**
@@ -108,10 +110,12 @@ public Object onResponseHeaders(long headerCount, boolean endStream, long[] stre
    * @param streamIntel, internal HTTP stream metrics, context, and other details.
    * @return Object[],   pair of HTTP filter status and optional modified data.
    */
-  public Object onResponseData(byte[] data, boolean endStream, long[] streamIntel) {
-    ByteBuffer dataBuffer = ByteBuffer.wrap(data);
+  public Object onResponseData(ByteBuffer data, boolean endStream, long[] streamIntel) {
+    // Create a copy of the `data` because the `data` uses direct `ByteBuffer` and the `data` will
+    // be destroyed after calling this callback.
+    ByteBuffer copiedData = ByteBuffers.copy(data);
     return toJniFilterDataStatus(
-        filter.onResponseData(dataBuffer, endStream, new EnvoyStreamIntelImpl(streamIntel)));
+        filter.onResponseData(copiedData, endStream, new EnvoyStreamIntelImpl(streamIntel)));
   }
 
   /**
@@ -138,22 +142,24 @@ public Object onResponseTrailers(long trailerCount, long[] streamIntel) {
    * @param streamIntel,  internal HTTP stream metrics, context, and other details.
    * @return Object[],    tuple of status with updated entities to be forwarded.
    */
-  public Object onResumeRequest(long headerCount, byte[] data, long trailerCount, boolean endStream,
-                                long[] streamIntel) {
+  public Object onResumeRequest(long headerCount, ByteBuffer data, long trailerCount,
+                                boolean endStream, long[] streamIntel) {
+    // Create a copy of the `data` because the `data` uses direct `ByteBuffer` and the `data` will
+    // be destroyed after calling this callback.
+    ByteBuffer copiedData = ByteBuffers.copy(data);
     // Headers are optional in this call, and a negative length indicates omission.
     Map<String, List<String>> headers = null;
     if (headerCount >= 0) {
       assert headerUtility.validateCount(headerCount);
       headers = headerUtility.retrieveHeaders();
     }
-    ByteBuffer dataBuffer = data == null ? null : ByteBuffer.wrap(data);
     // Trailers are optional in this call, and a negative length indicates omission.
     Map<String, List<String>> trailers = null;
     if (trailerCount >= 0) {
       assert trailerUtility.validateCount(trailerCount);
       trailers = trailerUtility.retrieveHeaders();
     }
-    return toJniFilterResumeStatus(filter.onResumeRequest(headers, dataBuffer, trailers, endStream,
+    return toJniFilterResumeStatus(filter.onResumeRequest(headers, copiedData, trailers, endStream,
                                                           new EnvoyStreamIntelImpl(streamIntel)));
   }
 
@@ -167,22 +173,24 @@ public Object onResumeRequest(long headerCount, byte[] data, long trailerCount,
    * @param streamIntel,  internal HTTP stream metrics, context, and other details.
    * @return Object[],    tuple of status with updated entities to be forwarded.
    */
-  public Object onResumeResponse(long headerCount, byte[] data, long trailerCount,
+  public Object onResumeResponse(long headerCount, ByteBuffer data, long trailerCount,
                                  boolean endStream, long[] streamIntel) {
+    // Create a copy of the `data` because the `data` uses direct `ByteBuffer` and the `data` will
+    // be destroyed after calling this callback.
+    ByteBuffer copiedData = ByteBuffers.copy(data);
     // Headers are optional in this call, and a negative length indicates omission.
     Map<String, List<String>> headers = null;
     if (headerCount >= 0) {
       assert headerUtility.validateCount(headerCount);
       headers = headerUtility.retrieveHeaders();
     }
-    ByteBuffer dataBuffer = data == null ? null : ByteBuffer.wrap(data);
     // Trailers are optional in this call, and a negative length indicates omission.
     Map<String, List<String>> trailers = null;
     if (trailerCount >= 0) {
       assert trailerUtility.validateCount(trailerCount);
       trailers = trailerUtility.retrieveHeaders();
     }
-    return toJniFilterResumeStatus(filter.onResumeResponse(headers, dataBuffer, trailers, endStream,
+    return toJniFilterResumeStatus(filter.onResumeResponse(headers, copiedData, trailers, endStream,
                                                            new EnvoyStreamIntelImpl(streamIntel)));
   }
 

mobile/library/java/io/envoyproxy/envoymobile/engine/types/EnvoyHTTPFilter.java

@@ -1,7 +1,6 @@
 package io.envoyproxy.envoymobile.engine.types;
 
 import java.nio.ByteBuffer;
-import java.util.concurrent.Executor;
 import java.util.List;
 import java.util.Map;
 

mobile/library/jni/jni_helper.cc

@@ -189,6 +189,13 @@ void JniHelper::callStaticVoidMethod(jclass clazz, jmethodID method_id, ...) {
   rethrowException();
 }
 
+LocalRefUniquePtr<jobject> JniHelper::newDirectByteBuffer(void* address, jlong capacity) {
+  LocalRefUniquePtr<jobject> result(env_->NewDirectByteBuffer(address, capacity),
+                                    LocalRefDeleter(env_));
+  rethrowException();
+  return result;
+}
+
 jlong JniHelper::getDirectBufferCapacity(jobject buffer) {
   return env_->GetDirectBufferCapacity(buffer);
 }

mobile/library/jni/jni_helper.h

@@ -369,6 +369,14 @@ class JniHelper {
     return result;
   }
 
+  /**
+   * Allocates and returns a direct `java.nio.ByteBuffer` referring to the block of memory starting
+   * at the memory address `address` and extending `capacity` bytes.
+   *
+   * https://docs.oracle.com/en/java/javase/17/docs/specs/jni/functions.html#newdirectbytebuffer
+   */
+  LocalRefUniquePtr<jobject> newDirectByteBuffer(void* address, jlong capacity);
+
   /**
    * Returns the capacity of the memory region referenced by the given `java.nio.Buffer` object.
    *

mobile/library/jni/jni_impl.cc

@@ -343,11 +343,11 @@ static Envoy::JNI::LocalRefUniquePtr<jobjectArray> jvm_on_data(const char* metho
 
   Envoy::JNI::LocalRefUniquePtr<jclass> jcls_JvmCallbackContext =
       jni_helper.getObjectClass(j_context);
-  jmethodID jmid_onData =
-      jni_helper.getMethodId(jcls_JvmCallbackContext.get(), method, "([BZ[J)Ljava/lang/Object;");
+  jmethodID jmid_onData = jni_helper.getMethodId(jcls_JvmCallbackContext.get(), method,
+                                                 "(Ljava/nio/ByteBuffer;Z[J)Ljava/lang/Object;");
 
-  Envoy::JNI::LocalRefUniquePtr<jbyteArray> j_data =
-      Envoy::JNI::envoyDataToJavaByteArray(jni_helper, data);
+  Envoy::JNI::LocalRefUniquePtr<jobject> j_data =
+      Envoy::JNI::envoyDataToJavaByteBuffer(jni_helper, data);
   Envoy::JNI::LocalRefUniquePtr<jlongArray> j_stream_intel =
       Envoy::JNI::envoyStreamIntelToJavaLongArray(jni_helper, stream_intel);
   Envoy::JNI::LocalRefUniquePtr<jobjectArray> result = jni_helper.callObjectMethod<jobjectArray>(
@@ -605,10 +605,10 @@ jvm_http_filter_on_resume(const char* method, envoy_headers* headers, envoy_data
     headers_length = static_cast<jlong>(headers->length);
     passHeaders("passHeader", *headers, j_context);
   }
-  Envoy::JNI::LocalRefUniquePtr<jbyteArray> j_in_data = Envoy::JNI::LocalRefUniquePtr<jbyteArray>(
+  Envoy::JNI::LocalRefUniquePtr<jobject> j_in_data = Envoy::JNI::LocalRefUniquePtr<jobject>(
       nullptr, Envoy::JNI::LocalRefDeleter(jni_helper.getEnv()));
   if (data) {
-    j_in_data = Envoy::JNI::envoyDataToJavaByteArray(jni_helper, *data);
+    j_in_data = Envoy::JNI::envoyDataToJavaByteBuffer(jni_helper, *data);
   }
   jlong trailers_length = -1;
   if (trailers) {
@@ -620,8 +620,8 @@ jvm_http_filter_on_resume(const char* method, envoy_headers* headers, envoy_data
 
   Envoy::JNI::LocalRefUniquePtr<jclass> jcls_JvmCallbackContext =
       jni_helper.getObjectClass(j_context);
-  jmethodID jmid_onResume =
-      jni_helper.getMethodId(jcls_JvmCallbackContext.get(), method, "(J[BJZ[J)Ljava/lang/Object;");
+  jmethodID jmid_onResume = jni_helper.getMethodId(
+      jcls_JvmCallbackContext.get(), method, "(JLjava/nio/ByteBuffer;JZ[J)Ljava/lang/Object;");
   // Note: be careful of JVM types. Before we casted to jlong we were getting integer problems.
   // TODO: make this cast safer.
   Envoy::JNI::LocalRefUniquePtr<jobjectArray> result = jni_helper.callObjectMethod<jobjectArray>(

mobile/library/jni/jni_utility.cc

@@ -79,6 +79,11 @@ LocalRefUniquePtr<jbyteArray> envoyDataToJavaByteArray(JniHelper& jni_helper, en
   return j_data;
 }
 
+LocalRefUniquePtr<jobject> envoyDataToJavaByteBuffer(JniHelper& jni_helper, envoy_data data) {
+  return jni_helper.newDirectByteBuffer(
+      const_cast<void*>(reinterpret_cast<const void*>(data.bytes)), data.length);
+}
+
 LocalRefUniquePtr<jlongArray> envoyStreamIntelToJavaLongArray(JniHelper& jni_helper,
                                                               envoy_stream_intel stream_intel) {
   LocalRefUniquePtr<jlongArray> j_array = jni_helper.newLongArray(4);
@@ -151,10 +156,10 @@ envoy_data javaByteBufferToEnvoyData(JniHelper& jni_helper, jobject j_data) {
     return native_data;
   }
 
-  return javaByteBufferToEnvoyData(jni_helper, j_data, static_cast<size_t>(data_length));
+  return javaByteBufferToEnvoyData(jni_helper, j_data, data_length);
 }
 
-envoy_data javaByteBufferToEnvoyData(JniHelper& jni_helper, jobject j_data, size_t data_length) {
+envoy_data javaByteBufferToEnvoyData(JniHelper& jni_helper, jobject j_data, jlong data_length) {
   // Returns nullptr if the buffer is not a direct buffer.
   uint8_t* direct_address = jni_helper.getDirectBufferAddress<uint8_t*>(j_data);
 

mobile/library/jni/jni_utility.h

@@ -57,6 +57,9 @@ envoy_data javaByteArrayToEnvoyData(JniHelper& jni_helper, jbyteArray j_data, si
 /** Converts from `envoy_data` to Java byte array. */
 LocalRefUniquePtr<jbyteArray> envoyDataToJavaByteArray(JniHelper& jni_helper, envoy_data data);
 
+/** Converts from `envoy_data to `java.nio.ByteBuffer`. */
+LocalRefUniquePtr<jobject> envoyDataToJavaByteBuffer(JniHelper& jni_helper, envoy_data data);
+
 /** Converts from `envoy_stream_intel` to Java long array. */
 LocalRefUniquePtr<jlongArray> envoyStreamIntelToJavaLongArray(JniHelper& jni_helper,
                                                               envoy_stream_intel stream_intel);
@@ -76,7 +79,7 @@ LocalRefUniquePtr<jstring> envoyDataToJavaString(JniHelper& jni_helper, envoy_da
 envoy_data javaByteBufferToEnvoyData(JniHelper& jni_helper, jobject j_data);
 
 /** Converts from Java `ByteBuffer` to `envoy_data` with the given length. */
-envoy_data javaByteBufferToEnvoyData(JniHelper& jni_helper, jobject j_data, size_t data_length);
+envoy_data javaByteBufferToEnvoyData(JniHelper& jni_helper, jobject j_data, jlong data_length);
 
 /** Returns the pointer of conversion from Java `ByteBuffer` to `envoy_data`. */
 envoy_data* javaByteBufferToEnvoyDataPtr(JniHelper& jni_helper, jobject j_data);

mobile/test/java/io/envoyproxy/envoymobile/engine/BUILD

@@ -67,3 +67,23 @@ envoy_mobile_android_test(
         "//test/kotlin/io/envoyproxy/envoymobile/mocks:mocks_lib",
     ],
 )
+
+envoy_mobile_android_test(
+    name = "byte_buffers_test",
+    srcs = [
+        "ByteBuffersTest.java",
+    ],
+    associates = ["//library/kotlin/io/envoyproxy/envoymobile:envoy_interfaces_lib"],
+    native_deps = [
+        "//test/jni:libenvoy_jni_with_test_extensions.so",
+    ] + select({
+        "@platforms//os:macos": [
+            "//test/jni:libenvoy_jni_with_test_extensions_jnilib",
+        ],
+        "//conditions:default": [],
+    }),
+    native_lib_name = "envoy_jni_with_test_extensions",
+    deps = [
+        "//library/java/io/envoyproxy/envoymobile/engine:envoy_base_engine_lib",
+    ],
+)

mobile/test/java/io/envoyproxy/envoymobile/engine/ByteBuffersTest.java

@@ -0,0 +1,26 @@
+package io.envoyproxy.envoymobile.engine;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.nio.ByteBuffer;
+
+@RunWith(AndroidJUnit4.class)
+public class ByteBuffersTest {
+  @Test
+  public void testCopy() {
+    ByteBuffer source = ByteBuffer.allocateDirect(3);
+    source.put((byte)1);
+    source.put((byte)2);
+    source.put((byte)3);
+    source.flip();
+
+    ByteBuffer dest = ByteBuffers.copy(source);
+    source.flip();
+    assertThat(dest).isEqualTo(source);
+  }
+}

mobile/test/java/io/envoyproxy/envoymobile/jni/JniHelperTest.java

@@ -7,6 +7,9 @@
 import org.junit.runner.RunWith;
 import org.robolectric.RobolectricTestRunner;
 
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+
 @RunWith(RobolectricTestRunner.class)
 public class JniHelperTest {
   public JniHelperTest() { System.loadLibrary("envoy_jni_helper_test"); }
@@ -84,6 +87,7 @@ public static native boolean callStaticBooleanMethod(Class<?> clazz, String name
                                                        String signature);
   public static native void callStaticVoidMethod(Class<?> clazz, String name, String signature);
   public static native Object callStaticObjectMethod(Class<?> clazz, String name, String signature);
+  public static native Object newDirectByteBuffer();
 
   //================================================================================
   // Object methods used for Call<Type>Method tests.
@@ -424,4 +428,13 @@ public void testCallStaticObjectMethod() {
         callStaticObjectMethod(JniHelperTest.class, "staticObjectMethod", "()Ljava/lang/String;"))
         .isEqualTo("Hello");
   }
+
+  @Test
+  public void testNewDirectByteBuffer() {
+    ByteBuffer byteBuffer = ((ByteBuffer)newDirectByteBuffer()).order(ByteOrder.LITTLE_ENDIAN);
+    assertThat(byteBuffer.capacity()).isEqualTo(3);
+    assertThat(byteBuffer.get(0)).isEqualTo(1);
+    assertThat(byteBuffer.get(1)).isEqualTo(2);
+    assertThat(byteBuffer.get(2)).isEqualTo(3);
+  }
 }

mobile/test/java/io/envoyproxy/envoymobile/utilities/ByteBuffersTest.java

@@ -0,0 +1,7 @@
+package io.envoyproxy.envoymobile.utilities;
+
+import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+
+@RunWith(RobolectricTestRunner.class)
+public class ByteBuffersTest {}

mobile/test/jni/jni_helper_test.cc

@@ -224,3 +224,13 @@ Java_io_envoyproxy_envoymobile_jni_JniHelperTest_callStaticObjectMethod(JNIEnv*
   jmethodID method_id = jni_helper.getStaticMethodId(clazz, name_ptr.get(), sig_ptr.get());
   return jni_helper.callStaticObjectMethod(clazz, method_id).release();
 }
+
+extern "C" JNIEXPORT jobject JNICALL
+Java_io_envoyproxy_envoymobile_jni_JniHelperTest_newDirectByteBuffer(JNIEnv* env, jclass) {
+  Envoy::JNI::JniHelper jni_helper(env);
+  char* bytes = new char[3];
+  bytes[0] = 1;
+  bytes[1] = 2;
+  bytes[2] = 3;
+  return jni_helper.newDirectByteBuffer(reinterpret_cast<void*>(bytes), sizeof(char) * 3).release();
+}