From 707e76effd808f514e3be6608e8140c54827e867 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Thu, 7 May 2026 21:17:45 +0300 Subject: [PATCH 1/8] security: add RestrictedUnpickler for safe v0 pickle deserialization Ref: google/adk-python#5634 --- .../adk/sessions/schemas/_safe_unpickle.py | 84 +++++++++++++++++++ 1 file changed, 84 insertions(+) create mode 100644 src/google/adk/sessions/schemas/_safe_unpickle.py diff --git a/src/google/adk/sessions/schemas/_safe_unpickle.py b/src/google/adk/sessions/schemas/_safe_unpickle.py new file mode 100644 index 0000000000..9018953692 --- /dev/null +++ b/src/google/adk/sessions/schemas/_safe_unpickle.py @@ -0,0 +1,84 @@ +# Copyright 2026 Google LLC +# +# 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. +"""Restricted unpickler for safe deserialization of v0 EventActions data. + +The v0 schema stored EventActions as pickle blobs. This module provides a +safe deserialization path that only allows known ADK and standard types, +blocking arbitrary code execution via crafted pickle payloads. + +See: https://docs.python.org/3/library/pickle.html#restricting-globals +""" + +from __future__ import annotations + +import io +import logging +import os +import pickle +from typing import Any + +logger = logging.getLogger("google_adk." + __name__) + +_ALLOWED_MODULE_PREFIXES: tuple[str, ...] = ( + "google.adk.", + "google.genai.", + "pydantic.", + "pydantic_core.", +) + +_ALLOWED_GLOBALS: dict[str, set[str]] = { + "builtins": { + "dict", "list", "set", "tuple", "frozenset", "bytes", + "bytearray", "True", "False", "None", "type", "object", + "complex", "slice", "range", "int", "float", "str", "bool", + }, + "collections": {"OrderedDict", "defaultdict"}, + "datetime": {"datetime", "date", "time", "timedelta", "timezone"}, + "copy_reg": {"_reconstructor"}, + "copyreg": {"_reconstructor", "__newobj__"}, + "_codecs": {"encode"}, +} + + +class _RestrictedUnpickler(pickle.Unpickler): + """Unpickler that only allows reconstruction of known-safe types.""" + + def find_class(self, module: str, name: str) -> Any: + for prefix in _ALLOWED_MODULE_PREFIXES: + if module.startswith(prefix): + return super().find_class(module, name) + allowed_names = _ALLOWED_GLOBALS.get(module) + if allowed_names and name in allowed_names: + return super().find_class(module, name) + raise pickle.UnpicklingError( + f"Blocked unsafe pickle global: {module}.{name}. " + f"If this is a legitimate ADK type, please file an issue at " + f"https://github.com/google/adk-python/issues" + ) + + +def safe_loads(data: bytes) -> Any: + """Deserialize pickle bytes using a restricted unpickler. + + If ADK_ALLOW_UNSAFE_V0_PICKLE=1 is set, falls back to unrestricted + pickle.loads() for compatibility. A deprecation warning is logged. + """ + if os.environ.get("ADK_ALLOW_UNSAFE_V0_PICKLE") == "1": + logger.warning( + "ADK_ALLOW_UNSAFE_V0_PICKLE is set - using unrestricted " + "pickle.loads(). This is unsafe and will be removed in a " + "future release. Migrate to the v1 JSON schema." + ) + return pickle.loads(data) # noqa: S301 + return _RestrictedUnpickler(io.BytesIO(data)).load() From a382b8e7515ebe6d6b77fdd31ef12e2a530234ba Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Thu, 7 May 2026 21:17:59 +0300 Subject: [PATCH 2/8] security: use RestrictedUnpickler in v0 runtime read path Replace raw pickle.loads() with safe_loads() in DynamicPickleType.process_result_value(). Ref: google/adk-python#5634 --- src/google/adk/sessions/schemas/v0.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/google/adk/sessions/schemas/v0.py b/src/google/adk/sessions/schemas/v0.py index e4a4368c6d..62ee72b607 100644 --- a/src/google/adk/sessions/schemas/v0.py +++ b/src/google/adk/sessions/schemas/v0.py @@ -31,6 +31,8 @@ import json import logging import pickle + +from ._safe_unpickle import safe_loads as _safe_pickle_loads from typing import Any from typing import Optional @@ -114,7 +116,7 @@ def process_result_value(self, value, dialect): """Ensures the raw bytes from the database are unpickled back into a Python object.""" if value is not None: if dialect.name in ("spanner+spanner", "mysql"): - return pickle.loads(value) + return _safe_pickle_loads(value) return value From ba04ed4a886c8dfe392578cc5ba45b649886c4e8 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Thu, 7 May 2026 21:18:22 +0300 Subject: [PATCH 3/8] security: use RestrictedUnpickler in v0 migration path Replace raw pickle.loads() with safe_loads() in _row_to_event() migration function. Ref: google/adk-python#5634 --- .../adk/sessions/migration/migrate_from_sqlalchemy_pickle.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py index a6d1ad2a78..48b9c6cc41 100644 --- a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py +++ b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py @@ -21,6 +21,8 @@ import json import logging import pickle + +from google.adk.sessions.schemas._safe_unpickle import safe_loads as _safe_pickle_loads import sys from typing import Any @@ -59,7 +61,7 @@ def _row_to_event(row: dict) -> Event: if actions_val is not None: try: if isinstance(actions_val, bytes): - actions = pickle.loads(actions_val) + actions = _safe_pickle_loads(actions_val) else: # for spanner - it might return object directly actions = actions_val except Exception as e: From b76b2f3142c00b5eafff87f341a8ab0372368428 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Thu, 7 May 2026 21:22:54 +0300 Subject: [PATCH 4/8] chore: retrigger CLA check after signing From d9a0f7ab37f51060312b571f14e2ba6b56990b91 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Thu, 7 May 2026 22:50:56 +0300 Subject: [PATCH 5/8] security: add enum types to RestrictedUnpickler allowlist Covers potential enum values in state_delta/agent_state dict[str, Any] fields. Ref: google/adk-python#5634 --- src/google/adk/sessions/schemas/_safe_unpickle.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/google/adk/sessions/schemas/_safe_unpickle.py b/src/google/adk/sessions/schemas/_safe_unpickle.py index 9018953692..8309462d4b 100644 --- a/src/google/adk/sessions/schemas/_safe_unpickle.py +++ b/src/google/adk/sessions/schemas/_safe_unpickle.py @@ -48,6 +48,7 @@ "copy_reg": {"_reconstructor"}, "copyreg": {"_reconstructor", "__newobj__"}, "_codecs": {"encode"}, + "enum": {"__new__", "Enum", "IntEnum", "StrEnum"}, } From 422e1d876e42afe80eb383dc65ca4263050f2ea5 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Sat, 9 May 2026 10:10:19 +0300 Subject: [PATCH 6/8] style: run pyink + isort on _safe_unpickle.py Ref: google/adk-python#5634 --- .../adk/sessions/schemas/_safe_unpickle.py | 72 +++++++++++-------- 1 file changed, 44 insertions(+), 28 deletions(-) diff --git a/src/google/adk/sessions/schemas/_safe_unpickle.py b/src/google/adk/sessions/schemas/_safe_unpickle.py index 8309462d4b..a61fcfb5b6 100644 --- a/src/google/adk/sessions/schemas/_safe_unpickle.py +++ b/src/google/adk/sessions/schemas/_safe_unpickle.py @@ -39,9 +39,25 @@ _ALLOWED_GLOBALS: dict[str, set[str]] = { "builtins": { - "dict", "list", "set", "tuple", "frozenset", "bytes", - "bytearray", "True", "False", "None", "type", "object", - "complex", "slice", "range", "int", "float", "str", "bool", + "dict", + "list", + "set", + "tuple", + "frozenset", + "bytes", + "bytearray", + "True", + "False", + "None", + "type", + "object", + "complex", + "slice", + "range", + "int", + "float", + "str", + "bool", }, "collections": {"OrderedDict", "defaultdict"}, "datetime": {"datetime", "date", "time", "timedelta", "timezone"}, @@ -53,33 +69,33 @@ class _RestrictedUnpickler(pickle.Unpickler): - """Unpickler that only allows reconstruction of known-safe types.""" + """Unpickler that only allows reconstruction of known-safe types.""" - def find_class(self, module: str, name: str) -> Any: - for prefix in _ALLOWED_MODULE_PREFIXES: - if module.startswith(prefix): - return super().find_class(module, name) - allowed_names = _ALLOWED_GLOBALS.get(module) - if allowed_names and name in allowed_names: - return super().find_class(module, name) - raise pickle.UnpicklingError( - f"Blocked unsafe pickle global: {module}.{name}. " - f"If this is a legitimate ADK type, please file an issue at " - f"https://github.com/google/adk-python/issues" - ) + def find_class(self, module: str, name: str) -> Any: + for prefix in _ALLOWED_MODULE_PREFIXES: + if module.startswith(prefix): + return super().find_class(module, name) + allowed_names = _ALLOWED_GLOBALS.get(module) + if allowed_names and name in allowed_names: + return super().find_class(module, name) + raise pickle.UnpicklingError( + f"Blocked unsafe pickle global: {module}.{name}. " + "If this is a legitimate ADK type, please file an issue at " + "https://github.com/google/adk-python/issues" + ) def safe_loads(data: bytes) -> Any: - """Deserialize pickle bytes using a restricted unpickler. + """Deserialize pickle bytes using a restricted unpickler. - If ADK_ALLOW_UNSAFE_V0_PICKLE=1 is set, falls back to unrestricted - pickle.loads() for compatibility. A deprecation warning is logged. - """ - if os.environ.get("ADK_ALLOW_UNSAFE_V0_PICKLE") == "1": - logger.warning( - "ADK_ALLOW_UNSAFE_V0_PICKLE is set - using unrestricted " - "pickle.loads(). This is unsafe and will be removed in a " - "future release. Migrate to the v1 JSON schema." - ) - return pickle.loads(data) # noqa: S301 - return _RestrictedUnpickler(io.BytesIO(data)).load() + If ADK_ALLOW_UNSAFE_V0_PICKLE=1 is set, falls back to unrestricted + pickle.loads() for compatibility. A deprecation warning is logged. + """ + if os.environ.get("ADK_ALLOW_UNSAFE_V0_PICKLE") == "1": + logger.warning( + "ADK_ALLOW_UNSAFE_V0_PICKLE is set - using unrestricted " + "pickle.loads(). This is unsafe and will be removed in a " + "future release. Migrate to the v1 JSON schema." + ) + return pickle.loads(data) # noqa: S301 + return _RestrictedUnpickler(io.BytesIO(data)).load() From 1407e8dfd816cd5d4f578682c9f59cfa9e767c98 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Sat, 9 May 2026 10:10:21 +0300 Subject: [PATCH 7/8] style: run pyink + isort on v0.py Ref: google/adk-python#5634 --- src/google/adk/sessions/schemas/v0.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/google/adk/sessions/schemas/v0.py b/src/google/adk/sessions/schemas/v0.py index 62ee72b607..5b378cb7a8 100644 --- a/src/google/adk/sessions/schemas/v0.py +++ b/src/google/adk/sessions/schemas/v0.py @@ -31,8 +31,6 @@ import json import logging import pickle - -from ._safe_unpickle import safe_loads as _safe_pickle_loads from typing import Any from typing import Optional @@ -59,6 +57,7 @@ from ...events.event import Event from ...events.event_actions import EventActions from ..session import Session +from ._safe_unpickle import safe_loads as _safe_pickle_loads from .shared import DEFAULT_MAX_KEY_LENGTH from .shared import DEFAULT_MAX_VARCHAR_LENGTH from .shared import DynamicJSON From ac60e57eef8de0fec0587f8d4c938931e9031d39 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Sat, 9 May 2026 10:10:22 +0300 Subject: [PATCH 8/8] style: run pyink + isort on migrate_from_sqlalchemy_pickle.py Ref: google/adk-python#5634 --- .../adk/sessions/migration/migrate_from_sqlalchemy_pickle.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py index 48b9c6cc41..ec834da52d 100644 --- a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py +++ b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py @@ -21,8 +21,6 @@ import json import logging import pickle - -from google.adk.sessions.schemas._safe_unpickle import safe_loads as _safe_pickle_loads import sys from typing import Any @@ -31,6 +29,7 @@ from google.adk.sessions import _session_util from google.adk.sessions.migration import _schema_check_utils from google.adk.sessions.schemas import v1 +from google.adk.sessions.schemas._safe_unpickle import safe_loads as _safe_pickle_loads from google.genai import types import sqlalchemy from sqlalchemy import create_engine