1 files changed, 50 insertions, 0 deletions

diff --git a/java/org/brotli/wrapper/dec/decoder_jni.cc b/java/org/brotli/wrapper/dec/decoder_jni.cc
index 268a10b..e10ffe3 100644
--- a/java/org/brotli/wrapper/dec/decoder_jni.cc
+++ b/java/org/brotli/wrapper/dec/decoder_jni.cc
@@ -15,6 +15,9 @@ namespace {
 typedef struct DecoderHandle {
   BrotliDecoderState* state;
 
+  jobject dictionary_refs[15];
+  size_t dictionary_count;
+
   uint8_t* input_start;
   size_t input_offset;
   size_t input_length;
@@ -54,6 +57,10 @@ Java_org_brotli_wrapper_dec_DecoderJNI_nativeCreate(
   ok = !!handle;
 
   if (ok) {
+    for (int i = 0; i < 15; ++i) {
+      handle->dictionary_refs[i] = nullptr;
+    }
+    handle->dictionary_count = 0;
     handle->input_offset = 0;
     handle->input_length = 0;
     handle->input_start = nullptr;
@@ -193,10 +200,53 @@ Java_org_brotli_wrapper_dec_DecoderJNI_nativeDestroy(
   env->GetLongArrayRegion(ctx, 0, 3, context);
   DecoderHandle* handle = getHandle(reinterpret_cast<void*>(context[0]));
   BrotliDecoderDestroyInstance(handle->state);
+  for (size_t i = 0; i < handle->dictionary_count; ++i) {
+    env->DeleteGlobalRef(handle->dictionary_refs[i]);
+  }
   delete[] handle->input_start;
   delete handle;
 }
 
+JNIEXPORT jboolean JNICALL
+Java_org_brotli_wrapper_dec_DecoderJNI_nativeAttachDictionary(
+    JNIEnv* env, jobject /*jobj*/, jlongArray ctx, jobject dictionary) {
+  jlong context[3];
+  env->GetLongArrayRegion(ctx, 0, 3, context);
+  DecoderHandle* handle = getHandle(reinterpret_cast<void*>(context[0]));
+  jobject ref = nullptr;
+  uint8_t* address = nullptr;
+  jlong capacity = 0;
+
+  bool ok = true;
+  if (ok && !dictionary) {
+    ok = false;
+  }
+  if (ok && handle->dictionary_count >= 15) {
+    ok = false;
+  }
+  if (ok) {
+    ref = env->NewGlobalRef(dictionary);
+    ok = !!ref;
+  }
+  if (ok) {
+    handle->dictionary_refs[handle->dictionary_count] = ref;
+    handle->dictionary_count++;
+    address = static_cast<uint8_t*>(env->GetDirectBufferAddress(ref));
+    ok = !!address;
+  }
+  if (ok) {
+    capacity = env->GetDirectBufferCapacity(ref);
+    ok = (capacity > 0) && (capacity < (1 << 30));
+  }
+  if (ok) {
+    size_t size = static_cast<size_t>(capacity);
+    ok = !!BrotliDecoderAttachDictionary(handle->state,
+        BROTLI_SHARED_DICTIONARY_RAW, size, address);
+  }
+
+  return static_cast<jboolean>(ok);
+}
+
 #ifdef __cplusplus
 }
 #endif