From d1d82065b9db9127b32a083f71a07b42d348d44c Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Thu, 5 Nov 2020 18:02:41 -0500
Subject: [PATCH] [Java] Fixes an error allocating large direct byte buffers
 during OnnxTensor creation (#5619)

* Fixing an error with allocating large direct byte buffers during tensor creation.

* Removing the redundant overflow check.
---
 java/src/main/java/ai/onnxruntime/OnnxTensor.java | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/java/src/main/java/ai/onnxruntime/OnnxTensor.java b/java/src/main/java/ai/onnxruntime/OnnxTensor.java
index 752652bf22..f345d5626d 100644
--- a/java/src/main/java/ai/onnxruntime/OnnxTensor.java
+++ b/java/src/main/java/ai/onnxruntime/OnnxTensor.java
@@ -704,15 +704,19 @@ public class OnnxTensor implements OnnxValue {
   private static OnnxTensor createTensor(
       OnnxJavaType type, OrtAllocator allocator, Buffer data, long[] shape) throws OrtException {
     int bufferPos;
-    int bufferSize = data.remaining() * type.size;
-    if ((bufferSize < 0) || (bufferSize > ((Integer.MAX_VALUE / type.size) - 10))) {
-      // overflowed as we can't make a direct byte buffer of that size
+    long bufferSizeLong = data.remaining() * (long) type.size;
+    if (bufferSizeLong > (Integer.MAX_VALUE - (8 * type.size))) {
+      // The maximum direct byte buffer size is a little below Integer.MAX_VALUE depending
+      // on the JVM, so we check for something 8 elements below the maximum size which
+      // should be allocatable (assuming there is enough memory) on all 64-bit JVMs.
       throw new IllegalStateException(
           "Cannot allocate a direct buffer of the requested size and type, size "
               + data.remaining()
               + ", type = "
               + type);
     }
+    // Now we know we're in range
+    int bufferSize = data.remaining() * type.size;
     Buffer tmp;
     if (data.isDirect()) {
       tmp = data;