Merge pull request #21858 from shaod2/py-cp-29

Cherrypick Pure Python recursion limit enforcement to 29.x

Author: shaod2

python/google/protobuf/internal/decoder.py

@@ -184,7 +184,10 @@ def SpecificDecoder(field_number, is_repeated, is_packed, key, new_default,
                       clear_if_default=False):
     if is_packed:
       local_DecodeVarint = _DecodeVarint
-      def DecodePackedField(buffer, pos, end, message, field_dict):
+      def DecodePackedField(
+          buffer, pos, end, message, field_dict, current_depth=0
+      ):
+        del current_depth  # unused
         value = field_dict.get(key)
         if value is None:
           value = field_dict.setdefault(key, new_default(message))
@@ -199,11 +202,15 @@ def DecodePackedField(buffer, pos, end, message, field_dict):
           del value[-1]   # Discard corrupt value.
           raise _DecodeError('Packed element was truncated.')
         return pos
+
       return DecodePackedField
     elif is_repeated:
       tag_bytes = encoder.TagBytes(field_number, wire_type)
       tag_len = len(tag_bytes)
-      def DecodeRepeatedField(buffer, pos, end, message, field_dict):
+      def DecodeRepeatedField(
+          buffer, pos, end, message, field_dict, current_depth=0
+      ):
+        del current_depth  # unused
         value = field_dict.get(key)
         if value is None:
           value = field_dict.setdefault(key, new_default(message))
@@ -218,9 +225,12 @@ def DecodeRepeatedField(buffer, pos, end, message, field_dict):
             if new_pos > end:
               raise _DecodeError('Truncated message.')
             return new_pos
+
       return DecodeRepeatedField
     else:
-      def DecodeField(buffer, pos, end, message, field_dict):
+
+      def DecodeField(buffer, pos, end, message, field_dict, current_depth=0):
+        del current_depth  # unused
         (new_value, pos) = decode_value(buffer, pos)
         if pos > end:
           raise _DecodeError('Truncated message.')
@@ -229,6 +239,7 @@ def DecodeField(buffer, pos, end, message, field_dict):
         else:
           field_dict[key] = new_value
         return pos
+
       return DecodeField
 
   return SpecificDecoder
@@ -364,7 +375,9 @@ def EnumDecoder(field_number, is_repeated, is_packed, key, new_default,
   enum_type = key.enum_type
   if is_packed:
     local_DecodeVarint = _DecodeVarint
-    def DecodePackedField(buffer, pos, end, message, field_dict):
+    def DecodePackedField(
+        buffer, pos, end, message, field_dict, current_depth=0
+    ):
       """Decode serialized packed enum to its value and a new position.
 
       Args:
@@ -377,6 +390,7 @@ def DecodePackedField(buffer, pos, end, message, field_dict):
       Returns:
         int, new position in serialized data.
       """
+      del current_depth  # unused
       value = field_dict.get(key)
       if value is None:
         value = field_dict.setdefault(key, new_default(message))
@@ -407,11 +421,14 @@ def DecodePackedField(buffer, pos, end, message, field_dict):
           # pylint: enable=protected-access
         raise _DecodeError('Packed element was truncated.')
       return pos
+
     return DecodePackedField
   elif is_repeated:
     tag_bytes = encoder.TagBytes(field_number, wire_format.WIRETYPE_VARINT)
     tag_len = len(tag_bytes)
-    def DecodeRepeatedField(buffer, pos, end, message, field_dict):
+    def DecodeRepeatedField(
+        buffer, pos, end, message, field_dict, current_depth=0
+    ):
       """Decode serialized repeated enum to its value and a new position.
 
       Args:
@@ -424,6 +441,7 @@ def DecodeRepeatedField(buffer, pos, end, message, field_dict):
       Returns:
         int, new position in serialized data.
       """
+      del current_depth  # unused
       value = field_dict.get(key)
       if value is None:
         value = field_dict.setdefault(key, new_default(message))
@@ -446,9 +464,11 @@ def DecodeRepeatedField(buffer, pos, end, message, field_dict):
           if new_pos > end:
             raise _DecodeError('Truncated message.')
           return new_pos
+
     return DecodeRepeatedField
   else:
-    def DecodeField(buffer, pos, end, message, field_dict):
+
+    def DecodeField(buffer, pos, end, message, field_dict, current_depth=0):
       """Decode serialized repeated enum to its value and a new position.
 
       Args:
@@ -461,6 +481,7 @@ def DecodeField(buffer, pos, end, message, field_dict):
       Returns:
         int, new position in serialized data.
       """
+      del current_depth  # unused
       value_start_pos = pos
       (enum_value, pos) = _DecodeSignedVarint32(buffer, pos)
       if pos > end:
@@ -480,6 +501,7 @@ def DecodeField(buffer, pos, end, message, field_dict):
             (tag_bytes, buffer[value_start_pos:pos].tobytes()))
         # pylint: enable=protected-access
       return pos
+
     return DecodeField
 
 
@@ -538,7 +560,10 @@ def _ConvertToUnicode(memview):
     tag_bytes = encoder.TagBytes(field_number,
                                  wire_format.WIRETYPE_LENGTH_DELIMITED)
     tag_len = len(tag_bytes)
-    def DecodeRepeatedField(buffer, pos, end, message, field_dict):
+    def DecodeRepeatedField(
+        buffer, pos, end, message, field_dict, current_depth=0
+    ):
+      del current_depth  # unused
       value = field_dict.get(key)
       if value is None:
         value = field_dict.setdefault(key, new_default(message))
@@ -553,9 +578,12 @@ def DecodeRepeatedField(buffer, pos, end, message, field_dict):
         if buffer[new_pos:pos] != tag_bytes or new_pos == end:
           # Prediction failed.  Return.
           return new_pos
+
     return DecodeRepeatedField
   else:
-    def DecodeField(buffer, pos, end, message, field_dict):
+
+    def DecodeField(buffer, pos, end, message, field_dict, current_depth=0):
+      del current_depth  # unused
       (size, pos) = local_DecodeVarint(buffer, pos)
       new_pos = pos + size
       if new_pos > end:
@@ -565,6 +593,7 @@ def DecodeField(buffer, pos, end, message, field_dict):
       else:
         field_dict[key] = _ConvertToUnicode(buffer[pos:new_pos])
       return new_pos
+
     return DecodeField
 
 
@@ -579,7 +608,10 @@ def BytesDecoder(field_number, is_repeated, is_packed, key, new_default,
     tag_bytes = encoder.TagBytes(field_number,
                                  wire_format.WIRETYPE_LENGTH_DELIMITED)
     tag_len = len(tag_bytes)
-    def DecodeRepeatedField(buffer, pos, end, message, field_dict):
+    def DecodeRepeatedField(
+        buffer, pos, end, message, field_dict, current_depth=0
+    ):
+      del current_depth  # unused
       value = field_dict.get(key)
       if value is None:
         value = field_dict.setdefault(key, new_default(message))
@@ -594,9 +626,12 @@ def DecodeRepeatedField(buffer, pos, end, message, field_dict):
         if buffer[new_pos:pos] != tag_bytes or new_pos == end:
           # Prediction failed.  Return.
           return new_pos
+
     return DecodeRepeatedField
   else:
-    def DecodeField(buffer, pos, end, message, field_dict):
+
+    def DecodeField(buffer, pos, end, message, field_dict, current_depth=0):
+      del current_depth  # unused
       (size, pos) = local_DecodeVarint(buffer, pos)
       new_pos = pos + size
       if new_pos > end:
@@ -606,6 +641,7 @@ def DecodeField(buffer, pos, end, message, field_dict):
       else:
         field_dict[key] = buffer[pos:new_pos].tobytes()
       return new_pos
+
     return DecodeField
 
 
@@ -621,7 +657,9 @@ def GroupDecoder(field_number, is_repeated, is_packed, key, new_default):
     tag_bytes = encoder.TagBytes(field_number,
                                  wire_format.WIRETYPE_START_GROUP)
     tag_len = len(tag_bytes)
-    def DecodeRepeatedField(buffer, pos, end, message, field_dict):
+    def DecodeRepeatedField(
+        buffer, pos, end, message, field_dict, current_depth=0
+    ):
       value = field_dict.get(key)
       if value is None:
         value = field_dict.setdefault(key, new_default(message))
@@ -630,7 +668,13 @@ def DecodeRepeatedField(buffer, pos, end, message, field_dict):
         if value is None:
           value = field_dict.setdefault(key, new_default(message))
         # Read sub-message.
-        pos = value.add()._InternalParse(buffer, pos, end)
+        current_depth += 1
+        if current_depth > _recursion_limit:
+          raise _DecodeError(
+              'Error parsing message: too many levels of nesting.'
+          )
+        pos = value.add()._InternalParse(buffer, pos, end, current_depth)
+        current_depth -= 1
         # Read end tag.
         new_pos = pos+end_tag_len
         if buffer[pos:new_pos] != end_tag_bytes or new_pos > end:
@@ -640,19 +684,26 @@ def DecodeRepeatedField(buffer, pos, end, message, field_dict):
         if buffer[new_pos:pos] != tag_bytes or new_pos == end:
           # Prediction failed.  Return.
           return new_pos
+
     return DecodeRepeatedField
   else:
-    def DecodeField(buffer, pos, end, message, field_dict):
+
+    def DecodeField(buffer, pos, end, message, field_dict, current_depth=0):
       value = field_dict.get(key)
       if value is None:
         value = field_dict.setdefault(key, new_default(message))
       # Read sub-message.
-      pos = value._InternalParse(buffer, pos, end)
+      current_depth += 1
+      if current_depth > _recursion_limit:
+        raise _DecodeError('Error parsing message: too many levels of nesting.')
+      pos = value._InternalParse(buffer, pos, end, current_depth)
+      current_depth -= 1
       # Read end tag.
       new_pos = pos+end_tag_len
       if buffer[pos:new_pos] != end_tag_bytes or new_pos > end:
         raise _DecodeError('Missing group end tag.')
       return new_pos
+
     return DecodeField
 
 
@@ -666,7 +717,9 @@ def MessageDecoder(field_number, is_repeated, is_packed, key, new_default):
     tag_bytes = encoder.TagBytes(field_number,
                                  wire_format.WIRETYPE_LENGTH_DELIMITED)
     tag_len = len(tag_bytes)
-    def DecodeRepeatedField(buffer, pos, end, message, field_dict):
+    def DecodeRepeatedField(
+        buffer, pos, end, message, field_dict, current_depth=0
+    ):
       value = field_dict.get(key)
       if value is None:
         value = field_dict.setdefault(key, new_default(message))
@@ -677,18 +730,29 @@ def DecodeRepeatedField(buffer, pos, end, message, field_dict):
         if new_pos > end:
           raise _DecodeError('Truncated message.')
         # Read sub-message.
-        if value.add()._InternalParse(buffer, pos, new_pos) != new_pos:
+        current_depth += 1
+        if current_depth > _recursion_limit:
+          raise _DecodeError(
+              'Error parsing message: too many levels of nesting.'
+          )
+        if (
+            value.add()._InternalParse(buffer, pos, new_pos, current_depth)
+            != new_pos
+        ):
           # The only reason _InternalParse would return early is if it
           # encountered an end-group tag.
           raise _DecodeError('Unexpected end-group tag.')
+        current_depth -= 1
         # Predict that the next tag is another copy of the same repeated field.
         pos = new_pos + tag_len
         if buffer[new_pos:pos] != tag_bytes or new_pos == end:
           # Prediction failed.  Return.
           return new_pos
+
     return DecodeRepeatedField
   else:
-    def DecodeField(buffer, pos, end, message, field_dict):
+
+    def DecodeField(buffer, pos, end, message, field_dict, current_depth=0):
       value = field_dict.get(key)
       if value is None:
         value = field_dict.setdefault(key, new_default(message))
@@ -698,11 +762,16 @@ def DecodeField(buffer, pos, end, message, field_dict):
       if new_pos > end:
         raise _DecodeError('Truncated message.')
       # Read sub-message.
-      if value._InternalParse(buffer, pos, new_pos) != new_pos:
+      current_depth += 1
+      if current_depth > _recursion_limit:
+        raise _DecodeError('Error parsing message: too many levels of nesting.')
+      if value._InternalParse(buffer, pos, new_pos, current_depth) != new_pos:
         # The only reason _InternalParse would return early is if it encountered
         # an end-group tag.
         raise _DecodeError('Unexpected end-group tag.')
+      current_depth -= 1
       return new_pos
+
     return DecodeField
 
 
@@ -851,7 +920,8 @@ def MapDecoder(field_descriptor, new_default, is_message_map):
   # Can't read _concrete_class yet; might not be initialized.
   message_type = field_descriptor.message_type
 
-  def DecodeMap(buffer, pos, end, message, field_dict):
+  def DecodeMap(buffer, pos, end, message, field_dict, current_depth=0):
+    del current_depth  # Unused.
     submsg = message_type._concrete_class()
     value = field_dict.get(key)
     if value is None:
@@ -934,7 +1004,16 @@ def _SkipGroup(buffer, pos, end):
     pos = new_pos
 
 
-def _DecodeUnknownFieldSet(buffer, pos, end_pos=None):
+DEFAULT_RECURSION_LIMIT = 100
+_recursion_limit = DEFAULT_RECURSION_LIMIT
+
+
+def SetRecursionLimit(new_limit):
+  global _recursion_limit
+  _recursion_limit = new_limit
+
+
+def _DecodeUnknownFieldSet(buffer, pos, end_pos=None, current_depth=0):
   """Decode UnknownFieldSet.  Returns the UnknownFieldSet and new position."""
 
   unknown_field_set = containers.UnknownFieldSet()
@@ -944,14 +1023,16 @@ def _DecodeUnknownFieldSet(buffer, pos, end_pos=None):
     field_number, wire_type = wire_format.UnpackTag(tag)
     if wire_type == wire_format.WIRETYPE_END_GROUP:
       break
-    (data, pos) = _DecodeUnknownField(buffer, pos, wire_type)
+    (data, pos) = _DecodeUnknownField(buffer, pos, wire_type, current_depth)
     # pylint: disable=protected-access
     unknown_field_set._add(field_number, wire_type, data)
 
   return (unknown_field_set, pos)
 
 
-def _DecodeUnknownField(buffer, pos, wire_type):
+def _DecodeUnknownField(
+    buffer, pos, wire_type, current_depth=0
+):
   """Decode a unknown field.  Returns the UnknownField and new position."""
 
   if wire_type == wire_format.WIRETYPE_VARINT:
@@ -965,7 +1046,11 @@ def _DecodeUnknownField(buffer, pos, wire_type):
     data = buffer[pos:pos+size].tobytes()
     pos += size
   elif wire_type == wire_format.WIRETYPE_START_GROUP:
-    (data, pos) = _DecodeUnknownFieldSet(buffer, pos)
+    current_depth += 1
+    if current_depth >= _recursion_limit:
+      raise _DecodeError('Error parsing message: too many levels of nesting.')
+    data, pos = _DecodeUnknownFieldSet(buffer, pos, None, current_depth)
+    current_depth -= 1
   elif wire_type == wire_format.WIRETYPE_END_GROUP:
     return (0, -1)
   else:

python/google/protobuf/internal/decoder_test.py

@@ -11,8 +11,10 @@
 import io
 import unittest
 
+from google.protobuf import message
 from google.protobuf.internal import decoder
 from google.protobuf.internal import testing_refleaks
+from google.protobuf.internal import wire_format
 
 
 _INPUT_BYTES = b'\x84r\x12'
@@ -52,6 +54,18 @@ def test_decode_varint_bytesio_empty(self):
     size = decoder._DecodeVarint(input_io)
     self.assertEqual(size, None)
 
+  def test_decode_unknown_group_field_too_many_levels(self):
+    data = memoryview(b'\023' * 5_000_000)
+    self.assertRaisesRegex(
+        message.DecodeError,
+        'Error parsing message',
+        decoder._DecodeUnknownField,
+        data,
+        1,
+        wire_format.WIRETYPE_START_GROUP,
+        1
+    )
+
 
 if __name__ == '__main__':
   unittest.main()