diff options
Diffstat (limited to 'protobuf-lite')
-rw-r--r-- | protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java | 17 | ||||
-rw-r--r-- | protobuf-lite/src/test/java/io/grpc/protobuf/lite/ProtoLiteUtilsTest.java | 33 |
2 files changed, 45 insertions, 5 deletions
diff --git a/protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java b/protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java index 891420134..2a88a3f89 100644 --- a/protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java +++ b/protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java @@ -128,12 +128,19 @@ public class ProtoLiteUtils { buf = new byte[size]; bufs.set(new WeakReference<byte[]>(buf)); } - int chunkSize; - int position = 0; - while ((chunkSize = stream.read(buf, position, size - position)) != -1) { - position += chunkSize; + + int remaining = size; + while (remaining > 0) { + int position = size - remaining; + int count = stream.read(buf, position, remaining); + if (count == -1) { + break; + } + remaining -= count; } - if (size != position) { + + if (remaining != 0) { + int position = size - remaining; throw new RuntimeException("size inaccurate: " + size + " != " + position); } cis = CodedInputStream.newInstance(buf, 0, size); diff --git a/protobuf-lite/src/test/java/io/grpc/protobuf/lite/ProtoLiteUtilsTest.java b/protobuf-lite/src/test/java/io/grpc/protobuf/lite/ProtoLiteUtilsTest.java index d47e089a0..c4a4d3217 100644 --- a/protobuf-lite/src/test/java/io/grpc/protobuf/lite/ProtoLiteUtilsTest.java +++ b/protobuf-lite/src/test/java/io/grpc/protobuf/lite/ProtoLiteUtilsTest.java @@ -29,6 +29,7 @@ import com.google.protobuf.Enum; import com.google.protobuf.InvalidProtocolBufferException; import com.google.protobuf.Type; import io.grpc.Drainable; +import io.grpc.KnownLength; import io.grpc.Metadata; import io.grpc.MethodDescriptor.Marshaller; import io.grpc.MethodDescriptor.PrototypeMarshaller; @@ -215,4 +216,36 @@ public class ProtoLiteUtilsTest { ProtoLiteUtils.setExtensionRegistry(null); } + + @Test + public void parseFromKnowLengthInputStream() throws Exception { + Marshaller<Type> marshaller = ProtoLiteUtils.marshaller(Type.getDefaultInstance()); + Type expect = Type.newBuilder().setName("expected name").build(); + + Type result = marshaller.parse(new CustomKnownLengthInputStream(expect.toByteArray())); + assertEquals(expect, result); + } + + private static class CustomKnownLengthInputStream extends InputStream implements KnownLength { + private int position = 0; + private byte[] source; + + private CustomKnownLengthInputStream(byte[] source) { + this.source = source; + } + + @Override + public int available() throws IOException { + return source.length - position; + } + + @Override + public int read() throws IOException { + if (position == source.length) { + return -1; + } + + return source[position++]; + } + } } |