Propagate recursion limit when parsing MessageSet extensions.

PiperOrigin-RevId: 949140156

Author: protobuf-github-bot

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

@@ -46,6 +46,11 @@ private CodedInputStreamReader(CodedInputStream input) {
     this.input.wrapper = this;
   }
 
+  /** Returns the remaining recursion budget for this CodedInputStream. */
+  int getRemainingInputRecursionLimit() {
+    return input.recursionLimit - input.messageDepth - input.groupDepth;
+  }
+
   public boolean shouldDiscardUnknownFields() {
     return input.shouldDiscardUnknownFields();
   }

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

@@ -66,11 +66,12 @@ abstract void parseLengthPrefixedMessageSetItem(
       throws IOException;
 
   /**
-   * Parses the entire content of a {@link ByteString} as one MessageSet item. Unlike {@link
-   * #parseLengthPrefixedMessageSetItem}, there isn't a length-prefix.
+   * Parses the entire content of a {@link CodedInputStream} that wraps a {@link ByteString} as one
+   * MessageSet item. Unlike {@link #parseLengthPrefixedMessageSetItem}, there isn't a
+   * length-prefix.
    */
   abstract void parseMessageSetItem(
-      ByteString data,
+      CodedInputStream input,
       Object extension,
       ExtensionRegistryLite extensionRegistry,
       FieldSet<T> extensions)

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

@@ -541,7 +541,7 @@ void parseLengthPrefixedMessageSetItem(
 
   @Override
   void parseMessageSetItem(
-      ByteString data,
+      CodedInputStream input,
       Object extensionObject,
       ExtensionRegistryLite extensionRegistry,
       FieldSet<ExtensionDescriptor> extensions)
@@ -551,8 +551,6 @@ void parseMessageSetItem(
 
     MessageLite.Builder builder = extension.getMessageDefaultInstance().newBuilderForType();
 
-    final CodedInputStream input = data.newCodedInput();
-
     builder.mergeFrom(input, extensionRegistry);
     extensions.setField(extension.descriptor, builder.buildPartial());
     input.checkLastTagWas(0);

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

@@ -872,7 +872,14 @@ private <Message2T extends MessageLite> void mergeMessageSetExtensionFromCodedSt
       // Process the raw bytes.
       if (rawBytes != null && typeId != 0) { // Zero is not a valid type ID.
         if (extension != null) { // We known the type
-          mergeMessageSetExtensionFromBytes(rawBytes, extensionRegistry, extension);
+          CodedInputStream subInput = rawBytes.newCodedInput();
+          int remainingInputRecursionLimit =
+              input.recursionLimit - input.messageDepth - input.groupDepth - 1;
+          if (remainingInputRecursionLimit < 0) {
+            throw InvalidProtocolBufferException.recursionLimitExceeded();
+          }
+          subInput.setRecursionLimit(remainingInputRecursionLimit);
+          mergeMessageSetExtension(subInput, extensionRegistry, extension);
         } else { // We don't know how to parse this. Ignore it.
           if (rawBytes != null) {
             mergeLengthDelimitedField(typeId, rawBytes);
@@ -893,8 +900,8 @@ private void eagerlyMergeMessageSetExtension(
       boolean unused = parseExtension(input, extensionRegistry, extension, tag, fieldNumber);
     }
 
-    private void mergeMessageSetExtensionFromBytes(
-        ByteString rawBytes,
+    private void mergeMessageSetExtension(
+        CodedInputStream input,
         ExtensionRegistryLite extensionRegistry,
         GeneratedExtension<?, ?> extension)
         throws IOException {
@@ -906,7 +913,7 @@ private void mergeMessageSetExtensionFromBytes(
       if (subBuilder == null) {
         subBuilder = extension.getMessageDefaultInstance().newBuilderForType();
       }
-      subBuilder.mergeFrom(rawBytes, extensionRegistry);
+      subBuilder.mergeFrom(input, extensionRegistry);
       MessageLite value = subBuilder.build();
 
       ensureExtensionsAreMutable()

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

@@ -294,19 +294,8 @@ Object parseMessage(
         CodedInputStream input,
         ExtensionRegistryLite registry,
         Descriptors.FieldDescriptor descriptor,
-        Message defaultInstance)
-        throws IOException;
-
-    /**
-     * Parse from a ByteString into a sub field message defined based on either FieldDescriptor or
-     * the default instance. There isn't a varint indicating the length of the message at the
-     * beginning of the input ByteString.
-     */
-    Object parseMessageFromBytes(
-        ByteString bytes,
-        ExtensionRegistryLite registry,
-        Descriptors.FieldDescriptor descriptor,
-        Message defaultInstance)
+        Message defaultInstance,
+        boolean inputLengthPrefixed)
         throws IOException;
 
     /**
@@ -494,7 +483,8 @@ public Object parseMessage(
         CodedInputStream input,
         ExtensionRegistryLite extensionRegistry,
         Descriptors.FieldDescriptor field,
-        Message defaultInstance)
+        Message defaultInstance,
+        boolean inputLengthPrefixed)
         throws IOException {
       Message.Builder subBuilder;
       // When default instance is not null. The field is an extension field.
@@ -509,31 +499,11 @@ public Object parseMessage(
           subBuilder.mergeFrom(originalMessage);
         }
       }
-      input.readMessage(subBuilder, extensionRegistry);
-      return subBuilder.buildPartial();
-    }
-
-    @Override
-    public Object parseMessageFromBytes(
-        ByteString bytes,
-        ExtensionRegistryLite extensionRegistry,
-        Descriptors.FieldDescriptor field,
-        Message defaultInstance)
-        throws IOException {
-      Message.Builder subBuilder;
-      // When default instance is not null. The field is an extension field.
-      if (defaultInstance != null) {
-        subBuilder = defaultInstance.newBuilderForType();
+      if (inputLengthPrefixed) {
+        input.readMessage(subBuilder, extensionRegistry);
       } else {
-        subBuilder = builder.newBuilderForField(field);
+        subBuilder.mergeFrom(input, extensionRegistry);
       }
-      if (!field.isRepeated()) {
-        Message originalMessage = (Message) getField(field);
-        if (originalMessage != null) {
-          subBuilder.mergeFrom(originalMessage);
-        }
-      }
-      subBuilder.mergeFrom(bytes, extensionRegistry);
       return subBuilder.buildPartial();
     }
 
@@ -762,7 +732,8 @@ public Object parseMessage(
         CodedInputStream input,
         ExtensionRegistryLite registry,
         Descriptors.FieldDescriptor field,
-        Message defaultInstance)
+        Message defaultInstance,
+        boolean inputLengthPrefixed)
         throws IOException {
       Message.Builder subBuilder = defaultInstance.newBuilderForType();
       if (!field.isRepeated()) {
@@ -771,7 +742,11 @@ public Object parseMessage(
           subBuilder.mergeFrom(originalMessage);
         }
       }
-      input.readMessage(subBuilder, registry);
+      if (inputLengthPrefixed) {
+        input.readMessage(subBuilder, registry);
+      } else {
+        subBuilder.mergeFrom(input, registry);
+      }
       return subBuilder.buildPartial();
     }
 
@@ -823,24 +798,6 @@ public void mergeMessage(
       }
     }
 
-    @Override
-    public Object parseMessageFromBytes(
-        ByteString bytes,
-        ExtensionRegistryLite registry,
-        Descriptors.FieldDescriptor field,
-        Message defaultInstance)
-        throws IOException {
-      Message.Builder subBuilder = defaultInstance.newBuilderForType();
-      if (!field.isRepeated()) {
-        Message originalMessage = (Message) getField(field);
-        if (originalMessage != null) {
-          subBuilder.mergeFrom(originalMessage);
-        }
-      }
-      subBuilder.mergeFrom(bytes, registry);
-      return subBuilder.buildPartial();
-    }
-
     @Override
     public MergeTarget newMergeTargetForField(
         Descriptors.FieldDescriptor descriptor, Message defaultInstance) {
@@ -977,7 +934,8 @@ public Object parseMessage(
         CodedInputStream input,
         ExtensionRegistryLite registry,
         Descriptors.FieldDescriptor field,
-        Message defaultInstance)
+        Message defaultInstance,
+        boolean inputLengthPrefixed)
         throws IOException {
       Message.Builder subBuilder = defaultInstance.newBuilderForType();
       if (!field.isRepeated()) {
@@ -986,7 +944,11 @@ public Object parseMessage(
           subBuilder.mergeFrom(originalMessage);
         }
       }
-      input.readMessage(subBuilder, registry);
+      if (inputLengthPrefixed) {
+        input.readMessage(subBuilder, registry);
+      } else {
+        subBuilder.mergeFrom(input, registry);
+      }
       return subBuilder.buildPartial();
     }
 
@@ -1067,24 +1029,6 @@ public void mergeMessage(
       }
     }
 
-    @Override
-    public Object parseMessageFromBytes(
-        ByteString bytes,
-        ExtensionRegistryLite registry,
-        Descriptors.FieldDescriptor field,
-        Message defaultInstance)
-        throws IOException {
-      Message.Builder subBuilder = defaultInstance.newBuilderForType();
-      if (!field.isRepeated()) {
-        Message originalMessage = (Message) getField(field);
-        if (originalMessage != null) {
-          subBuilder.mergeFrom(originalMessage);
-        }
-      }
-      subBuilder.mergeFrom(bytes, registry);
-      return subBuilder.buildPartial();
-    }
-
     @Override
     public MergeTarget newMergeTargetForField(
         Descriptors.FieldDescriptor descriptor, Message defaultInstance) {
@@ -1366,10 +1310,17 @@ private static void mergeMessageSetExtensionFromCodedStream(
     }
     input.checkLastTagWas(WireFormat.MESSAGE_SET_ITEM_END_TAG);
 
+    int remainingInputRecursionLimit =
+        input.recursionLimit - input.messageDepth - input.groupDepth - 1;
+    if (remainingInputRecursionLimit < 0) {
+      throw InvalidProtocolBufferException.recursionLimitExceeded();
+    }
+
     // Process the raw bytes.
     if (rawBytes != null && typeId != 0) { // Zero is not a valid type ID.
       if (extension != null) { // We known the type
-        mergeMessageSetExtensionFromBytes(rawBytes, extension, extensionRegistry, target);
+        mergeMessageSetExtensionFromBytes(
+            rawBytes, extension, extensionRegistry, target, remainingInputRecursionLimit);
       } else { // We don't know how to parse this. Ignore it.
         if (rawBytes != null && unknownFields != null) {
           unknownFields.mergeField(
@@ -1383,17 +1334,26 @@ private static void mergeMessageSetExtensionFromBytes(
       ByteString rawBytes,
       ExtensionRegistry.ExtensionInfo extension,
       ExtensionRegistryLite extensionRegistry,
-      MergeTarget target)
+      MergeTarget target,
+      int remainingInputRecursionLimit)
       throws IOException {
 
     Descriptors.FieldDescriptor field = extension.descriptor;
     boolean hasOriginalValue = target.hasField(field);
 
+    // TODO: b/535485217 - Consider letting InternalLazyField to concat the bytes to avoid parsing
+    // if the existing field is a lazy field with lazy bytes.
     if (hasOriginalValue || ExtensionRegistryLite.isEagerlyParseMessageSets()) {
       // If the field already exists, we just parse the field.
+      CodedInputStream input = rawBytes.newCodedInput();
+      input.setRecursionLimit(remainingInputRecursionLimit);
       Object value =
-          target.parseMessageFromBytes(
-              rawBytes, extensionRegistry, field, extension.defaultInstance);
+          target.parseMessage(
+              input,
+              extensionRegistry,
+              field,
+              extension.defaultInstance,
+              /* inputLengthPrefixed= */ false);
       target.setField(field, value);
     } else {
       // Use InternalLazyField to load MessageSet lazily.
@@ -1410,7 +1370,13 @@ private static void eagerlyMergeMessageSetExtension(
       MergeTarget target)
       throws IOException {
     Descriptors.FieldDescriptor field = extension.descriptor;
-    Object value = target.parseMessage(input, extensionRegistry, field, extension.defaultInstance);
+    Object value =
+        target.parseMessage(
+            input,
+            extensionRegistry,
+            field,
+            extension.defaultInstance,
+            /* inputLengthPrefixed= */ true);
     target.setField(field, value);
   }
 }

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

@@ -360,7 +360,15 @@ boolean parseMessageSetItemOrUnknownField(
       if (extension != null) { // We known the type
         // TODO: Instead of reading into a temporary ByteString, maybe there is a way
         // to read directly from CodedInputStreamReader to the submessage?
-        extensionSchema.parseMessageSetItem(rawBytes, extension, extensionRegistry, extensions);
+        int remainingInputRecursionLimit = reader.getRemainingInputRecursionLimit() - 1;
+        if (remainingInputRecursionLimit < 0) {
+          throw InvalidProtocolBufferException.recursionLimitExceeded();
+        }
+        CodedInputStream rawBytesInput = rawBytes.newCodedInput();
+        rawBytesInput.setRecursionLimit(remainingInputRecursionLimit);
+
+        extensionSchema.parseMessageSetItem(
+            rawBytesInput, extension, extensionRegistry, extensions);
       } else {
         unknownFieldSchema.addLengthDelimited(unknownFields, typeId, rawBytes);
       }

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

@@ -227,4 +227,64 @@ public void testLoadCorruptedLazyField_getSerializedSize() throws Exception {
     assertThat(container)
         .isEqualTo(TestMessageSetWireFormatContainer.parseFrom(bytes, extensionRegistry));
   }
+
+  private ByteString makeRecursivePayload(int depth) throws Exception {
+    ByteString payload = ByteString.copyFromUtf8("payload");
+    for (int i = 0; i < depth; ++i) {
+      ByteString.Output extOut = ByteString.newOutput();
+      CodedOutputStream extCout = CodedOutputStream.newInstance(extOut);
+      // Write to the `recursive` field.
+      extCout.writeBytes(16, payload);
+      extCout.flush();
+      ByteString extPayload = extOut.toByteString();
+
+      ByteString.Output out = ByteString.newOutput();
+      CodedOutputStream cout = CodedOutputStream.newInstance(out);
+
+      // Item 1: normal order, ID first
+      cout.writeTag(1, WireFormat.WIRETYPE_START_GROUP);
+      cout.writeUInt32(2, TestMessageSetExtension1.messageSetExtension.getNumber());
+      cout.writeBytes(3, ByteString.EMPTY);
+      cout.writeTag(1, WireFormat.WIRETYPE_END_GROUP);
+
+      // Item 2: reversed order, payload first
+      cout.writeTag(1, WireFormat.WIRETYPE_START_GROUP);
+      cout.writeBytes(3, extPayload);
+      cout.writeUInt32(2, TestMessageSetExtension1.messageSetExtension.getNumber());
+      cout.writeTag(1, WireFormat.WIRETYPE_END_GROUP);
+
+      cout.flush();
+      payload = out.toByteString();
+    }
+    return payload;
+  }
+
+  private void testMessageSetRecursionLimit(boolean eager) throws Exception {
+    boolean originalEagerlyParse = ExtensionRegistryLite.isEagerlyParseMessageSets();
+    ExtensionRegistryLite.setEagerlyParseMessageSets(eager);
+    try {
+      ByteString payload = makeRecursivePayload(2000);
+
+      ExtensionRegistry registry = ExtensionRegistry.newInstance();
+      registry.add(TestMessageSetExtension1.messageSetExtension);
+
+      Throwable exception =
+          assertThrows(
+              InvalidProtocolBufferException.class,
+              () -> TestMessageSet.parseFrom(payload, registry));
+      assertThat(exception).hasMessageThat().contains("too many levels of nesting");
+    } finally {
+      ExtensionRegistryLite.setEagerlyParseMessageSets(originalEagerlyParse);
+    }
+  }
+
+  @Test
+  public void testRecursionLimit_eagerMessageSet() throws Exception {
+    testMessageSetRecursionLimit(/* eager= */ true);
+  }
+
+  @Test
+  public void testRecursionLimit_lazyMessageSet() throws Exception {
+    testMessageSetRecursionLimit(/* eager= */ false);
+  }
 }