Remove proto size restriction when parsing protos.

This commit is contained in:
nmittler 2015-08-19 08:55:32 -07:00
parent 6ebf9b1373
commit 573f79a06d
2 changed files with 31 additions and 1 deletions

View File

@ -31,6 +31,7 @@
package io.grpc.protobuf; package io.grpc.protobuf;
import com.google.protobuf.CodedInputStream;
import com.google.protobuf.InvalidProtocolBufferException; import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.Message; import com.google.protobuf.Message;
import com.google.protobuf.MessageLite; import com.google.protobuf.MessageLite;
@ -73,12 +74,28 @@ public class ProtoUtils {
} }
} }
try { try {
return parser.parseFrom(stream); return parseFrom(stream);
} catch (InvalidProtocolBufferException ipbe) { } catch (InvalidProtocolBufferException ipbe) {
throw Status.INTERNAL.withDescription("Invalid protobuf byte sequence") throw Status.INTERNAL.withDescription("Invalid protobuf byte sequence")
.withCause(ipbe).asRuntimeException(); .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;
}
}
}; };
} }

View File

@ -35,6 +35,7 @@ import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertSame; import static org.junit.Assert.assertSame;
import com.google.common.io.ByteStreams; import com.google.common.io.ByteStreams;
import com.google.protobuf.ByteString;
import com.google.protobuf.Enum; import com.google.protobuf.Enum;
import com.google.protobuf.Type; import com.google.protobuf.Type;
@ -46,6 +47,7 @@ import org.junit.runners.JUnit4;
import java.io.ByteArrayInputStream; import java.io.ByteArrayInputStream;
import java.io.InputStream; import java.io.InputStream;
import java.util.Random;
/** Unit tests for {@link ProtoUtils}. */ /** Unit tests for {@link ProtoUtils}. */
@RunWith(JUnit4.class) @RunWith(JUnit4.class)
@ -72,4 +74,15 @@ public class ProtoUtilsTest {
Enum altProto = Enum.newBuilder().setName(proto.getName()).build(); Enum altProto = Enum.newBuilder().setName(proto.getName()).build();
assertEquals(proto, marshaller.parse(enumMarshaller.stream(altProto))); 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();
}
} }