Security: gate marshal-based function deserialization behind enable_opaque_pickle

pyglove/core/utils/json_conversion.py

@@ -387,10 +387,11 @@ def _type_name(
 def enable_opaque_pickle(enable: bool = True) -> ContextManager[None]:
   """Returns a context manager to enable or disable opaque-object pickle.
 
-  Deserializing ``_OpaqueObject`` uses ``pickle.loads``, which can execute
-  arbitrary code.  It is enabled by default for backward compatibility.
+  Deserializing ``_OpaqueObject`` uses ``pickle.loads`` and deserializing a
+  function from inline ``code`` uses ``marshal.loads``; both can execute
+  arbitrary code.  They are enabled by default for backward compatibility.
   Security-sensitive contexts (e.g. cloud services that process untrusted
-  JSON) should disable it::
+  JSON) should disable them::
 
     with json_conversion.enable_opaque_pickle(False):
       obj = json_conversion.from_json(untrusted_json)
@@ -1107,6 +1108,13 @@ def _fn(json_value: Dict[str, str], **kwargs) -> types.FunctionType:
     del kwargs
     function_name = json_value['name']
     if 'code' in json_value:
+      if not _opaque_pickle_enabled:
+        raise TypeError(
+            'Deserializing a function from inline code is disabled because it '
+            'uses `marshal.loads`, which can execute arbitrary code when the '
+            'JSON comes from an untrusted source. '
+            'Use json_conversion.enable_opaque_pickle() to opt in.'
+        )
       code = marshal.loads(
           base64.decodebytes(json_value['code'].encode('utf-8')))
       defaults = from_json(json_value['defaults'], _typename_resolved=True)

pyglove/core/utils/json_conversion_test.py

@@ -599,6 +599,29 @@ def test_opaque_from_json_works_inside_allow_context(self):
         result = json_conversion._OpaqueObject.from_json(json_value)
     self.assertEqual(result, x)
 
+  def test_function_code_deserialization_gated_when_disabled(self):
+    """Inline-code (marshal) function deser must be gated when disabled."""
+    # A lambda serializes via the inline-`code` (marshal.dumps) path.
+    fn = lambda x: x + 1
+    json_value = json_conversion.to_json(fn)
+    self.assertIn('code', json_value)
+    # When opaque pickle is disabled, untrusted JSON must not be turned into an
+    # executable function (the marshal.loads -> FunctionType RCE path).
+    with json_conversion.enable_opaque_pickle(False):
+      with self.assertRaisesRegex(TypeError, 'disabled'):
+        json_conversion.from_json(json_value)
+
+  def test_function_code_deserialization_works_by_default(self):
+    """Inline-code function deser works by default (backward compatible)."""
+    fn = lambda x: x + 1
+    restored = json_conversion.from_json(json_conversion.to_json(fn))
+    self.assertEqual(restored(3), 4)
+    # Re-enabling inside a disabled scope works too.
+    with json_conversion.enable_opaque_pickle(False):
+      with json_conversion.enable_opaque_pickle(True):
+        restored2 = json_conversion.from_json(json_conversion.to_json(fn))
+    self.assertEqual(restored2(3), 4)
+
   def test_enable_opaque_pickle_restores_on_exception(self):
     """Flag must be restored even if the body raises."""
     self.assertTrue(json_conversion._opaque_pickle_enabled)