fix(migration): restrict unpickling of v0 actions blobs

Port of GitHub PR: 9db48ce

Restrict unpickling of v0 actions blobs during session migration to prevent unsafe deserialization.
Add --allow-unsafe-unpickling option to allow unsafe unpickling for trusted databases.

Co-authored-by: Yifan Wang <wanyif@google.com>
PiperOrigin-RevId: 927430066

Author: wyf7107

src/google/adk/cli/cli_tools_click.py

@@ -2205,15 +2205,33 @@ def migrate():
     default="INFO",
     help="Optional. Set the logging level",
 )
+@click.option(  # type: ignore[untyped-decorator]
+    "--allow-unsafe-unpickling",
+    "--allow_unsafe_unpickling",
+    is_flag=True,
+    default=False,
+    help=(
+        "Optional. Allow unsafe pickle loading for trusted legacy session"
+        " databases."
+    ),
+)
 def cli_migrate_session(
-    *, source_db_url: str, dest_db_url: str, log_level: str
+    *,
+    source_db_url: str,
+    dest_db_url: str,
+    log_level: str,
+    allow_unsafe_unpickling: bool,
 ):
   """Migrates a session database to the latest schema version."""
   logs.setup_adk_logger(getattr(logging, log_level.upper()))
   try:
     from ..sessions.migration import migration_runner
 
-    migration_runner.upgrade(source_db_url, dest_db_url)
+    migration_runner.upgrade(
+        source_db_url,
+        dest_db_url,
+        allow_unsafe_unpickling=allow_unsafe_unpickling,
+    )
     click.secho("Migration check and upgrade process finished.", fg="green")
   except Exception as e:
     click.secho(f"Migration failed: {e}", fg="red", err=True)

src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py

@@ -18,6 +18,7 @@
 import argparse
 from datetime import datetime
 from datetime import timezone
+import io
 import json
 import logging
 import pickle
@@ -37,6 +38,93 @@
 
 logger = logging.getLogger("google_adk." + __name__)
 
+_ALLOWED_PICKLE_GLOBALS = {
+    # Basic types/containers
+    ("builtins", "dict"),
+    ("builtins", "list"),
+    ("builtins", "set"),
+    ("builtins", "tuple"),
+    ("builtins", "str"),
+    ("builtins", "bytes"),
+    ("builtins", "bytearray"),
+    ("builtins", "int"),
+    ("builtins", "float"),
+    ("builtins", "bool"),
+    ("datetime", "datetime"),
+    ("datetime", "timedelta"),
+    ("datetime", "timezone"),
+    # Expected pickled payload for v0 session schema events.
+    ("fastapi.openapi.models", "APIKey"),
+    ("fastapi.openapi.models", "APIKeyIn"),
+    ("fastapi.openapi.models", "HTTPBase"),
+    ("fastapi.openapi.models", "HTTPBearer"),
+    ("fastapi.openapi.models", "OAuth2"),
+    ("fastapi.openapi.models", "OAuthFlow"),
+    ("fastapi.openapi.models", "OAuthFlowAuthorizationCode"),
+    ("fastapi.openapi.models", "OAuthFlowClientCredentials"),
+    ("fastapi.openapi.models", "OAuthFlowImplicit"),
+    ("fastapi.openapi.models", "OAuthFlowPassword"),
+    ("fastapi.openapi.models", "OAuthFlows"),
+    ("fastapi.openapi.models", "OpenIdConnect"),
+    ("fastapi.openapi.models", "SecurityBase"),
+    ("fastapi.openapi.models", "SecurityScheme"),
+    ("fastapi.openapi.models", "SecuritySchemeType"),
+    ("google.adk.auth.auth_credential", "AuthCredential"),
+    ("google.adk.auth.auth_credential", "AuthCredentialTypes"),
+    ("google.adk.auth.auth_credential", "HttpAuth"),
+    ("google.adk.auth.auth_credential", "HttpCredentials"),
+    ("google.adk.auth.auth_credential", "OAuth2Auth"),
+    ("google.adk.auth.auth_credential", "ServiceAccountCredential"),
+    ("google.adk.auth.auth_schemes", "CustomAuthScheme"),
+    ("google.adk.auth.auth_schemes", "ExtendedOAuth2"),
+    ("google.adk.auth.auth_schemes", "OAuthGrantType"),
+    ("google.adk.auth.auth_schemes", "OpenIdConnectWithConfig"),
+    ("google.adk.auth.auth_tool", "AuthConfig"),
+    ("google.adk.events.event_actions", "EventActions"),
+    ("google.adk.events.event_actions", "EventCompaction"),
+    ("google.adk.events.ui_widget", "UiWidget"),
+    ("google.adk.tools.tool_confirmation", "ToolConfirmation"),
+    ("google.genai.types", "Blob"),
+    ("google.genai.types", "CodeExecutionResult"),
+    ("google.genai.types", "Content"),
+    ("google.genai.types", "ExecutableCode"),
+    ("google.genai.types", "FileData"),
+    ("google.genai.types", "FunctionCall"),
+    ("google.genai.types", "FunctionResponse"),
+    ("google.genai.types", "FunctionResponseBlob"),
+    ("google.genai.types", "FunctionResponseFileData"),
+    ("google.genai.types", "FunctionResponsePart"),
+    ("google.genai.types", "Part"),
+    ("google.genai.types", "PartMediaResolution"),
+    ("google.genai.types", "VideoMetadata"),
+}
+
+
+class _RestrictedUnpickler(pickle.Unpickler):
+  """Restricted unpickler for migrating legacy v0 schema actions.
+
+  The v0 session schema stored `EventActions` as a pickled blob. During
+  migration we treat the raw bytes read from the source DB as untrusted input
+  and only allow the minimum set of safe globals needed to reconstruct
+  `EventActions`.
+  """
+
+  def find_class(self, module: str, name: str) -> Any:  # noqa: ANN001
+    if (module, name) in _ALLOWED_PICKLE_GLOBALS:
+      return super().find_class(module, name)
+    raise pickle.UnpicklingError(
+        f"Blocked global during migration unpickle: {module}.{name}"
+    )
+
+
+def _restricted_pickle_loads(
+    data: bytes, *, allow_unsafe_unpickling: bool = False
+) -> Any:
+  """Load a pickle payload using the restricted unpickler by default."""
+  if allow_unsafe_unpickling:
+    return pickle.loads(data)
+  return _RestrictedUnpickler(io.BytesIO(data)).load()
+
 
 def _to_datetime_obj(val: Any) -> datetime | Any:
   """Converts string to datetime if needed."""
@@ -51,15 +139,19 @@ def _to_datetime_obj(val: Any) -> datetime | Any:
   return val
 
 
-def _row_to_event(row: dict) -> Event:
+def _row_to_event(
+    row: dict[str, Any], *, allow_unsafe_unpickling: bool = False
+) -> Event:
   """Converts event row (dict) to event object, handling missing columns and deserializing."""
 
   actions_val = row.get("actions")
   actions = None
   if actions_val is not None:
     try:
       if isinstance(actions_val, bytes):
-        actions = pickle.loads(actions_val)
+        actions = _restricted_pickle_loads(
+            actions_val, allow_unsafe_unpickling=allow_unsafe_unpickling
+        )
       else:  # for spanner - it might return object directly
         actions = actions_val
     except Exception as e:
@@ -75,17 +167,25 @@ def _row_to_event(row: dict) -> Event:
   else:
     actions = EventActions()
 
-  def _safe_json_load(val):
-    data = None
+  def _safe_json_load(val: Any) -> dict[str, Any] | None:
     if isinstance(val, str):
       try:
         data = json.loads(val)
       except json.JSONDecodeError:
         logger.warning(f"Failed to decode JSON for event {row.get('id')}")
         return None
     elif isinstance(val, dict):
-      data = val  # for postgres JSONB
-    return data
+      return val  # for postgres JSONB
+    else:
+      return None
+
+    if isinstance(data, dict):
+      return data
+    logger.warning(
+        f"Expected JSON object for event {row.get('id')}, got"
+        f" {type(data).__name__}."
+    )
+    return None
 
   content_dict = _safe_json_load(row.get("content"))
   grounding_metadata_dict = _safe_json_load(row.get("grounding_metadata"))
@@ -147,23 +247,31 @@ def _safe_json_load(val):
   )
 
 
-def _get_state_dict(state_val: Any) -> dict:
+def _get_state_dict(state_val: Any) -> dict[str, Any]:
   """Safely load dict from JSON string or return dict if already dict."""
   if isinstance(state_val, dict):
     return state_val
   if isinstance(state_val, str):
     try:
-      return json.loads(state_val)
+      data = json.loads(state_val)
     except json.JSONDecodeError:
       logger.warning(
           "Failed to parse state JSON string, defaulting to empty dict."
       )
       return {}
+    if isinstance(data, dict):
+      return data
+    logger.warning("State JSON was not an object, defaulting to empty dict.")
+    return {}
   return {}
 
 
 # --- Migration Logic ---
-def migrate(source_db_url: str, dest_db_url: str):
+def migrate(
+    source_db_url: str,
+    dest_db_url: str,
+    allow_unsafe_unpickling: bool = False,
+) -> None:
   """Migrates data from old pickle schema to new JSON schema."""
   # Convert async driver URLs to sync URLs for SQLAlchemy's synchronous engine.
   # This allows users to provide URLs like 'postgresql+asyncpg://...' and have
@@ -172,6 +280,11 @@ def migrate(source_db_url: str, dest_db_url: str):
   dest_sync_url = _schema_check_utils.to_sync_url(dest_db_url)
 
   logger.info(f"Connecting to source database: {source_db_url}")
+  if allow_unsafe_unpickling:
+    logger.warning(
+        "Unsafe pickle migration mode is enabled. Only use this with a trusted"
+        " source database."
+    )
   try:
     source_engine = create_engine(source_sync_url)
     SourceSession = sessionmaker(bind=source_engine)
@@ -265,7 +378,10 @@ def migrate(source_db_url: str, dest_db_url: str):
             text("SELECT * FROM events")
         ).mappings():
           try:
-            event_obj = _row_to_event(dict(row))
+            event_obj = _row_to_event(
+                dict(row),
+                allow_unsafe_unpickling=allow_unsafe_unpickling,
+            )
             new_event = v1.StorageEvent(
                 id=event_obj.id,
                 app_name=row["app_name"],
@@ -309,9 +425,22 @@ def migrate(source_db_url: str, dest_db_url: str):
       required=True,
       help="SQLAlchemy URL of destination database",
   )
+  parser.add_argument(
+      "--allow_unsafe_unpickling",
+      "--allow-unsafe-unpickling",
+      action="store_true",
+      help=(
+          "Allow legacy pickle payloads to use Python's unsafe pickle loader."
+          " Only use this with a trusted source database."
+      ),
+  )
   args = parser.parse_args()
   try:
-    migrate(args.source_db_url, args.dest_db_url)
+    migrate(
+        args.source_db_url,
+        args.dest_db_url,
+        allow_unsafe_unpickling=args.allow_unsafe_unpickling,
+    )
   except Exception as e:
     logger.error(f"Migration failed: {e}")
     sys.exit(1)

src/google/adk/sessions/migration/migration_runner.py

@@ -42,7 +42,11 @@
 LATEST_VERSION = _schema_check_utils.LATEST_SCHEMA_VERSION
 
 
-def upgrade(source_db_url: str, dest_db_url: str):
+def upgrade(
+    source_db_url: str,
+    dest_db_url: str,
+    allow_unsafe_unpickling: bool = False,
+) -> None:
   """Migrates a database from its current version to the latest version.
 
   If the source database schema is older than the latest version, this
@@ -61,6 +65,9 @@ def upgrade(source_db_url: str, dest_db_url: str):
     source_db_url: The SQLAlchemy URL of the database to migrate from.
     dest_db_url: The SQLAlchemy URL of the database to migrate to. This must be
       different from source_db_url.
+    allow_unsafe_unpickling: If true, use Python's unsafe pickle loader for the
+      legacy pickle migration step. Only use this with a trusted source
+      database.
 
   Raises:
     RuntimeError: If source_db_url and dest_db_url are the same, or if no
@@ -113,7 +120,14 @@ def upgrade(source_db_url: str, dest_db_url: str):
       logger.info(
           f"Migrating from {in_url} to {out_url} (schema v{end_version})..."
       )
-      migrate_func(in_url, out_url)
+      if migrate_func is migrate_from_sqlalchemy_pickle.migrate:
+        migrate_func(
+            in_url,
+            out_url,
+            allow_unsafe_unpickling=allow_unsafe_unpickling,
+        )
+      else:
+        migrate_func(in_url, out_url)
       logger.info("Finished migration step to schema %s.", end_version)
       # The output of this step becomes the input for the next step.
       in_url = out_url

tests/unittests/cli/utils/test_cli_tools_click.py

@@ -1193,6 +1193,53 @@ def test_cli_web_passes_service_uris(
   assert called_kwargs.get("memory_service_uri") == "rag://mycorpus"
 
 
+@pytest.mark.parametrize(
+    "flag",
+    ["--allow-unsafe-unpickling", "--allow_unsafe_unpickling"],
+)
+def test_cli_migrate_session_allows_unsafe_unpickling_flag(
+    monkeypatch: pytest.MonkeyPatch, flag: str
+) -> None:
+  calls: list[dict[str, Any]] = []
+
+  def fake_upgrade(
+      source_db_url: str,
+      dest_db_url: str,
+      *,
+      allow_unsafe_unpickling: bool = False,
+  ) -> None:
+    calls.append({
+        "source_db_url": source_db_url,
+        "dest_db_url": dest_db_url,
+        "allow_unsafe_unpickling": allow_unsafe_unpickling,
+    })
+
+  monkeypatch.setattr(
+      "google.adk.sessions.migration.migration_runner.upgrade",
+      fake_upgrade,
+  )
+
+  result = CliRunner().invoke(
+      cli_tools_click.main,
+      [
+          "migrate",
+          "session",
+          "--source_db_url",
+          "sqlite:///source.db",
+          "--dest_db_url",
+          "sqlite:///dest.db",
+          flag,
+      ],
+  )
+
+  assert result.exit_code == 0, (result.output, repr(result.exception))
+  assert calls == [{
+      "source_db_url": "sqlite:///source.db",
+      "dest_db_url": "sqlite:///dest.db",
+      "allow_unsafe_unpickling": True,
+  }]
+
+
 def test_cli_eval_with_eval_set_file_path(
     mock_load_eval_set_from_file,
     mock_get_root_agent,