Add recursion check when parsing unknown fields in Java. (#18388)

* Internal change

PiperOrigin-RevId: 663919912

* Internal change

PiperOrigin-RevId: 653615736

* Add recursion check when parsing unknown fields in Java.

PiperOrigin-RevId: 675657198

---------

Co-authored-by: Protobuf Team Bot <protobuf-github-bot@google.com>

Author: zhangskz

java/core/BUILD.bazel

@@ -608,6 +608,7 @@ junit_tests(
             "src/test/java/com/google/protobuf/DescriptorsTest.java",
             "src/test/java/com/google/protobuf/DebugFormatTest.java",
             "src/test/java/com/google/protobuf/CodedOutputStreamTest.java",
+            "src/test/java/com/google/protobuf/CodedInputStreamTest.java",
             "src/test/java/com/google/protobuf/ProtobufToStringOutputTest.java",
             # Excluded in core_tests
             "src/test/java/com/google/protobuf/DecodeUtf8Test.java",
@@ -656,6 +657,7 @@ junit_tests(
             "src/test/java/com/google/protobuf/DescriptorsTest.java",
             "src/test/java/com/google/protobuf/DebugFormatTest.java",
             "src/test/java/com/google/protobuf/CodedOutputStreamTest.java",
+            "src/test/java/com/google/protobuf/CodedInputStreamTest.java",
             "src/test/java/com/google/protobuf/ProtobufToStringOutputTest.java",
             # Excluded in core_tests
             "src/test/java/com/google/protobuf/DecodeUtf8Test.java",

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

@@ -23,9 +23,12 @@
  */
 @CheckReturnValue
 final class ArrayDecoders {
+  static final int DEFAULT_RECURSION_LIMIT = 100;
 
-  private ArrayDecoders() {
-  }
+  @SuppressWarnings("NonFinalStaticField")
+  private static volatile int recursionLimit = DEFAULT_RECURSION_LIMIT;
+
+  private ArrayDecoders() {}
 
   /**
    * A helper used to return multiple values in a Java function. Java doesn't natively support
@@ -38,6 +41,7 @@ static final class Registers {
     public long long1;
     public Object object1;
     public final ExtensionRegistryLite extensionRegistry;
+    public int recursionDepth;
 
     Registers() {
       this.extensionRegistry = ExtensionRegistryLite.getEmptyRegistry();
@@ -245,7 +249,10 @@ static int mergeMessageField(
     if (length < 0 || length > limit - position) {
       throw InvalidProtocolBufferException.truncatedMessage();
     }
+    registers.recursionDepth++;
+    checkRecursionLimit(registers.recursionDepth);
     schema.mergeFrom(msg, data, position, position + length, registers);
+    registers.recursionDepth--;
     registers.object1 = msg;
     return position + length;
   }
@@ -263,8 +270,11 @@ static int mergeGroupField(
     // A group field must has a MessageSchema (the only other subclass of Schema is MessageSetSchema
     // and it can't be used in group fields).
     final MessageSchema messageSchema = (MessageSchema) schema;
+    registers.recursionDepth++;
+    checkRecursionLimit(registers.recursionDepth);
     final int endPosition =
         messageSchema.parseMessage(msg, data, position, limit, endGroup, registers);
+    registers.recursionDepth--;
     registers.object1 = msg;
     return endPosition;
   }
@@ -1025,6 +1035,8 @@ static int decodeUnknownField(
         final UnknownFieldSetLite child = UnknownFieldSetLite.newInstance();
         final int endGroup = (tag & ~0x7) | WireFormat.WIRETYPE_END_GROUP;
         int lastTag = 0;
+        registers.recursionDepth++;
+        checkRecursionLimit(registers.recursionDepth);
         while (position < limit) {
           position = decodeVarint32(data, position, registers);
           lastTag = registers.int1;
@@ -1033,6 +1045,7 @@ static int decodeUnknownField(
           }
           position = decodeUnknownField(lastTag, data, position, limit, child, registers);
         }
+        registers.recursionDepth--;
         if (position > limit || lastTag != endGroup) {
           throw InvalidProtocolBufferException.parseFailure();
         }
@@ -1079,4 +1092,18 @@ static int skipField(int tag, byte[] data, int position, int limit, Registers re
         throw InvalidProtocolBufferException.invalidTag();
     }
   }
+
+  /**
+   * Set the maximum recursion limit that ArrayDecoders will allow. An exception will be thrown if
+   * the depth of the message exceeds this limit.
+   */
+  public static void setRecursionLimit(int limit) {
+    recursionLimit = limit;
+  }
+
+  private static void checkRecursionLimit(int depth) throws InvalidProtocolBufferException {
+    if (depth >= recursionLimit) {
+      throw InvalidProtocolBufferException.recursionLimitExceeded();
+    }
+  }
 }

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

@@ -224,13 +224,41 @@ public abstract boolean skipField(final int tag, final CodedOutputStream output)
    * Reads and discards an entire message. This will read either until EOF or until an endgroup tag,
    * whichever comes first.
    */
-  public abstract void skipMessage() throws IOException;
+  public void skipMessage() throws IOException {
+    while (true) {
+      final int tag = readTag();
+      if (tag == 0) {
+        return;
+      }
+      checkRecursionLimit();
+      ++recursionDepth;
+      boolean fieldSkipped = skipField(tag);
+      --recursionDepth;
+      if (!fieldSkipped) {
+        return;
+      }
+    }
+  }
 
   /**
    * Reads an entire message and writes it to output in wire format. This will read either until EOF
    * or until an endgroup tag, whichever comes first.
    */
-  public abstract void skipMessage(CodedOutputStream output) throws IOException;
+  public void skipMessage(CodedOutputStream output) throws IOException {
+    while (true) {
+      final int tag = readTag();
+      if (tag == 0) {
+        return;
+      }
+      checkRecursionLimit();
+      ++recursionDepth;
+      boolean fieldSkipped = skipField(tag, output);
+      --recursionDepth;
+      if (!fieldSkipped) {
+        return;
+      }
+    }
+  }
 
   // -----------------------------------------------------------------
 
@@ -700,26 +728,6 @@ public boolean skipField(final int tag, final CodedOutputStream output) throws I
       }
     }
 
-    @Override
-    public void skipMessage() throws IOException {
-      while (true) {
-        final int tag = readTag();
-        if (tag == 0 || !skipField(tag)) {
-          return;
-        }
-      }
-    }
-
-    @Override
-    public void skipMessage(CodedOutputStream output) throws IOException {
-      while (true) {
-        final int tag = readTag();
-        if (tag == 0 || !skipField(tag, output)) {
-          return;
-        }
-      }
-    }
-
     // -----------------------------------------------------------------
 
     @Override
@@ -1412,26 +1420,6 @@ public boolean skipField(final int tag, final CodedOutputStream output) throws I
       }
     }
 
-    @Override
-    public void skipMessage() throws IOException {
-      while (true) {
-        final int tag = readTag();
-        if (tag == 0 || !skipField(tag)) {
-          return;
-        }
-      }
-    }
-
-    @Override
-    public void skipMessage(CodedOutputStream output) throws IOException {
-      while (true) {
-        final int tag = readTag();
-        if (tag == 0 || !skipField(tag, output)) {
-          return;
-        }
-      }
-    }
-
     // -----------------------------------------------------------------
 
     @Override
@@ -2178,26 +2166,6 @@ public boolean skipField(final int tag, final CodedOutputStream output) throws I
       }
     }
 
-    @Override
-    public void skipMessage() throws IOException {
-      while (true) {
-        final int tag = readTag();
-        if (tag == 0 || !skipField(tag)) {
-          return;
-        }
-      }
-    }
-
-    @Override
-    public void skipMessage(CodedOutputStream output) throws IOException {
-      while (true) {
-        final int tag = readTag();
-        if (tag == 0 || !skipField(tag, output)) {
-          return;
-        }
-      }
-    }
-
     /** Collects the bytes skipped and returns the data in a ByteBuffer. */
     private class SkippedDataSink implements RefillCallback {
       private int lastPos = pos;
@@ -3322,26 +3290,6 @@ public boolean skipField(final int tag, final CodedOutputStream output) throws I
       }
     }
 
-    @Override
-    public void skipMessage() throws IOException {
-      while (true) {
-        final int tag = readTag();
-        if (tag == 0 || !skipField(tag)) {
-          return;
-        }
-      }
-    }
-
-    @Override
-    public void skipMessage(CodedOutputStream output) throws IOException {
-      while (true) {
-        final int tag = readTag();
-        if (tag == 0 || !skipField(tag, output)) {
-          return;
-        }
-      }
-    }
-
     // -----------------------------------------------------------------
 
     @Override

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

@@ -132,7 +132,7 @@ public InvalidWireTypeException(String description) {
   static InvalidProtocolBufferException recursionLimitExceeded() {
     return new InvalidProtocolBufferException(
         "Protocol message had too many levels of nesting.  May be malicious.  "
-            + "Use CodedInputStream.setRecursionLimit() to increase the depth limit.");
+            + "Use setRecursionLimit() to increase the recursion depth limit.");
   }
 
   static InvalidProtocolBufferException sizeLimitExceeded() {

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

@@ -3006,7 +3006,8 @@ private <UT, UB, ET extends FieldDescriptorLite<ET>> void mergeFromHelper(
               unknownFields = unknownFieldSchema.getBuilderFromMessage(message);
             }
             // Unknown field.
-            if (unknownFieldSchema.mergeOneFieldFrom(unknownFields, reader)) {
+            if (unknownFieldSchema.mergeOneFieldFrom(
+                unknownFields, reader, /* currentDepth= */ 0)) {
               continue;
             }
           }
@@ -3381,7 +3382,8 @@ private <UT, UB, ET extends FieldDescriptorLite<ET>> void mergeFromHelper(
               if (unknownFields == null) {
                 unknownFields = unknownFieldSchema.getBuilderFromMessage(message);
               }
-              if (!unknownFieldSchema.mergeOneFieldFrom(unknownFields, reader)) {
+              if (!unknownFieldSchema.mergeOneFieldFrom(
+                  unknownFields, reader, /* currentDepth= */ 0)) {
                 return;
               }
               break;
@@ -3397,7 +3399,8 @@ private <UT, UB, ET extends FieldDescriptorLite<ET>> void mergeFromHelper(
             if (unknownFields == null) {
               unknownFields = unknownFieldSchema.getBuilderFromMessage(message);
             }
-            if (!unknownFieldSchema.mergeOneFieldFrom(unknownFields, reader)) {
+            if (!unknownFieldSchema.mergeOneFieldFrom(
+                unknownFields, reader, /* currentDepth= */ 0)) {
               return;
             }
           }

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

@@ -278,7 +278,7 @@ boolean parseMessageSetItemOrUnknownField(
               reader, extension, extensionRegistry, extensions);
           return true;
         } else {
-          return unknownFieldSchema.mergeOneFieldFrom(unknownFields, reader);
+          return unknownFieldSchema.mergeOneFieldFrom(unknownFields, reader, /* currentDepth= */ 0);
         }
       } else {
         return reader.skipField();

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

@@ -13,6 +13,11 @@
 @CheckReturnValue
 abstract class UnknownFieldSchema<T, B> {
 
+  static final int DEFAULT_RECURSION_LIMIT = 100;
+
+  @SuppressWarnings("NonFinalStaticField")
+  private static volatile int recursionLimit = DEFAULT_RECURSION_LIMIT;
+
   /** Whether unknown fields should be dropped. */
   abstract boolean shouldDiscardUnknownFields(Reader reader);
 
@@ -56,7 +61,8 @@ abstract class UnknownFieldSchema<T, B> {
   abstract void makeImmutable(Object message);
 
   /** Merges one field into the unknown fields. */
-  final boolean mergeOneFieldFrom(B unknownFields, Reader reader) throws IOException {
+  final boolean mergeOneFieldFrom(B unknownFields, Reader reader, int currentDepth)
+      throws IOException {
     int tag = reader.getTag();
     int fieldNumber = WireFormat.getTagFieldNumber(tag);
     switch (WireFormat.getTagWireType(tag)) {
@@ -75,7 +81,12 @@ final boolean mergeOneFieldFrom(B unknownFields, Reader reader) throws IOExcepti
       case WireFormat.WIRETYPE_START_GROUP:
         final B subFields = newBuilder();
         int endGroupTag = WireFormat.makeTag(fieldNumber, WireFormat.WIRETYPE_END_GROUP);
-        mergeFrom(subFields, reader);
+        currentDepth++;
+        if (currentDepth >= recursionLimit) {
+          throw InvalidProtocolBufferException.recursionLimitExceeded();
+        }
+        mergeFrom(subFields, reader, currentDepth);
+        currentDepth--;
         if (endGroupTag != reader.getTag()) {
           throw InvalidProtocolBufferException.invalidEndTag();
         }
@@ -88,10 +99,11 @@ final boolean mergeOneFieldFrom(B unknownFields, Reader reader) throws IOExcepti
     }
   }
 
-  final void mergeFrom(B unknownFields, Reader reader) throws IOException {
+  private final void mergeFrom(B unknownFields, Reader reader, int currentDepth)
+      throws IOException {
     while (true) {
       if (reader.getFieldNumber() == Reader.READ_DONE
-          || !mergeOneFieldFrom(unknownFields, reader)) {
+          || !mergeOneFieldFrom(unknownFields, reader, currentDepth)) {
         break;
       }
     }
@@ -108,4 +120,12 @@ final void mergeFrom(B unknownFields, Reader reader) throws IOException {
   abstract int getSerializedSizeAsMessageSet(T message);
 
   abstract int getSerializedSize(T unknowns);
+
+  /**
+   * Set the maximum recursion limit that ArrayDecoders will allow. An exception will be thrown if
+   * the depth of the message exceeds this limit.
+   */
+  public void setRecursionLimit(int limit) {
+    recursionLimit = limit;
+  }
 }