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:44:01 UTC
(spark) branch branch-4.1 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 branch-4.1
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.1 by this push:
new 772e361eb0bb [SPARK-59508][PYTHON] Harden client safety with restricted unpickle serializer
772e361eb0bb is described below
commit 772e361eb0bbe10c01f9993f8921b756b07b157b
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
* 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`.
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.
No.
A new test is added to test the restricted unpickler.
No.
Closes #58798 from gaogaotiantian/safe-pickle.
Authored-by: Tian Gao <ga...@hotmail.com>
Signed-off-by: Tian Gao <ga...@hotmail.com>
(cherry picked from commit 6886811d1b38b2452e394486026e61be9281f3d9)
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 | 15 +++++++++++----
python/pyspark/tests/test_serializers.py | 11 +++++++++++
4 files changed, 54 insertions(+), 5 deletions(-)
diff --git a/python/pyspark/accumulators.py b/python/pyspark/accumulators.py
index e557fe1cd8fb..2f273b6354d8 100644
--- a/python/pyspark/accumulators.py
+++ b/python/pyspark/accumulators.py
@@ -23,7 +23,7 @@ import socketserver
import threading
from typing import Callable, Dict, Generic, Tuple, Type, TYPE_CHECKING, TypeVar, Union, Optional
-from pyspark.serializers import read_int, CPickleSerializer
+from pyspark.serializers import read_int, CPickleSerializer, CloudPickleSerializer
from pyspark.errors import PySparkRuntimeError
if TYPE_CHECKING:
@@ -38,6 +38,10 @@ U = TypeVar("U", bound="SupportsIAdd")
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 821e142304aa..cdef1eb065c7 100644
--- a/python/pyspark/serializers.py
+++ b/python/pyspark/serializers.py
@@ -55,6 +55,7 @@ which contains two batches of two objects:
import sys
import os
+import io
from itertools import chain, product
import marshal
import struct
@@ -434,6 +435,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):
"""
@@ -453,6 +465,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
@@ -470,6 +494,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 80d83c69c45c..a37e9d977359 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -30,6 +30,7 @@ import logging
import threading
import os
import copy
+import pickle
import platform
import urllib.parse
import uuid
@@ -63,7 +64,7 @@ from google.protobuf import text_format, any_pb2
from google.rpc import error_details_pb2
from pyspark.util import is_remote_only
-from pyspark.accumulators import SpecialAccumulatorIds
+from pyspark.accumulators import SpecialAccumulatorIds, specialAccumulatorSer
from pyspark.version import __version__
from pyspark.resource.information import ResourceInformation
from pyspark.sql.metrics import MetricValue, PlanMetrics, ExecutionInfo, ObservedMetrics
@@ -1520,10 +1521,16 @@ class SparkConnectClient(object):
logger.debug("Received observed metric batch.")
for observed_metrics in self._build_observed_metrics(b.observed_metrics):
if observed_metrics.name == "__python_accumulator__":
- from pyspark.worker_util import pickleSer
-
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:
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 230723e10572..9b352222bf73 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
@@ -158,6 +160,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