From 573f79a06dbee6e7a346023d3b881c4b874023d8 Mon Sep 17 00:00:00 2001
From: nmittler <nathanmittler@google.com>
Date: Wed, 19 Aug 2015 08:55:32 -0700
Subject: [PATCH] Remove proto size restriction when parsing protos.

---
 .../java/io/grpc/protobuf/ProtoUtils.java     | 19 ++++++++++++++++++-
 .../java/io/grpc/protobuf/ProtoUtilsTest.java | 13 +++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/protobuf/src/main/java/io/grpc/protobuf/ProtoUtils.java b/protobuf/src/main/java/io/grpc/protobuf/ProtoUtils.java
index 6a2cd3ea8b..e2ff57ddb2 100644
--- a/protobuf/src/main/java/io/grpc/protobuf/ProtoUtils.java
+++ b/protobuf/src/main/java/io/grpc/protobuf/ProtoUtils.java
@@ -31,6 +31,7 @@
 
 package io.grpc.protobuf;
 
+import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.Message;
 import com.google.protobuf.MessageLite;
@@ -73,12 +74,28 @@ public class ProtoUtils {
           }
         }
         try {
-          return parser.parseFrom(stream);
+          return parseFrom(stream);
         } catch (InvalidProtocolBufferException ipbe) {
           throw Status.INTERNAL.withDescription("Invalid protobuf byte sequence")
             .withCause(ipbe).asRuntimeException();
         }
       }
+
+      private T parseFrom(InputStream stream) throws InvalidProtocolBufferException {
+        // Pre-create the CodedInputStream so that we can remove the size limit restriction
+        // when parsing.
+        CodedInputStream codedInput = CodedInputStream.newInstance(stream);
+        codedInput.setSizeLimit(Integer.MAX_VALUE);
+
+        T message = parser.parseFrom(codedInput);
+        try {
+          codedInput.checkLastTagWas(0);
+          return message;
+        } catch (InvalidProtocolBufferException e) {
+          e.setUnfinishedMessage(message);
+          throw e;
+        }
+      }
     };
   }
 
diff --git a/protobuf/src/test/java/io/grpc/protobuf/ProtoUtilsTest.java b/protobuf/src/test/java/io/grpc/protobuf/ProtoUtilsTest.java
index 35bc4f5944..c6212f4eec 100644
--- a/protobuf/src/test/java/io/grpc/protobuf/ProtoUtilsTest.java
+++ b/protobuf/src/test/java/io/grpc/protobuf/ProtoUtilsTest.java
@@ -35,6 +35,7 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertSame;
 
 import com.google.common.io.ByteStreams;
+import com.google.protobuf.ByteString;
 import com.google.protobuf.Enum;
 import com.google.protobuf.Type;
 
@@ -46,6 +47,7 @@ import org.junit.runners.JUnit4;
 
 import java.io.ByteArrayInputStream;
 import java.io.InputStream;
+import java.util.Random;
 
 /** Unit tests for {@link ProtoUtils}. */
 @RunWith(JUnit4.class)
@@ -72,4 +74,15 @@ public class ProtoUtilsTest {
     Enum altProto = Enum.newBuilder().setName(proto.getName()).build();
     assertEquals(proto, marshaller.parse(enumMarshaller.stream(altProto)));
   }
+
+  @Test
+  public void marshallerShouldNotLimitProtoSize() throws Exception {
+    byte[] bigName = new byte[100 * 1024 * 1024];
+    new Random().nextBytes(bigName);
+
+    proto = Type.newBuilder().setNameBytes(ByteString.copyFrom(bigName)).build();
+
+    // Just perform a round trip to verify that it works.
+    testRoundtrip();
+  }
 }