From f7f1a61efb227536aeba1190c1c0286a539952c7 Mon Sep 17 00:00:00 2001
From: Max Riedel <m_riedel_95@gmx.de>
Date: Sat, 15 Jul 2023 15:26:36 +0200
Subject: [PATCH] KAFKA-15123: Add tests for ChunkedBytesStream (#13941)

Reviewers: Divij Vaidya <diviv@amazon.com>
---
 .../common/utils/ChunkedBytesStreamTest.java  | 50 +++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/clients/src/test/java/org/apache/kafka/common/utils/ChunkedBytesStreamTest.java b/clients/src/test/java/org/apache/kafka/common/utils/ChunkedBytesStreamTest.java
index a1292386451e..10300a2e9568 100644
--- a/clients/src/test/java/org/apache/kafka/common/utils/ChunkedBytesStreamTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/utils/ChunkedBytesStreamTest.java
@@ -25,6 +25,7 @@
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.Mockito.spy;
 
 import java.io.IOException;
@@ -64,6 +65,16 @@ public void testCorrectnessForMethodReadFully(ByteBuffer input) throws IOExcepti
         assertArrayEquals(input.array(), got);
     }
 
+    @ParameterizedTest
+    @MethodSource("provideCasesWithInvalidInputsForMethodRead")
+    public void testInvalidInputsForMethodRead(byte[] b, int off, int len) throws IOException {
+        ByteBuffer buffer = ByteBuffer.allocate(16);
+
+        try (InputStream is = new ChunkedBytesStream(new ByteBufferInputStream(buffer.duplicate()), supplier, 10, false)) {
+            assertThrows(IndexOutOfBoundsException.class, () -> is.read(b, off, len));
+        }
+    }
+
     @ParameterizedTest
     @MethodSource("provideSourceBytebuffersForTest")
     public void testCorrectnessForMethodReadByte(ByteBuffer input) throws IOException {
@@ -149,6 +160,45 @@ public void testCorrectnessForMethodSkip(int bytesToPreRead, ByteBuffer inputBuf
         }
     }
 
+    @ParameterizedTest
+    @MethodSource("provideEdgeCaseInputForMethodSkip")
+    public void testEdgeCaseInputForMethodSkip(int bufferLength, long toSkip, boolean delegateSkipToSourceStream, long expected) throws IOException {
+        ByteBuffer inputBuf = ByteBuffer.allocate(bufferLength);
+        RANDOM.nextBytes(inputBuf.array());
+        inputBuf.rewind();
+
+        try (InputStream is = new ChunkedBytesStream(new ByteBufferInputStream(inputBuf.duplicate()), supplier, 10, delegateSkipToSourceStream)) {
+            assertEquals(expected, is.skip(toSkip));
+        }
+    }
+
+    private static Stream<Arguments> provideEdgeCaseInputForMethodSkip() {
+        int bufferLength = 16;
+        // Test toSkip larger than int and negative for both delegateToSourceStream true and false
+        return Stream.of(
+            Arguments.of(bufferLength, Integer.MAX_VALUE + 1L, true, bufferLength),
+            Arguments.of(bufferLength, -1, true, 0),
+            Arguments.of(bufferLength, Integer.MAX_VALUE + 1L, false, bufferLength),
+            Arguments.of(bufferLength, -1, false, 0)
+        );
+    }
+
+    private static Stream<Arguments> provideCasesWithInvalidInputsForMethodRead() {
+        byte[] b = new byte[16];
+        return Stream.of(
+            // negative off
+            Arguments.of(b, -1, b.length),
+            // negative len
+            Arguments.of(b, 0, -1),
+            // overflow off + len
+            Arguments.of(b, Integer.MAX_VALUE, 10),
+            // len greater than size of target array
+            Arguments.of(b, 0, b.length + 1),
+            // off + len greater than size of target array
+            Arguments.of(b, b.length - 1, 2)
+        );
+    }
+
     private static List<Arguments> provideSourceSkipValuesForTest() {
         ByteBuffer bufGreaterThanIntermediateBuf = ByteBuffer.allocate(16);
         RANDOM.nextBytes(bufGreaterThanIntermediateBuf.array());