Fix bug in UnsafeContentInputStream where read() corrupts "marked" content

2 files changed, 20 insertions, 5 deletions

diff --git a/jdisc_core/src/main/java/com/yahoo/jdisc/handler/UnsafeContentInputStream.java b/jdisc_core/src/main/java/com/yahoo/jdisc/handler/UnsafeContentInputStream.java
index 4d54c40209c..4af40a22447 100644
--- a/jdisc_core/src/main/java/com/yahoo/jdisc/handler/UnsafeContentInputStream.java
+++ b/jdisc_core/src/main/java/com/yahoo/jdisc/handler/UnsafeContentInputStream.java
@@ -71,10 +71,9 @@ public class UnsafeContentInputStream extends InputStream {
             read += toRead;
         }
         if (marked != null) {
-            if (readSinceMarked + len <= marked.length) {
-                for (int i=0; i < len; i++) {
-                    marked[readSinceMarked++] = buf[off+i];
-                }
+            if (readSinceMarked + read <= marked.length) {
+                System.arraycopy(buf, off, marked, readSinceMarked, read);
+                readSinceMarked += read;
             } else {
                 marked = null;
             }
diff --git a/jdisc_core/src/test/java/com/yahoo/jdisc/handler/UnsafeContentInputStreamTestCase.java b/jdisc_core/src/test/java/com/yahoo/jdisc/handler/UnsafeContentInputStreamTestCase.java
index 134c34641a5..0c3670dd56e 100644
--- a/jdisc_core/src/test/java/com/yahoo/jdisc/handler/UnsafeContentInputStreamTestCase.java
+++ b/jdisc_core/src/test/java/com/yahoo/jdisc/handler/UnsafeContentInputStreamTestCase.java
@@ -10,9 +10,10 @@ import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.nio.ByteBuffer;
 
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 /**
@@ -70,6 +71,21 @@ public class UnsafeContentInputStreamTestCase {
     }
 
     @Test
+    public void requireThatReadAfterResetIncludesDataAfterMark() throws IOException {
+        ReadableContentChannel content = new ReadableContentChannel();
+        UnsafeContentInputStream in = new UnsafeContentInputStream(content);
+        byte[] outBuf = new byte[] {1, 2, 3};
+        content.write(ByteBuffer.wrap(outBuf), null);
+        in.mark(4);
+        assertEquals(3, in.read(new byte[] {101, 102, 103, 104}));
+        in.reset();
+        byte[] inBuf = new byte[4];
+        int read = in.read(inBuf);
+        assertEquals(3, read);
+        assertArrayEquals(new byte[]{1, 2, 3, 0}, inBuf);
+    }
+
+    @Test
     public void requireThatCompletionsAreCalledWithDeprecatedContentWriter() throws IOException {
         BufferedContentChannel channel = new BufferedContentChannel();
         FastContentWriter writer = new FastContentWriter(channel);