fix(spanner_dbapi): replace insecure pickle with json for partition deserialization (#17014)

sinhasubham · olavloite · web-flow · commit 86e57cb9b6ec · 2026-06-01T21:20:41.000+05:30
This PR resolves a critical Insecure Deserialization vulnerability
(potential Remote Code Execution) in the `spanner_dbapi` module
[b/510871112](b/510871112) . Previously, the module utilized
`pickle.loads()` to decode partition IDs provided by users via the `RUN
PARTITION` statement, creating a possibility for arbitrary code
execution attack payloads.
We have fully eliminated `pickle` usage in this module and migrated to
standard `json` serialization.

---------

Co-authored-by: Knut Olav L&oslash;ite &lt;koloite@gmail.com&gt;
diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/partition_helper.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/partition_helper.py
@@ -13,23 +13,145 @@
 # limitations under the License.
 
 import base64
+import copy
+import datetime
 import gzip
-import pickle
+import json
 from dataclasses import dataclass
 from typing import Any
 
+from google.protobuf.json_format import MessageToDict, ParseDict
+from google.protobuf.message import Message
+from google.protobuf.struct_pb2 import Struct
+
 from google.cloud.spanner_v1 import BatchTransactionId
+from google.cloud.spanner_v1._helpers import _make_value_pb
+from google.cloud.spanner_v1.types import DirectedReadOptions, ExecuteSqlRequest, Type
+
+_PROTO_CLASS_MAP = {
+    "QueryOptions": ExecuteSqlRequest.QueryOptions,
+    "DirectedReadOptions": DirectedReadOptions,
+    "Struct": Struct,
+    "Type": Type,
+}
+
+
+def _serialize_value(val: Any) -> Any:
+    if isinstance(val, bytes):
+        return {"__type__": "bytes", "value": base64.b64encode(val).decode("utf-8")}
+    elif isinstance(val, datetime.datetime):
+        return {"__type__": "datetime", "value": val.isoformat()}
+    elif hasattr(val, "_pb"):
+        return {
+            "__type__": "protobuf",
+            "class": val.__class__.__name__,
+            "value": MessageToDict(val._pb, preserving_proto_field_name=True),
+        }
+    elif isinstance(val, Message):
+        return {
+            "__type__": "protobuf",
+            "class": val.__class__.__name__,
+            "value": MessageToDict(val, preserving_proto_field_name=True),
+        }
+    elif isinstance(val, dict):
+        return {k: _serialize_value(v) for k, v in val.items()}
+    elif isinstance(val, list):
+        return [_serialize_value(v) for v in val]
+    elif isinstance(val, tuple):
+        return {"__type__": "tuple", "value": [_serialize_value(v) for v in val]}
+    return val
+
+
+def _deserialize_value(val: Any) -> Any:
+    if isinstance(val, dict):
+        if "__type__" in val:
+            t = val["__type__"]
+            if t == "bytes":
+                return base64.b64decode(val["value"])
+            elif t == "datetime":
+                dt_str = val["value"]
+                if dt_str.endswith("Z"):
+                    dt_str = dt_str[:-1] + "+00:00"
+                return datetime.datetime.fromisoformat(dt_str)
+            elif t == "tuple":
+                return tuple(_deserialize_value(x) for x in val["value"])
+            elif t == "protobuf":
+                cls_name = val.get("class")
+                dict_val = val["value"]
+                if cls_name in _PROTO_CLASS_MAP:
+                    cls = _PROTO_CLASS_MAP[cls_name]
+                    msg = cls()._pb if hasattr(cls(), "_pb") else cls()
+                    ParseDict(dict_val, msg)
+                    return cls(msg) if hasattr(cls(), "_pb") else msg
+                return _deserialize_value(dict_val)
+        return {k: _deserialize_value(v) for k, v in val.items()}
+    elif isinstance(val, list):
+        return [_deserialize_value(v) for v in val]
+    return val
+
+
+def _unpack_value_pb(value):
+    which = value.WhichOneof("kind")
+    if which == "null_value":
+        return None
+    elif which == "number_value":
+        return value.number_value
+    elif which == "string_value":
+        return value.string_value
+    elif which == "bool_value":
+        return value.bool_value
+    elif which == "struct_value":
+        return {k: _unpack_value_pb(v) for k, v in value.struct_value.fields.items()}
+    elif which == "list_value":
+        return [_unpack_value_pb(v) for v in value.list_value.values]
+    return None
 
 
 def decode_from_string(encoded_partition_id):
     gzip_bytes = base64.b64decode(bytes(encoded_partition_id, "utf-8"))
     partition_id_bytes = gzip.decompress(gzip_bytes)
-    return pickle.loads(partition_id_bytes)
+
+    data = json.loads(partition_id_bytes.decode("utf-8"))
+    btid_data = data["batch_transaction_id"]
+    btid = BatchTransactionId(
+        transaction_id=_deserialize_value(btid_data["transaction_id"]),
+        session_id=btid_data["session_id"],
+        read_timestamp=_deserialize_value(btid_data["read_timestamp"]),
+    )
+    partition_result = _deserialize_value(data["partition_result"])
+
+    # Post-process query params back from Protobuf Struct to Python primitives
+    if "query" in partition_result and "params" in partition_result["query"]:
+        params_pb = partition_result["query"]["params"]
+        if params_pb:
+            partition_result["query"]["params"] = {
+                k: _unpack_value_pb(v) for k, v in params_pb.fields.items()
+            }
+
+    return PartitionId(btid, partition_result)
 
 
 def encode_to_string(batch_transaction_id, partition_result):
-    partition_id = PartitionId(batch_transaction_id, partition_result)
-    partition_id_bytes = pickle.dumps(partition_id)
+    # Copy to avoid modifying the caller's dictionary in connection.py
+    partition_result = copy.deepcopy(partition_result)
+
+    # Pre-process query params into a Protobuf Struct
+    if "query" in partition_result and "params" in partition_result["query"]:
+        params = partition_result["query"]["params"]
+        if params:
+            params_pb = Struct(fields={k: _make_value_pb(v) for k, v in params.items()})
+            partition_result["query"]["params"] = params_pb
+
+    data = {
+        "batch_transaction_id": {
+            "transaction_id": _serialize_value(batch_transaction_id.transaction_id),
+            "session_id": batch_transaction_id.session_id,
+            "read_timestamp": _serialize_value(batch_transaction_id.read_timestamp),
+        },
+        "partition_result": _serialize_value(partition_result),
+    }
+
+    partition_id_bytes = json.dumps(data).encode("utf-8")
     gzip_bytes = gzip.compress(partition_id_bytes)
     return str(base64.b64encode(gzip_bytes), "utf-8")
 
diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/testing/mock_spanner.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/testing/mock_spanner.py
@@ -36,16 +36,21 @@ class MockSpanner:
     def __init__(self):
         self.results = {}
         self.execute_streaming_sql_results = {}
+        self.partition_results = {}
         self.errors = {}
 
     def clear_results(self):
         self.results = {}
         self.execute_streaming_sql_results = {}
+        self.partition_results = {}
         self.errors = {}
 
     def add_result(self, sql: str, result: result_set.ResultSet):
         self.results[sql.lower().strip()] = result
 
+    def add_partition_result(self, sql: str, result: spanner.PartitionResponse):
+        self.partition_results[sql.lower().strip()] = result
+
     def add_execute_streaming_sql_results(
         self, sql: str, partial_result_sets: list[result_set.PartialResultSet]
     ):
@@ -57,6 +62,12 @@ def get_result(self, sql: str) -> result_set.ResultSet:
             raise ValueError(f"No result found for {sql}")
         return result
 
+    def get_partition_result(self, sql: str) -> spanner.PartitionResponse:
+        result = self.partition_results.get(sql.lower().strip())
+        if result is None:
+            return spanner.PartitionResponse()
+        return result
+
     def add_error(self, method: str, error: _Status):
         if not hasattr(self, "_errors_list"):
             self._errors_list = {}
@@ -300,11 +311,12 @@ def Rollback(self, request, context):
 
     def PartitionQuery(self, request, context):
         self._requests.append(request)
-        return spanner.PartitionResponse()
+        return self.mock_spanner.get_partition_result(request.sql)
 
     def PartitionRead(self, request, context):
         self._requests.append(request)
-        return spanner.PartitionResponse()
+        # For reads, look up by target table name
+        return self.mock_spanner.get_partition_result(request.table)
 
     def BatchWrite(self, request, context):
         self._requests.append(request)
diff --git a/packages/google-cloud-spanner/tests/_helpers.py b/packages/google-cloud-spanner/tests/_helpers.py
@@ -1,7 +1,10 @@
 from os import getenv
 from unittest import IsolatedAsyncioTestCase
 
-import mock
+try:
+    import mock
+except ImportError:
+    import unittest.mock as mock
 
 from google.cloud.spanner_v1 import gapic_version
 from google.cloud.spanner_v1.database_sessions_manager import TransactionType
diff --git a/packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_partition_query.py b/packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_partition_query.py
@@ -0,0 +1,134 @@
+# Copyright 2024 Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from google.cloud.spanner_dbapi.connection import Connection
+from google.cloud.spanner_dbapi.parsed_statement import ParsedStatement, Statement
+from google.cloud.spanner_v1 import TypeCode
+from google.cloud.spanner_v1.types import spanner as spanner_types
+from tests.mockserver_tests.mock_server_test_base import (
+    MockServerTestBase,
+    add_single_result,
+)
+
+
+class TestDbapiPartitionQuery(MockServerTestBase):
+    def test_partition_query_and_run_partition(self):
+        sql = "SELECT name FROM users WHERE active = true"
+
+        # 1. Set up mock results for PartitionQuery RPC in the mock servicer
+        partition_response = spanner_types.PartitionResponse()
+        partition_response.partitions.extend(
+            [
+                spanner_types.Partition(partition_token=b"mock-token-1"),
+                spanner_types.Partition(partition_token=b"mock-token-2"),
+            ]
+        )
+        self.spanner_service.mock_spanner.add_partition_result(sql, partition_response)
+
+        # 2. Set up mock results for ExecuteSql when executing the partitions
+        add_single_result(sql, "name", TypeCode.STRING, [("Alice",), ("Bob",)])
+
+        # 3. Connect via DB-API and mark connection as read-only (required for partitioning)
+        connection = Connection(self.instance, self.database)
+        connection._read_only = True
+
+        # Define partitioning parameters inside DB-API Statement
+        from google.cloud.spanner_dbapi.parsed_statement import (
+            ClientSideStatementType,
+            StatementType,
+        )
+
+        parsed = ParsedStatement(
+            statement_type=StatementType.CLIENT_SIDE,
+            statement=Statement(sql),
+            client_side_statement_type=ClientSideStatementType.PARTITION_QUERY,
+            client_side_statement_params=["SELECT name FROM users WHERE active = true"],
+        )
+
+        # Generate serialized token strings (Base64 + GZip JSON)
+        partition_ids = connection.partition_query(parsed)
+        self.assertEqual(2, len(partition_ids))
+
+        # 4. Reconstruct & Execute the partitions by deserializing their tokens
+        all_names = []
+        for token in partition_ids:
+            result_stream = connection.run_partition(token)
+            for row in result_stream:
+                all_names.append(row[0])
+
+        # Verify results are successfully round-tripped and parsed
+        self.assertIn("Alice", all_names)
+        self.assertIn("Bob", all_names)
+
+    def test_partition_query_with_complex_parameters(self):
+        import datetime
+        import decimal
+
+        sql = "SELECT name FROM users WHERE active = @active AND salary > @salary AND signup_time = @signup_time"
+
+        # Set up complex parameter values (bool, Decimal, datetime)
+        params = {
+            "active": True,
+            "salary": decimal.Decimal("75000.50"),
+            "signup_time": datetime.datetime(
+                2026, 5, 10, 12, 34, 56, tzinfo=datetime.timezone.utc
+            ),
+        }
+        from google.cloud.spanner_v1 import Type
+
+        param_types = {
+            "active": Type(code=TypeCode.BOOL),
+            "salary": Type(code=TypeCode.NUMERIC),
+            "signup_time": Type(code=TypeCode.TIMESTAMP),
+        }
+
+        # 1. Mock results for the partition generation RPC
+        partition_response = spanner_types.PartitionResponse()
+        partition_response.partitions.extend(
+            [spanner_types.Partition(partition_token=b"complex-mock-token-1")]
+        )
+        self.spanner_service.mock_spanner.add_partition_result(sql, partition_response)
+
+        # 2. Mock results for execution of partition streaming SQL
+        add_single_result(sql, "name", TypeCode.STRING, [("Charlie",)])
+
+        # 3. Establish Connection
+        connection = Connection(self.instance, self.database)
+        connection._read_only = True
+
+        from google.cloud.spanner_dbapi.parsed_statement import (
+            ClientSideStatementType,
+            StatementType,
+        )
+
+        parsed = ParsedStatement(
+            statement_type=StatementType.CLIENT_SIDE,
+            statement=Statement(sql, params=params, param_types=param_types),
+            client_side_statement_type=ClientSideStatementType.PARTITION_QUERY,
+            client_side_statement_params=[sql],
+        )
+
+        # Execute partition generation - this serializes query parameters!
+        partition_ids = connection.partition_query(parsed)
+        self.assertEqual(1, len(partition_ids))
+
+        # 4. Reconstruct and run the partition E2E
+        all_names = []
+        for token in partition_ids:
+            result_stream = connection.run_partition(token)
+            for row in result_stream:
+                all_names.append(row[0])
+
+        self.assertEqual(["Charlie"], all_names)
diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_partition_helper.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_partition_helper.py
