From 3fbd456c9d83e914a986a5c8459f0a4478bd939e Mon Sep 17 00:00:00 2001
From: David Tobin <dtobin6@gmail.com>
Date: Wed, 20 Oct 2021 07:13:43 -0700
Subject: [PATCH] Reduce byte[] allocations during reads via
 ApacheCtrDecryptingSeekableInput (#561) (#576)

Co-authored-by: Jacek Lach <jacekklach@gmail.com>
---
 .../io/ApacheCtrDecryptingSeekableInput.java  | 37 +++++++---
 ...ApacheCtrDecryptingSeekableInputTests.java | 71 +++++++++++++++++++
 2 files changed, 97 insertions(+), 11 deletions(-)
 create mode 100644 crypto-core/src/test/java/com/palantir/crypto2/io/ApacheCtrDecryptingSeekableInputTests.java

diff --git a/crypto-core/src/main/java/com/palantir/crypto2/io/ApacheCtrDecryptingSeekableInput.java b/crypto-core/src/main/java/com/palantir/crypto2/io/ApacheCtrDecryptingSeekableInput.java
index 812bfe440..7d5c6074b 100644
--- a/crypto-core/src/main/java/com/palantir/crypto2/io/ApacheCtrDecryptingSeekableInput.java
+++ b/crypto-core/src/main/java/com/palantir/crypto2/io/ApacheCtrDecryptingSeekableInput.java
@@ -16,6 +16,7 @@
 
 package com.palantir.crypto2.io;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.palantir.crypto2.cipher.ApacheCiphers;
 import com.palantir.crypto2.keys.KeyMaterial;
 import com.palantir.seekio.SeekableInput;
@@ -66,11 +67,12 @@ public void close() throws IOException {
         super.close();
     }
 
-    private static final class InputAdapter implements Input {
+    @VisibleForTesting
+    static final class InputAdapter implements Input {
+        private final SeekableInput input;
+        private final byte[] readBuffer = new byte[BUFFER_SIZE];
 
-        private SeekableInput input;
-
-        private InputAdapter(SeekableInput input) {
+        InputAdapter(SeekableInput input) {
             this.input = input;
         }
 
@@ -82,14 +84,28 @@ public int read(long position, byte[] buffer, int offset, int length) throws IOE
 
         @Override
         public int read(ByteBuffer dst) throws IOException {
-            byte[] bytes = new byte[dst.remaining()];
-            int read = input.read(bytes, 0, bytes.length);
-
-            if (read != -1) {
-                dst.put(bytes, 0, read);
+            int toRead = dst.remaining();
+            int totalRead = 0;
+
+            while (toRead > 0) {
+                int chunk = Math.min(toRead, readBuffer.length);
+                int read = input.read(readBuffer, 0, chunk);
+
+                if (read == -1) {
+                    if (totalRead == 0) {
+                        // first read hit EOF
+                        return -1;
+                    } else {
+                        return totalRead;
+                    }
+                } else {
+                    dst.put(readBuffer, 0, read);
+                    totalRead += read;
+                    toRead -= read;
+                }
             }
 
-            return read;
+            return totalRead;
         }
 
         @Override
@@ -113,5 +129,4 @@ public void close() throws IOException {
             input.close();
         }
     }
-
 }
diff --git a/crypto-core/src/test/java/com/palantir/crypto2/io/ApacheCtrDecryptingSeekableInputTests.java b/crypto-core/src/test/java/com/palantir/crypto2/io/ApacheCtrDecryptingSeekableInputTests.java
new file mode 100644
index 000000000..07a916839
--- /dev/null
+++ b/crypto-core/src/test/java/com/palantir/crypto2/io/ApacheCtrDecryptingSeekableInputTests.java
@@ -0,0 +1,71 @@
+/*
+ * (c) Copyright 2021 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.palantir.crypto2.io;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.palantir.seekio.InMemorySeekableDataInput;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Random;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+public final class ApacheCtrDecryptingSeekableInputTests {
+    private static final int NUM_BYTES = 1024 * 1024;
+    private static final Random random = new Random(0);
+    private static byte[] data;
+
+    @BeforeClass
+    public static void beforeClass() throws IOException {
+        data = new byte[NUM_BYTES];
+        random.nextBytes(data);
+    }
+
+    @Test
+    public void testEmptyRead() throws IOException {
+        ByteBuffer dst = ByteBuffer.allocate(1024);
+        byte[] emptyData = new byte[] {};
+
+        ApacheCtrDecryptingSeekableInput.InputAdapter adapter = inputAdapter(emptyData);
+        assertThat(adapter.read(dst)).isEqualTo(-1);
+        assertThat(dst.position()).isEqualTo(0);
+    }
+
+    @Test
+    public void testFullRead() throws IOException {
+        ByteBuffer dst = ByteBuffer.allocate(2 * NUM_BYTES);
+
+        ApacheCtrDecryptingSeekableInput.InputAdapter adapter = inputAdapter(data);
+        assertThat(adapter.read(dst)).isEqualTo(NUM_BYTES);
+        assertThat(dst.position()).isEqualTo(NUM_BYTES);
+    }
+
+    @Test
+    public void testPartialRead() throws IOException {
+        int toRead = NUM_BYTES / 2;
+        ByteBuffer dst = ByteBuffer.allocate(toRead);
+
+        ApacheCtrDecryptingSeekableInput.InputAdapter adapter = inputAdapter(data);
+        assertThat(adapter.read(dst)).isEqualTo(toRead);
+        assertThat(dst.position()).isEqualTo(toRead);
+    }
+
+    private ApacheCtrDecryptingSeekableInput.InputAdapter inputAdapter(byte[] inputData) {
+        return new ApacheCtrDecryptingSeekableInput.InputAdapter(new InMemorySeekableDataInput(inputData));
+    }
+}