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;
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;
}
}
};
}

View File

@ -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();
}
}