diff --git a/pyglove/core/utils/json_conversion.py b/pyglove/core/utils/json_conversion.py index 4a1ec21..5198e2d 100644 --- a/pyglove/core/utils/json_conversion.py +++ b/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) diff --git a/pyglove/core/utils/json_conversion_test.py b/pyglove/core/utils/json_conversion_test.py index 0265058..50aa2d6 100644 --- a/pyglove/core/utils/json_conversion_test.py +++ b/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)