1 files changed, 123 insertions, 0 deletions
diff --git a/java/core/src/test/java/com/google/protobuf/CodedInputStreamTest.java b/java/core/src/test/java/com/google/protobuf/CodedInputStreamTest.java
index e440c7db..da2c067e 100644
--- a/java/core/src/test/java/com/google/protobuf/CodedInputStreamTest.java
+++ b/java/core/src/test/java/com/google/protobuf/CodedInputStreamTest.java
@@ -41,6 +41,7 @@ import java.io.FilterInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
+import java.util.Arrays;
 import junit.framework.TestCase;
 
 /**
@@ -613,6 +614,82 @@ public class CodedInputStreamTest extends TestCase {
       checkSizeLimitExceeded(expected);
     }
   }
+  
+  public void testRefillBufferWithCorrectSize() throws Exception {
+    // NOTE: refillBuffer only applies to the stream-backed CIS.
+    byte[] bytes = "123456789".getBytes("UTF-8");
+    ByteArrayOutputStream rawOutput = new ByteArrayOutputStream();
+    CodedOutputStream output = CodedOutputStream.newInstance(rawOutput, bytes.length);
+
+    int tag = WireFormat.makeTag(1, WireFormat.WIRETYPE_LENGTH_DELIMITED);
+    output.writeRawVarint32(tag);
+    output.writeRawVarint32(bytes.length);
+    output.writeRawBytes(bytes);
+    output.writeRawVarint32(tag);
+    output.writeRawVarint32(bytes.length);
+    output.writeRawBytes(bytes);
+    output.writeRawByte(4);
+    output.flush();
+
+    // Input is two string with length 9 and one raw byte.
+    byte[] rawInput = rawOutput.toByteArray(); 
+    for (int inputStreamBufferLength = 8; 
+        inputStreamBufferLength <= rawInput.length + 1; inputStreamBufferLength++) {
+      CodedInputStream input = CodedInputStream.newInstance(
+              new ByteArrayInputStream(rawInput), inputStreamBufferLength);
+      input.setSizeLimit(rawInput.length - 1);
+      input.readString();
+      input.readString(); 
+      try {
+        input.readRawByte(); // Hits limit.
+        fail("Should have thrown an exception!");
+      } catch (InvalidProtocolBufferException expected) {
+        checkSizeLimitExceeded(expected);
+      }
+    }
+  }
+
+  public void testIsAtEnd() throws Exception {
+    CodedInputStream input = CodedInputStream.newInstance(
+        new ByteArrayInputStream(new byte[5]));
+    try {   
+      for (int i = 0; i < 5; i++) {
+        assertEquals(false, input.isAtEnd());
+        input.readRawByte();
+      }
+      assertEquals(true, input.isAtEnd());
+    } catch (Exception e) {
+      fail("Catch exception in the testIsAtEnd");
+    }
+  }
+
+  public void testCurrentLimitExceeded() throws Exception {
+    byte[] bytes = "123456789".getBytes("UTF-8");
+    ByteArrayOutputStream rawOutput = new ByteArrayOutputStream();
+    CodedOutputStream output = CodedOutputStream.newInstance(rawOutput, bytes.length);
+
+    int tag = WireFormat.makeTag(1, WireFormat.WIRETYPE_LENGTH_DELIMITED);
+    output.writeRawVarint32(tag);
+    output.writeRawVarint32(bytes.length);
+    output.writeRawBytes(bytes);
+    output.flush();
+
+    byte[] rawInput = rawOutput.toByteArray();
+    CodedInputStream input = CodedInputStream.newInstance(
+              new ByteArrayInputStream(rawInput));
+    // The length of the whole rawInput
+    input.setSizeLimit(11);
+    // Some number that is smaller than the rawInput's length
+    // but larger than 2 
+    input.pushLimit(5);
+    try {
+      input.readString();
+      fail("Should have thrown an exception");
+    } catch (InvalidProtocolBufferException expected) {
+      assertEquals(expected.getMessage(), 
+          InvalidProtocolBufferException.truncatedMessage().getMessage());
+    }
+  }
 
   public void testSizeLimitMultipleMessages() throws Exception {
     // NOTE: Size limit only applies to the stream-backed CIS.
@@ -807,6 +884,52 @@ public class CodedInputStreamTest extends TestCase {
     }
   }
 
+  public void testReadLargeByteStringFromInputStream() throws Exception {
+    byte[] bytes = new byte[1024 * 1024];
+    for (int i = 0; i < bytes.length; i++) {
+      bytes[i] = (byte) (i & 0xFF);
+    }
+    ByteString.Output rawOutput = ByteString.newOutput();
+    CodedOutputStream output = CodedOutputStream.newInstance(rawOutput);
+    output.writeRawVarint32(bytes.length);
+    output.writeRawBytes(bytes);
+    output.flush();
+    byte[] data = rawOutput.toByteString().toByteArray();
+
+    CodedInputStream input = CodedInputStream.newInstance(
+        new ByteArrayInputStream(data) {
+          @Override
+          public synchronized int available() {
+            return 0;
+          }
+        });
+    ByteString result = input.readBytes();
+    assertEquals(ByteString.copyFrom(bytes), result);
+  }
+
+  public void testReadLargeByteArrayFromInputStream() throws Exception {
+    byte[] bytes = new byte[1024 * 1024];
+    for (int i = 0; i < bytes.length; i++) {
+      bytes[i] = (byte) (i & 0xFF);
+    }
+    ByteString.Output rawOutput = ByteString.newOutput();
+    CodedOutputStream output = CodedOutputStream.newInstance(rawOutput);
+    output.writeRawVarint32(bytes.length);
+    output.writeRawBytes(bytes);
+    output.flush();
+    byte[] data = rawOutput.toByteString().toByteArray();
+
+    CodedInputStream input = CodedInputStream.newInstance(
+        new ByteArrayInputStream(data) {
+          @Override
+          public synchronized int available() {
+            return 0;
+          }
+        });
+    byte[] result = input.readByteArray();
+    assertTrue(Arrays.equals(bytes, result));
+  }
+
   public void testReadByteBuffer() throws Exception {
     ByteString.Output rawOutput = ByteString.newOutput();
     CodedOutputStream output = CodedOutputStream.newInstance(rawOutput);