Skip to content

Commit 0c18340

Browse files
protobuf-github-botcopybara-github
protobuf-github-bot
authored and
copybara-github
committed
Propagate recursion limit when parsing MessageSet extensions.
PiperOrigin-RevId: 949140156
1 parent c64603b commit 0c18340

9 files changed

Lines changed: 176 additions & 21 deletions

File tree

‎java/core/src/main/java/com/google/protobuf/CodedInputStreamReader.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,11 @@ private CodedInputStreamReader(CodedInputStream input) {
4646
this.input.wrapper = this;
4747
}
4848

49+
/** Returns the remaining recursion budget for this CodedInputStream. */
50+
int getRemainingInputRecursionLimit() {
51+
return input.recursionLimit - input.messageDepth - input.groupDepth;
52+
}
53+
4954
public boolean shouldDiscardUnknownFields() {
5055
return input.shouldDiscardUnknownFields();
5156
}

‎java/core/src/main/java/com/google/protobuf/ExtensionSchema.java‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -66,11 +66,12 @@ abstract void parseLengthPrefixedMessageSetItem(
6666
throws IOException;
6767

6868
/**
69-
* Parses the entire content of a {@link ByteString} as one MessageSet item. Unlike {@link
70-
* #parseLengthPrefixedMessageSetItem}, there isn't a length-prefix.
69+
* Parses the entire content of a {@link CodedInputStream} that wraps a {@link ByteString} as one
70+
* MessageSet item. Unlike {@link #parseLengthPrefixedMessageSetItem}, there isn't a
71+
* length-prefix.
7172
*/
7273
abstract void parseMessageSetItem(
73-
ByteString data,
74+
CodedInputStream input,
7475
Object extension,
7576
ExtensionRegistryLite extensionRegistry,
7677
FieldSet<T> extensions)

‎java/core/src/main/java/com/google/protobuf/ExtensionSchemaLite.java‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -541,7 +541,7 @@ void parseLengthPrefixedMessageSetItem(
541541

542542
@Override
543543
void parseMessageSetItem(
544-
ByteString data,
544+
CodedInputStream input,
545545
Object extensionObject,
546546
ExtensionRegistryLite extensionRegistry,
547547
FieldSet<ExtensionDescriptor> extensions)
@@ -551,8 +551,6 @@ void parseMessageSetItem(
551551

552552
MessageLite.Builder builder = extension.getMessageDefaultInstance().newBuilderForType();
553553

554-
final CodedInputStream input = data.newCodedInput();
555-
556554
builder.mergeFrom(input, extensionRegistry);
557555
extensions.setField(extension.descriptor, builder.buildPartial());
558556
input.checkLastTagWas(0);

‎java/core/src/main/java/com/google/protobuf/GeneratedMessageLite.java‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -872,7 +872,14 @@ private <Message2T extends MessageLite> void mergeMessageSetExtensionFromCodedSt
872872
// Process the raw bytes.
873873
if (rawBytes != null && typeId != 0) { // Zero is not a valid type ID.
874874
if (extension != null) { // We known the type
875-
mergeMessageSetExtensionFromBytes(rawBytes, extensionRegistry, extension);
875+
CodedInputStream subInput = rawBytes.newCodedInput();
876+
int remainingInputRecursionLimit =
877+
input.recursionLimit - input.messageDepth - input.groupDepth - 1;
878+
if (remainingInputRecursionLimit < 0) {
879+
throw InvalidProtocolBufferException.recursionLimitExceeded();
880+
}
881+
subInput.setRecursionLimit(remainingInputRecursionLimit);
882+
mergeMessageSetExtensionFromBytes(subInput, extensionRegistry, extension);
876883
} else { // We don't know how to parse this. Ignore it.
877884
if (rawBytes != null) {
878885
mergeLengthDelimitedField(typeId, rawBytes);
@@ -894,7 +901,7 @@ private void eagerlyMergeMessageSetExtension(
894901
}
895902

896903
private void mergeMessageSetExtensionFromBytes(
897-
ByteString rawBytes,
904+
CodedInputStream input,
898905
ExtensionRegistryLite extensionRegistry,
899906
GeneratedExtension<?, ?> extension)
900907
throws IOException {
@@ -906,7 +913,7 @@ private void mergeMessageSetExtensionFromBytes(
906913
if (subBuilder == null) {
907914
subBuilder = extension.getMessageDefaultInstance().newBuilderForType();
908915
}
909-
subBuilder.mergeFrom(rawBytes, extensionRegistry);
916+
subBuilder.mergeFrom(input, extensionRegistry);
910917
MessageLite value = subBuilder.build();
911918

912919
ensureExtensionsAreMutable()

‎java/core/src/main/java/com/google/protobuf/MessageReflection.java‎

Lines changed: 43 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -294,7 +294,8 @@ Object parseMessage(
294294
CodedInputStream input,
295295
ExtensionRegistryLite registry,
296296
Descriptors.FieldDescriptor descriptor,
297-
Message defaultInstance)
297+
Message defaultInstance,
298+
boolean inputLengthPrefixed)
298299
throws IOException;
299300

300301
/**
@@ -494,7 +495,8 @@ public Object parseMessage(
494495
CodedInputStream input,
495496
ExtensionRegistryLite extensionRegistry,
496497
Descriptors.FieldDescriptor field,
497-
Message defaultInstance)
498+
Message defaultInstance,
499+
boolean inputLengthPrefixedUnused)
498500
throws IOException {
499501
Message.Builder subBuilder;
500502
// When default instance is not null. The field is an extension field.
@@ -762,7 +764,8 @@ public Object parseMessage(
762764
CodedInputStream input,
763765
ExtensionRegistryLite registry,
764766
Descriptors.FieldDescriptor field,
765-
Message defaultInstance)
767+
Message defaultInstance,
768+
boolean inputLengthPrefixed)
766769
throws IOException {
767770
Message.Builder subBuilder = defaultInstance.newBuilderForType();
768771
if (!field.isRepeated()) {
@@ -771,7 +774,11 @@ public Object parseMessage(
771774
subBuilder.mergeFrom(originalMessage);
772775
}
773776
}
774-
input.readMessage(subBuilder, registry);
777+
if (inputLengthPrefixed) {
778+
input.readMessage(subBuilder, registry);
779+
} else {
780+
subBuilder.mergeFrom(input, registry);
781+
}
775782
return subBuilder.buildPartial();
776783
}
777784

@@ -977,7 +984,8 @@ public Object parseMessage(
977984
CodedInputStream input,
978985
ExtensionRegistryLite registry,
979986
Descriptors.FieldDescriptor field,
980-
Message defaultInstance)
987+
Message defaultInstance,
988+
boolean inputLengthPrefixed)
981989
throws IOException {
982990
Message.Builder subBuilder = defaultInstance.newBuilderForType();
983991
if (!field.isRepeated()) {
@@ -986,7 +994,11 @@ public Object parseMessage(
986994
subBuilder.mergeFrom(originalMessage);
987995
}
988996
}
989-
input.readMessage(subBuilder, registry);
997+
if (inputLengthPrefixed) {
998+
input.readMessage(subBuilder, registry);
999+
} else {
1000+
subBuilder.mergeFrom(input, registry);
1001+
}
9901002
return subBuilder.buildPartial();
9911003
}
9921004

@@ -1366,10 +1378,17 @@ private static void mergeMessageSetExtensionFromCodedStream(
13661378
}
13671379
input.checkLastTagWas(WireFormat.MESSAGE_SET_ITEM_END_TAG);
13681380

1381+
int remainingInputRecursionLimit =
1382+
input.recursionLimit - input.messageDepth - input.groupDepth - 1;
1383+
if (remainingInputRecursionLimit < 0) {
1384+
throw InvalidProtocolBufferException.recursionLimitExceeded();
1385+
}
1386+
13691387
// Process the raw bytes.
13701388
if (rawBytes != null && typeId != 0) { // Zero is not a valid type ID.
13711389
if (extension != null) { // We known the type
1372-
mergeMessageSetExtensionFromBytes(rawBytes, extension, extensionRegistry, target);
1390+
mergeMessageSetExtensionFromBytes(
1391+
rawBytes, extension, extensionRegistry, target, remainingInputRecursionLimit);
13731392
} else { // We don't know how to parse this. Ignore it.
13741393
if (rawBytes != null && unknownFields != null) {
13751394
unknownFields.mergeField(
@@ -1383,17 +1402,24 @@ private static void mergeMessageSetExtensionFromBytes(
13831402
ByteString rawBytes,
13841403
ExtensionRegistry.ExtensionInfo extension,
13851404
ExtensionRegistryLite extensionRegistry,
1386-
MergeTarget target)
1405+
MergeTarget target,
1406+
int remainingInputRecursionLimit)
13871407
throws IOException {
13881408

13891409
Descriptors.FieldDescriptor field = extension.descriptor;
13901410
boolean hasOriginalValue = target.hasField(field);
13911411

13921412
if (hasOriginalValue || ExtensionRegistryLite.isEagerlyParseMessageSets()) {
13931413
// If the field already exists, we just parse the field.
1414+
CodedInputStream input = rawBytes.newCodedInput();
1415+
input.setRecursionLimit(remainingInputRecursionLimit);
13941416
Object value =
1395-
target.parseMessageFromBytes(
1396-
rawBytes, extensionRegistry, field, extension.defaultInstance);
1417+
target.parseMessage(
1418+
input,
1419+
extensionRegistry,
1420+
field,
1421+
extension.defaultInstance,
1422+
/* inputLengthPrefixed= */ false);
13971423
target.setField(field, value);
13981424
} else {
13991425
// Use InternalLazyField to load MessageSet lazily.
@@ -1410,7 +1436,13 @@ private static void eagerlyMergeMessageSetExtension(
14101436
MergeTarget target)
14111437
throws IOException {
14121438
Descriptors.FieldDescriptor field = extension.descriptor;
1413-
Object value = target.parseMessage(input, extensionRegistry, field, extension.defaultInstance);
1439+
Object value =
1440+
target.parseMessage(
1441+
input,
1442+
extensionRegistry,
1443+
field,
1444+
extension.defaultInstance,
1445+
/* inputLengthPrefixed= */ true);
14141446
target.setField(field, value);
14151447
}
14161448
}

‎java/core/src/main/java/com/google/protobuf/MessageSetSchema.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -360,7 +360,15 @@ boolean parseMessageSetItemOrUnknownField(
360360
if (extension != null) { // We known the type
361361
// TODO: Instead of reading into a temporary ByteString, maybe there is a way
362362
// to read directly from CodedInputStreamReader to the submessage?
363-
extensionSchema.parseMessageSetItem(rawBytes, extension, extensionRegistry, extensions);
363+
int remainingInputRecursionLimit = reader.getRemainingInputRecursionLimit() - 1;
364+
if (remainingInputRecursionLimit < 0) {
365+
throw InvalidProtocolBufferException.recursionLimitExceeded();
366+
}
367+
CodedInputStream rawBytesInput = rawBytes.newCodedInput();
368+
rawBytesInput.setRecursionLimit(remainingInputRecursionLimit);
369+
370+
extensionSchema.parseMessageSetItem(
371+
rawBytesInput, extension, extensionRegistry, extensions);
364372
} else {
365373
unknownFieldSchema.addLengthDelimited(unknownFields, typeId, rawBytes);
366374
}

‎java/core/src/test/java/com/google/protobuf/LazilyParsedMessageSetTest.java‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,9 @@
1111
import static org.junit.Assert.assertThrows;
1212

1313
import com.google.protobuf.ExtensionRegistryLite.LazyExtensionMode;
14+
import proto2_unittest.UnittestMset;
1415
import proto2_unittest.UnittestMset.RawMessageSet;
16+
import proto2_unittest.UnittestMset.RecursiveMessageSet;
1517
import proto2_unittest.UnittestMset.TestMessageSetExtension1;
1618
import proto2_unittest.UnittestMset.TestMessageSetExtension2;
1719
import proto2_unittest.UnittestMset.TestMessageSetExtension3;
@@ -227,4 +229,57 @@ public void testLoadCorruptedLazyField_getSerializedSize() throws Exception {
227229
assertThat(container)
228230
.isEqualTo(TestMessageSetWireFormatContainer.parseFrom(bytes, extensionRegistry));
229231
}
232+
233+
private ByteString makeRecursivePayload(int depth) throws Exception {
234+
ByteString payload = ByteString.copyFromUtf8("payload");
235+
for (int i = 0; i < depth; ++i) {
236+
ByteString.Output out = ByteString.newOutput();
237+
CodedOutputStream cout = CodedOutputStream.newInstance(out);
238+
239+
// Item 1: normal order, ID first
240+
cout.writeTag(1, WireFormat.WIRETYPE_START_GROUP);
241+
cout.writeUInt32(2, 100);
242+
cout.writeBytes(3, ByteString.EMPTY);
243+
cout.writeTag(1, WireFormat.WIRETYPE_END_GROUP);
244+
245+
// Item 2: reversed order, payload first
246+
cout.writeTag(1, WireFormat.WIRETYPE_START_GROUP);
247+
cout.writeBytes(3, payload);
248+
cout.writeUInt32(2, 100);
249+
cout.writeTag(1, WireFormat.WIRETYPE_END_GROUP);
250+
251+
cout.flush();
252+
payload = out.toByteString();
253+
}
254+
return payload;
255+
}
256+
257+
private void testMessageSetRecursionLimit(boolean eager) throws Exception {
258+
boolean originalEagerlyParse = ExtensionRegistryLite.isEagerlyParseMessageSets();
259+
ExtensionRegistryLite.setEagerlyParseMessageSets(eager);
260+
try {
261+
ByteString payload = makeRecursivePayload(2000);
262+
263+
ExtensionRegistry registry = ExtensionRegistry.newInstance();
264+
registry.add(UnittestMset.messageSetExtension);
265+
266+
Throwable exception =
267+
assertThrows(
268+
InvalidProtocolBufferException.class,
269+
() -> RecursiveMessageSet.parseFrom(payload, registry));
270+
assertThat(exception).hasMessageThat().contains("too many levels of nesting");
271+
} finally {
272+
ExtensionRegistryLite.setEagerlyParseMessageSets(originalEagerlyParse);
273+
}
274+
}
275+
276+
@Test
277+
public void testRecursionLimit_eagerMessageSet() throws Exception {
278+
testMessageSetRecursionLimit(/* eager= */ true);
279+
}
280+
281+
@Test
282+
public void testRecursionLimit_lazyMessageSet() throws Exception {
283+
testMessageSetRecursionLimit(/* eager= */ false);
284+
}
230285
}

‎java/lite/src/test/java/com/google/protobuf/LiteTest.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@
3636
import map_lite_test.MapTestProto.TestMap.MessageValue;
3737
import proto2_unittest.NestedExtensionLite;
3838
import proto2_unittest.NonNestedExtensionLite;
39+
import proto2_unittest.UnittestMset;
40+
import proto2_unittest.UnittestMset.RecursiveMessageSet;
3941
import proto2_unittest.UnittestProto.TestOneof2;
4042
import proto2_unittest.lite_equals_and_hash.LiteEqualsAndHash.Bar;
4143
import proto2_unittest.lite_equals_and_hash.LiteEqualsAndHash.BarPrime;
@@ -3205,4 +3207,41 @@ public void testLazyParserRegistration() throws Exception {
32053207
assertThat(after).isInstanceOf(Parser.class);
32063208
assertThat(after).isSameInstanceAs(parser);
32073209
}
3210+
3211+
private ByteString makeRecursivePayload(int depth) throws IOException {
3212+
ByteString payload = ByteString.copyFromUtf8("payload");
3213+
for (int i = 0; i < depth; ++i) {
3214+
ByteString.Output out = ByteString.newOutput();
3215+
CodedOutputStream cout = CodedOutputStream.newInstance(out);
3216+
3217+
// Item 1: normal order, ID first
3218+
cout.writeTag(1, WireFormat.WIRETYPE_START_GROUP);
3219+
cout.writeUInt32(2, 100);
3220+
cout.writeBytes(3, ByteString.EMPTY);
3221+
cout.writeTag(1, WireFormat.WIRETYPE_END_GROUP);
3222+
3223+
// Item 2: reversed order, payload first
3224+
cout.writeTag(1, WireFormat.WIRETYPE_START_GROUP);
3225+
cout.writeBytes(3, payload);
3226+
cout.writeUInt32(2, 100);
3227+
cout.writeTag(1, WireFormat.WIRETYPE_END_GROUP);
3228+
3229+
cout.flush();
3230+
payload = out.toByteString();
3231+
}
3232+
return payload;
3233+
}
3234+
3235+
@Test
3236+
public void testMessageSetRecursionLimit() throws Exception {
3237+
ByteString payload = makeRecursivePayload(2000);
3238+
ExtensionRegistryLite registry = ExtensionRegistryLite.newInstance();
3239+
registry.add(UnittestMset.messageSetExtension);
3240+
3241+
Throwable exception =
3242+
assertThrows(
3243+
InvalidProtocolBufferException.class,
3244+
() -> RecursiveMessageSet.parseFrom(payload, registry));
3245+
assertThat(exception).hasMessageThat().contains("too many levels of nesting");
3246+
}
32083247
}

‎src/google/protobuf/unittest_mset.proto‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,3 +79,13 @@ message RawMessageSet {
7979
required bytes message = 3;
8080
}
8181
}
82+
83+
message RecursiveMessageSet {
84+
option message_set_wire_format = true;
85+
86+
extensions 100 to max;
87+
}
88+
89+
extend RecursiveMessageSet {
90+
optional RecursiveMessageSet message_set_extension = 100;
91+
}

0 commit comments

Comments
 (0)