From f96d1ef744a435d3557ee725a63407c4530501bf Mon Sep 17 00:00:00 2001
From: Diogo Monteiro <diogo.monteiro12@gmail.com>
Date: Thu, 17 Sep 2020 19:55:24 +0100
Subject: [PATCH] [Java] allowing larger buffer sizes when growing a byte
 buffer (#6118)

Related with issue #6113.

`(old_buf_size & 0xC0000000) != 0` checks if we can duplicate old_buf_size
 and still be under 2GB (by checking if bit 30 or 31 is 1). This doesn't
allow buffers larger than 1GB.

The strategy now is to allocate a buffer with the maximum array size
when we detect that we are overflowing the 2GB.

Also changed default buffer size to 1024.
---
 .../google/flatbuffers/FlatBufferBuilder.java | 33 ++++++++++++++++---
 1 file changed, 28 insertions(+), 5 deletions(-)

diff --git a/java/com/google/flatbuffers/FlatBufferBuilder.java b/java/com/google/flatbuffers/FlatBufferBuilder.java
index 939e1b0cd..a954d9fbb 100644
--- a/java/com/google/flatbuffers/FlatBufferBuilder.java
+++ b/java/com/google/flatbuffers/FlatBufferBuilder.java
@@ -53,6 +53,19 @@ public class FlatBufferBuilder {
     Map<String, Integer> string_pool; // map used to cache shared strings.
     /// @endcond
 
+
+    /**
+     * Maximum size of buffer to allocate. If we're allocating arrays on the heap,
+     * the header size of the array counts towards its maximum size.
+     */
+    private static final int MAX_BUFFER_SIZE = Integer.MAX_VALUE - 8;
+
+    /**
+     * Default buffer size that is allocated if an initial size is not given, or is
+     * non positive.
+     */
+    private static final int DEFAULT_BUFFER_SIZE = 1024;
+
     /**
      * Start with a buffer of size `initial_size`, then grow as required.
      *
@@ -74,7 +87,7 @@ public class FlatBufferBuilder {
     public FlatBufferBuilder(int initial_size, ByteBufferFactory bb_factory,
                              ByteBuffer existing_bb, Utf8 utf8) {
         if (initial_size <= 0) {
-          initial_size = 1;
+          initial_size = DEFAULT_BUFFER_SIZE;
         }
         this.bb_factory = bb_factory;
         if (existing_bb != null) {
@@ -101,7 +114,7 @@ public class FlatBufferBuilder {
      * Start with a buffer of 1KiB, then grow as required.
      */
     public FlatBufferBuilder() {
-        this(1024);
+        this(DEFAULT_BUFFER_SIZE);
     }
 
     /**
@@ -247,9 +260,19 @@ public class FlatBufferBuilder {
      */
     static ByteBuffer growByteBuffer(ByteBuffer bb, ByteBufferFactory bb_factory) {
         int old_buf_size = bb.capacity();
-        if ((old_buf_size & 0xC0000000) != 0)  // Ensure we don't grow beyond what fits in an int.
-            throw new AssertionError("FlatBuffers: cannot grow buffer beyond 2 gigabytes.");
-        int new_buf_size = old_buf_size == 0 ? 1 : old_buf_size << 1;
+
+        int new_buf_size;
+
+        if (old_buf_size == 0) {
+            new_buf_size = DEFAULT_BUFFER_SIZE;
+        }
+        else {
+            if (old_buf_size == MAX_BUFFER_SIZE) { // Ensure we don't grow beyond what fits in an int.
+                throw new AssertionError("FlatBuffers: cannot grow buffer beyond 2 gigabytes.");
+            }
+            new_buf_size = (old_buf_size & 0xC0000000) != 0 ? MAX_BUFFER_SIZE : old_buf_size << 1;
+        }
+
         bb.position(0);
         ByteBuffer nbb = bb_factory.newByteBuffer(new_buf_size);
         new_buf_size = nbb.clear().capacity(); // Ensure the returned buffer is treated as empty