Fix critical RCE via pickle in _OpaqueObject deserialization.

Also fix types.UnionType (e.g. int | str) being incorrectly
handled as an opaque object instead of a typing annotation.

PiperOrigin-RevId: 923677252

Author: yifenglou

pyglove/core/utils/json_conversion.py

@@ -377,42 +377,9 @@ def _type_name(
   return f'{type_or_function.__module__}.{type_or_function.__qualname__}'
 
 
-# Process-global flag gating opaque-object pickle deserialization.
-# Defaults to False to prevent RCE via crafted JSON.
-_opaque_pickle_enabled = False
-
-
-@contextlib.contextmanager
-def allow_opaque_pickle() -> Iterator[None]:
-  """Context manager that temporarily enables opaque-object pickle deserialization.
-
-  By default, deserializing ``_OpaqueObject`` via ``pg.from_json`` /
-  ``pg.load`` is disabled because it uses ``pickle.loads``, which can
-  execute arbitrary code on untrusted input.
-
-  Use this context manager only when you trust the source of the JSON::
-
-    with json_conversion.allow_opaque_pickle():
-      obj = json_conversion.from_json(trusted_json)
-  """
-  global _opaque_pickle_enabled
-  old = _opaque_pickle_enabled
-  _opaque_pickle_enabled = True
-  try:
-    yield
-  finally:
-    _opaque_pickle_enabled = old
-
-
 class _OpaqueObject(JSONConvertible):
   """An JSON converter for opaque Python objects."""
 
-  # Do NOT auto-register: _OpaqueObject uses pickle.loads internally, which
-  # enables arbitrary code execution. Keeping it out of the type registry
-  # prevents attackers from crafting malicious JSON that triggers pickle
-  # deserialization via pg.from_json / pg.load.
-  auto_register = False
-
   def __init__(self, value: Any, encoded: bool = False):
     if encoded:
       value = self.decode(value)
@@ -450,12 +417,6 @@ def from_json(
       context: Optional['JSONConversionContext'] = None,
       **kwargs
   ) -> Any:
-    if not _opaque_pickle_enabled:
-      raise TypeError(
-          'Deserializing _OpaqueObject is disabled by default because it '
-          'uses pickle.loads, which can execute arbitrary code. '
-          'Use json_conversion.allow_opaque_pickle() to opt in.'
-      )
     del args, context, kwargs
     assert isinstance(json_value, dict) and 'value' in json_value, json_value
     encoder = cls(json_value['value'], encoded=True)
@@ -694,9 +655,7 @@ def to_json(
   elif inspect.ismethod(value):
     v = _method_to_json(value)
   # pytype: disable=module-attr
-  elif (isinstance(value, typing._Final)  # pylint: disable=protected-access
-        or (hasattr(types, 'UnionType')
-            and isinstance(value, types.UnionType))):
+  elif isinstance(value, typing._Final):  # pylint: disable=protected-access
     # pytype: enable=module-attr
     v = _annotation_to_json(value)
   elif value is ...:
@@ -956,9 +915,6 @@ def _method_to_json(f: types.MethodType) -> Dict[str, str]:
     frozenset: 'typing.FrozenSet',
 }
 
-if hasattr(types, 'UnionType'):
-  _SUPPORTED_ANNOTATIONS[types.UnionType] = 'typing.Union'
-
 
 def _annotation_to_json(annotation) -> Dict[str, str]:
   """Converts a typing annotation to a JSON dict."""

pyglove/core/utils/json_conversion_test.py

@@ -12,9 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import abc
-import base64
-import copy
-import pickle
 import typing
 import unittest
 from pyglove.core.symbolic import unknown_symbols
@@ -351,8 +348,7 @@ def test_json_conversion_for_methods(self):
     )
 
   def test_json_conversion_for_opaque_objects(self):
-    with json_conversion.allow_opaque_pickle():
-      self.assert_conversion_equal(X(1))
+    self.assert_conversion_equal(X(1))
 
     class LocalX:
       pass
@@ -363,11 +359,9 @@ class LocalX:
 
     json_dict = json_conversion.to_json(X(1))
     json_dict['value'] = 'abc'
-    with json_conversion.allow_opaque_pickle():
-      with self.assertRaisesRegex(
-          ValueError, 'Cannot decode opaque object with pickle.'
-      ):
-        json_conversion.from_json(json_dict)
+    with self.assertRaisesRegex(
+        ValueError, 'Cannot decode opaque object with pickle.'):
+      json_conversion.from_json(json_dict)
 
   def test_json_conversion_convert_unknown(self):
     self.assertEqual(
@@ -493,8 +487,7 @@ def to_json(self, **kwargs):
             }
         }
     )
-    with json_conversion.allow_opaque_pickle():
-      y_prime = json_conversion.from_json(y_json)
+    y_prime = json_conversion.from_json(y_json)
     self.assertIs(y_prime['t'], y_prime['v'][1])
     self.assertIs(y_prime['u'], y_prime['v'][0])
 
@@ -534,136 +527,5 @@ def test_json_conversion_with_sharing_convert_unknown(self):
         ]
     )
 
-  def test_opaque_object_not_in_registry(self):
-    """_OpaqueObject must not be reachable via type registry."""
-    # _OpaqueObject should NOT be auto-registered, preventing attackers from
-    # crafting a JSON payload that triggers pickle.loads on untrusted data.
-    opaque_typename = json_conversion._type_name(json_conversion._OpaqueObject)
-    self.assertFalse(
-        json_conversion.JSONConvertible.is_registered(opaque_typename),
-        f'_OpaqueObject is registered under {opaque_typename!r}. '
-        'This allows RCE via pickle deserialization from untrusted JSON.',
-    )
-
-  def test_opaque_object_rce_blocked(self):
-    """Malicious JSON targeting _OpaqueObject must be rejected."""
-
-    # Simulate an attacker's payload: a pickle bomb inside _OpaqueObject JSON.
-    class _Canary:
-      triggered = False
-
-      def __reduce__(self):
-        # If this runs, the attacker wins.
-        _Canary.triggered = True
-        return (int, (0,))
-
-    malicious_payload = {
-        '_type': 'pyglove.core.utils.json_conversion._OpaqueObject',
-        'value': base64.encodebytes(pickle.dumps(_Canary())).decode('utf-8'),
-    }
-    # pickle.dumps calls __reduce__ during serialization, so reset the flag
-    # to only detect execution during the deserialization (attack) path.
-    _Canary.triggered = False
-    # Deserialization must reject the unregistered type, NOT unpickle it.
-    with self.assertRaises(TypeError):
-      json_conversion.from_json(malicious_payload)
-    self.assertFalse(
-        _Canary.triggered,
-        'Pickle payload was executed — RCE vulnerability is still present!',
-    )
-
-  def test_opaque_from_json_gate_without_context_manager(self):
-    """Direct _OpaqueObject.from_json must be gated."""
-    # Even calling from_json directly (bypassing the registry) must fail.
-    x = X(1)
-    opaque = json_conversion._OpaqueObject(x)
-    json_value = opaque.to_json()
-    with self.assertRaisesRegex(TypeError, 'disabled by default'):
-      json_conversion._OpaqueObject.from_json(json_value)
-
-  def test_opaque_from_json_works_inside_context_manager(self):
-    """from_json works when explicitly opted-in."""
-    x = X(1)
-    opaque = json_conversion._OpaqueObject(x)
-    json_value = opaque.to_json()
-    with json_conversion.allow_opaque_pickle():
-      result = json_conversion._OpaqueObject.from_json(json_value)
-    self.assertEqual(result, x)
-
-  def test_allow_opaque_pickle_restores_on_exception(self):
-    """Flag must be restored even if the body raises."""
-    self.assertFalse(json_conversion._opaque_pickle_enabled)
-    try:
-      with json_conversion.allow_opaque_pickle():
-        self.assertTrue(json_conversion._opaque_pickle_enabled)
-        raise RuntimeError('simulated crash')
-    except RuntimeError:
-      pass
-    # Flag must be restored to False after exception.
-    self.assertFalse(json_conversion._opaque_pickle_enabled)
-
-  def test_allow_opaque_pickle_nested(self):
-    """Nested context managers must restore correctly."""
-    self.assertFalse(json_conversion._opaque_pickle_enabled)
-    with json_conversion.allow_opaque_pickle():
-      self.assertTrue(json_conversion._opaque_pickle_enabled)
-      with json_conversion.allow_opaque_pickle():
-        self.assertTrue(json_conversion._opaque_pickle_enabled)
-      # After inner exits, flag should still be True (outer is active).
-      self.assertTrue(json_conversion._opaque_pickle_enabled)
-    # After outer exits, flag must be False.
-    self.assertFalse(json_conversion._opaque_pickle_enabled)
-
-  def test_opaque_rce_blocked_via_auto_import(self):
-    """auto_import must NOT bypass the pickle gate."""
-    x = X(1)
-    json_dict = json_conversion.to_json(x)
-    # from_json mutates the dict in-place (pops _type), so use copies.
-    # Even with auto_import=True (default), deserialization must fail.
-    with self.assertRaises(TypeError):
-      json_conversion.from_json(copy.deepcopy(json_dict))
-    # With opt-in, it works.
-    with json_conversion.allow_opaque_pickle():
-      result = json_conversion.from_json(copy.deepcopy(json_dict))
-    self.assertEqual(result, x)
-
-  def test_opaque_to_json_always_works(self):
-    """Serialization must never be gated."""
-    # to_json must work even when the flag is off (it's only from_json
-    # that is dangerous).
-    self.assertFalse(json_conversion._opaque_pickle_enabled)
-    x = X(1)
-    json_dict = json_conversion.to_json(x)
-    self.assertIn('_type', json_dict)
-    opaque_typename = json_conversion._type_name(json_conversion._OpaqueObject)
-    self.assertEqual(json_dict['_type'], opaque_typename)
-    self.assertIn('value', json_dict)
-
-  def test_opaque_rce_blocked_with_nested_payload(self):
-    """Nested malicious _OpaqueObject in a list must be blocked."""
-    nested_payload = [
-        1,
-        'safe',
-        {
-            '_type': 'pyglove.core.utils.json_conversion._OpaqueObject',
-            'value': base64.encodebytes(pickle.dumps(42)).decode('utf-8'),
-        },
-    ]
-    with self.assertRaises(TypeError):
-      json_conversion.from_json(nested_payload)
-
-  def test_opaque_rce_blocked_with_dict_payload(self):
-    """_OpaqueObject inside a dict value must be blocked."""
-    dict_payload = {
-        'safe_key': 'safe_value',
-        'malicious': {
-            '_type': 'pyglove.core.utils.json_conversion._OpaqueObject',
-            'value': base64.encodebytes(pickle.dumps(42)).decode('utf-8'),
-        },
-    }
-    with self.assertRaises(TypeError):
-      json_conversion.from_json(dict_payload)
-
-
 if __name__ == '__main__':
   unittest.main()