More tests for LimitedInputStream

digipost · Aug 2, 2024 · fb5d6c8 · fb5d6c8
1 parent ecb3a8e
commit fb5d6c8
Show file tree

Hide file tree

Showing 2 changed files with 103 additions and 7 deletions.
diff --git a/src/main/java/no/digipost/io/LimitedInputStream.java b/src/main/java/no/digipost/io/LimitedInputStream.java
@@ -137,12 +137,21 @@ public int read() throws IOException {
      */
     @Override
     public int read(byte[] b, int off, int len) throws IOException {
-        int allowedRemaing = (int)(maxBytesCount - count);
+        int allowedRemaining = (int)(maxBytesCount - count);
+        if (len == 0) {
+            return allowedRemaining < 0 ? reachedLimit() : 0;
+        }
         int res;
-        if (allowedRemaing > 0) {
-            res = super.read(b, off, min(len, allowedRemaing));
+        if (allowedRemaining > 0) {
+            // reads at most what is allowed according to the limit
+            res = super.read(b, off, min(len, allowedRemaining));
             count += max(res, 1);
         } else {
+            // The stream is not allowed to read any more bytes.
+            // Delegating to the single byte read method which handles
+            // if the stream is already beyond its set limit, and in
+            // any case at most reads one byte to determine if it has
+            // reached the EOF or contains more data.
             res = read();
         }
         return res;

diff --git a/src/test/java/no/digipost/io/LimitedInputStreamTest.java b/src/test/java/no/digipost/io/LimitedInputStreamTest.java
@@ -87,6 +87,41 @@ void neverReadsMoreThanTheSetLimit() {
                     }
                 });
         }
+
+        @Test
+        void discardsBytesAfterLimit() throws IOException {
+            byte[] sixBytes = new byte[] {65, 66, 67, 68, 69, 70};
+            try (
+                    InputStream source = new ByteArrayInputStream(sixBytes);
+                    InputStream limitedToTwoBytes = limit(source, bytes(4))) {
+
+
+                byte[] readBytes = new byte[6];
+                byte[] expectedEmpty = new byte[2];
+                assertAll(
+                        () -> assertThat("first read yields 4 bytes", limitedToTwoBytes.read(readBytes), is(4)),
+                        () -> assertArrayEquals(new byte[] {65, 66, 67, 68, 0, 0}, readBytes),
+                        () -> assertThat("reading single read yields EOF", limitedToTwoBytes.read(), is(-1)),
+                        () -> assertThat("reading chunk yields EOF", limitedToTwoBytes.read(expectedEmpty), is(-1)),
+                        () -> assertArrayEquals(new byte[] {0, 0}, expectedEmpty));
+            }
+        }
+
+        @Test
+        void readingZeroBytesYieldsZeroUntilEofHasBeenDetectedByANonZeroReadOperation() throws IOException {
+            byte[] twoBytes = new byte[] {65, 66};
+            try (
+                    InputStream source = new ByteArrayInputStream(twoBytes);
+                    InputStream limitedToOneByte = limit(source, bytes(1))) {
+                assertThat(limitedToOneByte.read(new byte[0]), is(0));
+                limitedToOneByte.read();
+                assertThat(limitedToOneByte.read(new byte[0]), is(0));
+
+                assertThat(limitedToOneByte.read(), is(-1));
+                assertThat(limitedToOneByte.read(new byte[0]), is(-1));
+            }
+        }
+
     }
 
 
@@ -117,6 +152,53 @@ public void wrapsOtherCheckedExceptionsThanIOExceptionAsRuntimeException() throw
             assertThat(assertThrows(RuntimeException.class, () -> testLimitedStream("xyz", () -> tooManyBytes)), where(Exception::getCause, sameInstance(tooManyBytes)));
         }
 
+        @Test
+        void cutsOffLastChunkIfReachingLimit() throws IOException {
+            byte[] sixBytes = new byte[] {65, 66, 67, 68, 69, 70};
+            try (
+                    InputStream source = new ByteArrayInputStream(sixBytes);
+                    InputStream limitedToTwoBytes = limit(source, bytes(4), () -> new IllegalStateException("reached limit!"))) {
+
+                byte[] consumed = new byte[6];
+                assertAll(
+                        () -> assertThat(limitedToTwoBytes.read(consumed), is(4)),
+                        () -> assertArrayEquals(new byte[] {65, 66, 67, 68, 0, 0}, consumed),
+                        () -> assertThrows(IllegalStateException.class, () -> limitedToTwoBytes.read(consumed)),
+                        () -> assertArrayEquals(new byte[] {65, 66, 67, 68, 0, 0}, consumed)
+                        );
+            }
+        }
+
+        @Test
+        void doesNotThrowWhenAttemptingToReadChunkExactlyEndingAtTheLimit() throws IOException {
+            byte[] sixBytes = new byte[] {65, 66, 67, 68, 69, 70};
+            try (
+                    InputStream source = new ByteArrayInputStream(sixBytes);
+                    InputStream limitedToTwoBytes = limit(source, bytes(6), () -> new IllegalStateException("reached limit!"))) {
+
+                limitedToTwoBytes.read(new byte[3]);
+                byte[] lastThreeBytes = new byte[3];
+                limitedToTwoBytes.read(lastThreeBytes);
+                assertArrayEquals(new byte[]{68, 69, 70}, lastThreeBytes);
+            }
+        }
+
+        @Test
+        void readingZeroBytesNeverThrowsUntilTryingToActuallyReadNonZeroBytes() throws IOException {
+            byte[] twoBytes = new byte[] {65, 66};
+            try (
+                    InputStream source = new ByteArrayInputStream(twoBytes);
+                    InputStream limitedToOneByte = limit(source, bytes(1), () -> new IOException("Should not be thrown"))) {
+                assertThat(limitedToOneByte.read(new byte[0]), is(0));
+                limitedToOneByte.read();
+                assertThat(limitedToOneByte.read(new byte[0]), is(0));
+                assertAll(
+                    () -> assertThrows(IOException.class, () -> limitedToOneByte.read()),
+                    () -> assertThrows(IOException.class, () -> limitedToOneByte.read(new byte[0])),
+                    () -> assertThrows(IOException.class, () -> limitedToOneByte.read(new byte[1])));
+            }
+        }
+
     }
 
 
@@ -181,9 +263,10 @@ void doesNotConsumeMoreFromUnderlyingInputStreamThanGivenLimit() throws IOExcept
                 InputStream threeBytes = new ByteArrayInputStream(new byte[] {65, 66, 67});
                 InputStream maxTwoBytes = limit(threeBytes, bytes(2))) {
 
-            assertThat(maxTwoBytes.read(readBytes), is(2));
+            assertAll(
+                    () -> assertThat(maxTwoBytes.read(readBytes), is(2)),
+                    () -> assertArrayEquals(new byte[] {65, 66, 0}, readBytes));
         }
-        assertArrayEquals(new byte[] {65, 66, 0}, readBytes);
     }
 
     @Test
@@ -199,7 +282,7 @@ void ableToResetBufferedStreamWhenLimitedStreamIsExhausted() throws IOException
                 InputStream bufferedSource = new BufferedInputStream(source, 400)) {
 
             bufferedSource.mark(limit);
-            try (InputStream limitedStream = limit(bufferedSource, DataSize.bytes(limit), () -> new IllegalStateException("Reached limit!"))) {
+            try (InputStream limitedStream = limit(bufferedSource, DataSize.bytes(limit))) {
                 assertThat(limitedStream.read(readFromLimitedStream), is(limit));
                 bufferedSource.reset();
                 bufferedSource.read(readFromBufferedStream);
@@ -224,6 +307,11 @@ void rewindWhenReachingLimit() throws IOException {
                 InputStream source = new ByteArrayInputStream(twoKiloByte);
                 InputStream bufferedSource = new BufferedInputStream(source, 512)) {
 
+            // consuming a LimitedInputStream until EOF requires consuming one "extra" byte
+            // to determine if the underlying stream is also at EOF or has more data available
+            // in order to distinguish if it should EOF or throw an exception. It would in theory
+            // be possible to avoid reading this extra byte if the LimitedInputStream is not
+            // set to throw an exception on reaching the limit, but currently it does not support this.
             bufferedSource.mark(limit + 1); //  <-- :(
             try (InputStream limitedStream = limit(bufferedSource, DataSize.bytes(limit))) {
                 byte[] readFromLimitedStream = toByteArray(limitedStream);
@@ -234,7 +322,6 @@ void rewindWhenReachingLimit() throws IOException {
                 assertArrayEquals(readFromBufferedStream, twoKiloByte);
             }
         }
-
     }
 
 }