You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ga...@apache.org on 2026/09/23 23:23:38 UTC
(spark) branch master updated: [SPARK-59508][PYTHON] Harden client safety with restricted unpickle serializer
This is an automated email from the ASF dual-hosted git repository.
gaogaotiantian pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 6886811d1b38 [SPARK-59508][PYTHON] Harden client safety with restricted unpickle serializer
6886811d1b38 is described below
commit 6886811d1b38b2452e394486026e61be9281f3d9
Author: Tian Gao <ga...@hotmail.com>
AuthorDate: Wed Sep 23 16:23:23 2026 -0700
[SPARK-59508][PYTHON] Harden client safety with restricted unpickle serializer
### What changes were proposed in this pull request?
* Add an optional argument `allowed_names` for `CloudPickleSerializer`. If specified, it will be the allowlist for unpickling.
* Replace the unrestricted `pickleSer` with a restricted one which only allows unpickling `pstat.Stats`.
### Why are the changes needed?
The client unpickles arbitrary bytes from the server. If the server is compromised, this will behave as a RCE. We should only allow pickling on the data structure we actually use.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
A new test is added to test the restricted unpickler.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #58798 from gaogaotiantian/safe-pickle.
Authored-by: Tian Gao <ga...@hotmail.com>
Signed-off-by: Tian Gao <ga...@hotmail.com>
---
python/pyspark/accumulators.py | 6 +++++-
python/pyspark/serializers.py | 27 +++++++++++++++++++++++++++
python/pyspark/sql/connect/client/core.py | 13 +++++++++++--
python/pyspark/tests/test_serializers.py | 11 +++++++++++
4 files changed, 54 insertions(+), 3 deletions(-)
diff --git a/python/pyspark/accumulators.py b/python/pyspark/accumulators.py
index e09e987ff417..872a357be183 100644
--- a/python/pyspark/accumulators.py
+++ b/python/pyspark/accumulators.py
@@ -25,7 +25,7 @@ import threading
from typing import TYPE_CHECKING, Callable, Dict, Generic, Optional, Tuple, Type, TypeVar, Union
from pyspark.errors import PySparkRuntimeError
-from pyspark.serializers import CPickleSerializer, read_int
+from pyspark.serializers import CloudPickleSerializer, CPickleSerializer, read_int
if TYPE_CHECKING:
from socketserver import BaseRequestHandler
@@ -40,6 +40,10 @@ U = TypeVar("U", bound=Union["SupportsIAdd", int, float, complex])
pickleSer = CPickleSerializer()
+# Safe serializer for special accumulators. The client will unpickle the metrics from the server
+# unconditionally, so we need to restrict the allowed classes.
+specialAccumulatorSer = CloudPickleSerializer(allowed_names=[("pstats", "Stats")])
+
# Holds accumulators registered on the current machine, keyed by ID. This is then used to send
# the local accumulator updates back to the driver program at the end of a task.
_accumulatorRegistry: Dict[int, "Accumulator"] = {}
diff --git a/python/pyspark/serializers.py b/python/pyspark/serializers.py
index f00cbf7d8a6b..4266020ee999 100644
--- a/python/pyspark/serializers.py
+++ b/python/pyspark/serializers.py
@@ -62,6 +62,7 @@ which contains two batches of two objects:
import codecs
import collections
+import io
import itertools
import marshal
import os
@@ -414,6 +415,17 @@ if os.environ.get("PYSPARK_ENABLE_NAMEDTUPLE_PATCH") == "1":
_hijack_namedtuple()
+class _RestrictedUnpickler(pickle.Unpickler):
+ def __init__(self, file, allowed_names, *, encoding="bytes"):
+ super().__init__(file, encoding=encoding)
+ self.allowed_names = allowed_names
+
+ def find_class(self, module, name):
+ if (module, name) not in self.allowed_names:
+ raise pickle.UnpicklingError(f"Unpickling {module}.{name} is not allowed")
+ return super().find_class(module, name)
+
+
class PickleSerializer(FramedSerializer):
"""
Serializes objects using Python's pickle serializer:
@@ -432,6 +444,18 @@ class PickleSerializer(FramedSerializer):
class CloudPickleSerializer(FramedSerializer):
+ # Class-level default for `allowed_names`. Instances of this serializer are pickled and
+ # shipped across the driver and workers, which may run different Spark versions. Pickle
+ # restores an instance via __new__ + __dict__ without calling __init__, so an instance
+ # pickled by a version that predates `allowed_names` would otherwise lack the attribute
+ # entirely. Declaring the default here keeps `self.allowed_names` resolvable in that
+ # cross-version case; only a version-matched, explicitly restricted instance sets it.
+ allowed_names = None
+
+ def __init__(self, allowed_names=None):
+ super().__init__()
+ self.allowed_names = allowed_names
+
def dumps(self, obj):
from pyspark.util import print_exec
@@ -449,6 +473,9 @@ class CloudPickleSerializer(FramedSerializer):
raise pickle.PicklingError(msg)
def loads(self, obj, encoding="bytes"):
+ if self.allowed_names is not None:
+ unpickler = _RestrictedUnpickler(io.BytesIO(obj), self.allowed_names, encoding=encoding)
+ return unpickler.load()
return cloudpickle.loads(obj, encoding=encoding)
diff --git a/python/pyspark/sql/connect/client/core.py b/python/pyspark/sql/connect/client/core.py
index e5826887f195..855d2d3f21d4 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -26,6 +26,7 @@ import concurrent.futures
import copy
import logging
import os
+import pickle
import platform
import sys
import threading
@@ -64,7 +65,7 @@ import pyspark
import pyspark.sql.connect.proto as pb2
import pyspark.sql.connect.proto.base_pb2_grpc as grpc_lib
import pyspark.sql.connect.types as types
-from pyspark.accumulators import SpecialAccumulatorIds, pickleSer
+from pyspark.accumulators import SpecialAccumulatorIds, specialAccumulatorSer
from pyspark.errors import (
PySparkAssertionError,
PySparkNotImplementedError,
@@ -1794,7 +1795,15 @@ class SparkConnectClient(object):
else:
if observed_metrics.name == "__python_accumulator__":
for metric in observed_metrics.metrics:
- aid, update = pickleSer.loads(LiteralExpression._to_value(metric))
+ try:
+ aid, update = specialAccumulatorSer.loads(
+ LiteralExpression._to_value(metric)
+ )
+ except pickle.UnpicklingError as e:
+ # We found unexpected class/function in the accumulator metric.
+ # We will ignore this metric and continue.
+ logger.warning(f"Error unpickling accumulator metric: {e}")
+ continue
if aid == SpecialAccumulatorIds.SQL_UDF_PROFIER_V2:
self._profiler_collector._update(update)
elif observed_metrics.name in observations:
diff --git a/python/pyspark/tests/test_serializers.py b/python/pyspark/tests/test_serializers.py
index 243ecc89fbb4..10fb77b65b00 100644
--- a/python/pyspark/tests/test_serializers.py
+++ b/python/pyspark/tests/test_serializers.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+import collections
import math
+import pickle
import sys
import unittest
@@ -157,6 +159,15 @@ class SerializationTestCase(unittest.TestCase):
hash(CompressedSerializer(CPickleSerializer()))
hash(FlattenedValuesSerializer(CPickleSerializer()))
+ def test_restricted_unpickler(self):
+ ser = CloudPickleSerializer(allowed_names=[("collections", "deque")])
+ q = collections.deque([1, 2, 3])
+ q2 = ser.loads(ser.dumps(q))
+ self.assertEqual(q, q2)
+
+ with self.assertRaises(pickle.UnpicklingError):
+ ser.loads(ser.dumps(collections.defaultdict(int)))
+
@unittest.skipIf(not have_scipy, "SciPy not installed")
class SciPyTests(PySparkTestCase):
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org