TemporaryBuffer: fix toByteArray(limit)

Heap always copied whole blocks, which leads to AIOOBEs. LocalFile
didn't overwrite the method and thus caused NPEs.

Change-Id: Ia37d4a875df9f25d4825e6bc95fed7f0dff42afb
Signed-off-by: Thomas Wolf <thomas.wolf@paranor.ch>

2 files changed, 97 insertions, 3 deletions

diff --git a/org.eclipse.jgit.test/tst/org/eclipse/jgit/util/TemporaryBufferTest.java b/org.eclipse.jgit.test/tst/org/eclipse/jgit/util/TemporaryBufferTest.java
index 4e65ca7a4b..01dcde29bd 100644
--- a/org.eclipse.jgit.test/tst/org/eclipse/jgit/util/TemporaryBufferTest.java
+++ b/org.eclipse.jgit.test/tst/org/eclipse/jgit/util/TemporaryBufferTest.java
@@ -406,4 +406,69 @@ public class TemporaryBufferTest {
 			}
 		}
 	}
+
+	@Test
+	public void testHeapToByteArrayWithLimit() throws IOException {
+		int sz = 2 * Block.SZ;
+		try (TemporaryBuffer b = new TemporaryBuffer.Heap(sz / 2, sz)) {
+			for (int i = 0; i < sz; i++) {
+				b.write('a' + i % 26);
+			}
+			byte[] prefix = b.toByteArray(5);
+			assertEquals(5, prefix.length);
+			for (int i = 0; i < prefix.length; i++) {
+				assertEquals('a' + i % 26, prefix[i]);
+			}
+			prefix = b.toByteArray(Block.SZ + 37);
+			assertEquals(Block.SZ + 37, prefix.length);
+			for (int i = 0; i < prefix.length; i++) {
+				assertEquals('a' + i % 26, prefix[i]);
+			}
+			prefix = b.toByteArray(sz);
+			assertEquals(sz, prefix.length);
+			for (int i = 0; i < prefix.length; i++) {
+				assertEquals('a' + i % 26, prefix[i]);
+			}
+			prefix = b.toByteArray(sz + 37);
+			assertEquals(sz, prefix.length);
+			for (int i = 0; i < prefix.length; i++) {
+				assertEquals('a' + i % 26, prefix[i]);
+			}
+		}
+	}
+
+	@Test
+	public void testFileToByteArrayWithLimit() throws IOException {
+		@SuppressWarnings("resource") // Buffer is explicitly destroyed in finally block
+		TemporaryBuffer b = new TemporaryBuffer.LocalFile(null, 2 * Block.SZ);
+		int sz = 3 * Block.SZ;
+		try {
+			for (int i = 0; i < sz; i++) {
+				b.write('a' + i % 26);
+			}
+			b.close();
+			byte[] prefix = b.toByteArray(5);
+			assertEquals(5, prefix.length);
+			for (int i = 0; i < prefix.length; i++) {
+				assertEquals('a' + i % 26, prefix[i]);
+			}
+			prefix = b.toByteArray(Block.SZ + 37);
+			assertEquals(Block.SZ + 37, prefix.length);
+			for (int i = 0; i < prefix.length; i++) {
+				assertEquals('a' + i % 26, prefix[i]);
+			}
+			prefix = b.toByteArray(sz);
+			assertEquals(sz, prefix.length);
+			for (int i = 0; i < prefix.length; i++) {
+				assertEquals('a' + i % 26, prefix[i]);
+			}
+			prefix = b.toByteArray(sz + 37);
+			assertEquals(sz, prefix.length);
+			for (int i = 0; i < prefix.length; i++) {
+				assertEquals('a' + i % 26, prefix[i]);
+			}
+		} finally {
+			b.destroy();
+		}
+	}
 }
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/util/TemporaryBuffer.java b/org.eclipse.jgit/src/org/eclipse/jgit/util/TemporaryBuffer.java
index 1f0fedda6b..562eb05dd9 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/util/TemporaryBuffer.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/util/TemporaryBuffer.java
@@ -230,11 +230,16 @@ public abstract class TemporaryBuffer extends OutputStream {
 		if (Integer.MAX_VALUE < len)
 			throw new OutOfMemoryError(
 					JGitText.get().lengthExceedsMaximumArraySize);
-		final byte[] out = new byte[(int) len];
+		int length = (int) len;
+		final byte[] out = new byte[length];
 		int outPtr = 0;
 		for (Block b : blocks) {
-			System.arraycopy(b.buffer, 0, out, outPtr, b.count);
-			outPtr += b.count;
+			int toCopy = Math.min(length - outPtr, b.count);
+			System.arraycopy(b.buffer, 0, out, outPtr, toCopy);
+			outPtr += toCopy;
+			if (outPtr == length) {
+				break;
+			}
 		}
 		return out;
 	}
@@ -461,6 +466,30 @@ public abstract class TemporaryBuffer extends OutputStream {
 		}
 
 		@Override
+		public byte[] toByteArray(int limit) throws IOException {
+			if (onDiskFile == null) {
+				return super.toByteArray(limit);
+			}
+			final long len = Math.min(length(), limit);
+			if (Integer.MAX_VALUE < len) {
+				throw new OutOfMemoryError(
+						JGitText.get().lengthExceedsMaximumArraySize);
+			}
+			final byte[] out = new byte[(int) len];
+			try (FileInputStream in = new FileInputStream(onDiskFile)) {
+				int read = 0;
+				int chunk;
+				while ((chunk = in.read(out, read, out.length - read)) >= 0) {
+					read += chunk;
+					if (read == out.length) {
+						break;
+					}
+				}
+			}
+			return out;
+		}
+
+		@Override
 		public void writeTo(OutputStream os, ProgressMonitor pm)
 				throws IOException {
 			if (onDiskFile == null) {