fix(spanner): catch recursion and decode errors in proto parsing to prevent DoS

Author: sinhasubham

packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py

@@ -28,7 +28,7 @@
 from google.api_core.exceptions import Aborted
 from google.cloud._helpers import _date_from_iso8601_date
 from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
-from google.protobuf.message import Message
+from google.protobuf.message import DecodeError, Message
 from google.protobuf.struct_pb2 import ListValue, Value
 from google.rpc.error_details_pb2 import RetryInfo
 
@@ -603,8 +603,14 @@ def _parse_proto(value_pb, column_info, field_name):
         default_proto_message = column_info.get(field_name)
         if isinstance(default_proto_message, Message):
             proto_message = type(default_proto_message)()
-            proto_message.ParseFromString(bytes_value)
-            return proto_message
+            try:
+                proto_message.ParseFromString(bytes_value)
+                return proto_message
+            except (DecodeError, RecursionError):
+                log.warning(
+                    "Warning: Field could not be parsed as Proto due to excessive nesting/corruption. Returning raw bytes."
+                )
+                return bytes_value
     return bytes_value
 
 

packages/google-cloud-spanner/tests/unit/test__helpers.py

@@ -771,6 +771,57 @@ def test_w_proto_message(self):
             self._callFUT(value_pb, field_type, field_name, column_info), VALUE
         )
 
+    def test_w_proto_message_decode_error(self):
+        import base64
+        from unittest import mock
+
+        from google.protobuf.message import DecodeError
+        from google.protobuf.struct_pb2 import Value
+
+        from google.cloud.spanner_v1 import Type, TypeCode
+
+        from .testdata import singer_pb2
+
+        VALUE = singer_pb2.SingerInfo()
+        field_type = Type(code=TypeCode.PROTO)
+        field_name = "proto_message_column"
+        raw_bytes = VALUE.SerializeToString()
+        value_pb = Value(string_value=base64.b64encode(raw_bytes))
+        column_info = {"proto_message_column": singer_pb2.SingerInfo()}
+
+        # Mock ParseFromString to raise DecodeError
+        with mock.patch(
+            "google.protobuf.message.Message.ParseFromString",
+            side_effect=DecodeError("Mock Decode Error"),
+        ):
+            result = self._callFUT(value_pb, field_type, field_name, column_info)
+            # Should return raw bytes
+            self.assertEqual(result, raw_bytes)
+
+    def test_w_proto_message_recursion_error(self):
+        import base64
+        from unittest import mock
+
+        from google.protobuf.struct_pb2 import Value
+
+        from google.cloud.spanner_v1 import Type, TypeCode
+
+        from .testdata import singer_pb2
+
+        VALUE = singer_pb2.SingerInfo()
+        field_type = Type(code=TypeCode.PROTO)
+        field_name = "proto_message_column"
+        raw_bytes = VALUE.SerializeToString()
+        value_pb = Value(string_value=base64.b64encode(raw_bytes))
+        column_info = {"proto_message_column": singer_pb2.SingerInfo()}
+
+        with mock.patch(
+            "google.protobuf.message.Message.ParseFromString",
+            side_effect=RecursionError("Mock Recursion Error"),
+        ):
+            result = self._callFUT(value_pb, field_type, field_name, column_info)
+            self.assertEqual(result, raw_bytes)
+
     def test_w_proto_enum(self):
         from google.protobuf.struct_pb2 import Value