Security: Remediation of RCE via Insecure Pickle Deserialization and Hardening of I/O URIs

pyglove/core/coding/execution.py

@@ -17,6 +17,7 @@
 import contextlib
 import io
 import multiprocessing
+import msgpack
 import pickle
 import queue
 from typing import Any, Callable, Dict, Optional, Union
@@ -187,7 +188,7 @@ def _call(q, *args, **kwargs):
     def _run():
       r = func(*args, **kwargs)
       try:
-        return pickle.dumps(r)
+        return msgpack.packb(utils.to_json(r))
       except BaseException as e:
         raise errors.SerializationError(
             f'Cannot serialize sandbox result: {r}', e
@@ -218,7 +219,7 @@ def _run():
   if isinstance(x, Exception):
     raise x
   try:
-    return pickle.loads(x)
+    return utils.from_json(msgpack.unpackb(x))
   except Exception as e:
     raise errors.SerializationError(
         'Cannot deserialize the output from sandbox.', e

pyglove/core/io/file_system.py

@@ -836,8 +836,28 @@ def copy(
 #
 
 
-def open(path: Union[str, os.PathLike[str]], mode: str = 'r', **kwargs) -> File:  # pylint:disable=redefined-builtin
+def _check_remote_access(
+    path: Union[str, os.PathLike[str]], allow_remote: bool
+) -> None:
+  """Checks if remote access is allowed for a path."""
+  if not allow_remote:
+    path_str = resolve_path(path)
+    if '://' in path_str and not path_str.startswith('file://'):
+      raise RuntimeError(
+          f'SECURITY ERROR: Remote URI loading is blocked for {path_str!r}. '
+          'By default, PyGlove only allows local file access. If you trust '
+          'the source, use allow_remote=True.'
+      )
+
+
+def open(  # pylint:disable=redefined-builtin
+    path: Union[str, os.PathLike[str]],
+    mode: str = 'r',
+    allow_remote: bool = False,
+    **kwargs,
+) -> File:
   """Opens a file with a path."""
+  _check_remote_access(path, allow_remote)
   return _fs.get(path).open(path, mode, **kwargs)
 
 
@@ -850,9 +870,11 @@ def readfile(
     path: Union[str, os.PathLike[str]],
     mode: str = 'r',
     nonexist_ok: bool = False,
+    allow_remote: bool = False,
     **kwargs,
 ) -> Union[bytes, str, None]:
   """Reads content from a file."""
+  _check_remote_access(path, allow_remote)
   try:
     with _fs.get(path).open(path, mode=mode, **kwargs) as f:
       return f.read()
@@ -867,10 +889,12 @@ def writefile(
     content: Union[str, bytes],
     *,
     mode: str = 'w',
-    perms: Optional[int] = 0o664,   # Default to world-readable.
+    perms: Optional[int] = 0o664,  # Default to world-readable.
+    allow_remote: bool = False,
     **kwargs,
 ) -> None:
   """Writes content to a file."""
+  _check_remote_access(path, allow_remote)
   with _fs.get(path).open(path, mode=mode, **kwargs) as f:
     f.write(content)
   if perms is not None:

pyglove/core/io/sequence.py

@@ -94,6 +94,7 @@ def open(
       perms: Optional[int],
       serializer: Optional[Callable[[Any], Union[bytes, str]]],
       deserializer: Optional[Callable[[Union[bytes, str]], Any]],
+      allow_remote: bool = False,
       **kwargs
   ) -> Sequence:
     """Opens a sequence for reading or writing."""
@@ -148,6 +149,7 @@ def open_sequence(
         Callable[[Union[bytes, str]], Any]
     ] = None,
     make_dirs_if_not_exist: bool = True,
+    allow_remote: bool = False,
 ) -> Sequence:
   """Open sequence for reading or writing.
 
@@ -170,7 +172,8 @@ def open_sequence(
     if make_dirs_if_not_exist:
       file_system.mkdirs(parent_dir, exist_ok=True)
   return _registry.get(path).open(
-      path, mode, perms=perms, serializer=serializer, deserializer=deserializer
+      path, mode, perms=perms, serializer=serializer, deserializer=deserializer,
+      allow_remote=allow_remote,
   )
 
 
@@ -243,6 +246,7 @@ def open(
       perms: Optional[int],
       serializer: Optional[Callable[[Any], Union[bytes, str]]],
       deserializer: Optional[Callable[[Union[bytes, str]], Any]],
+      allow_remote: bool = False,
       **kwargs
   ) -> Sequence:
     """Opens a reader for a sequence."""
@@ -268,11 +272,12 @@ def __init__(
       perms: Optional[int],
       serializer: Optional[Callable[[Any], Union[bytes, str]]],
       deserializer: Optional[Callable[[Union[bytes, str]], Any]],
+      allow_remote: bool = False,
   ) -> None:
     super().__init__(perms, serializer, deserializer)
     self._path = path
     self._mode = mode
-    self._file = file_system.open(path, mode)
+    self._file = file_system.open(path, mode, allow_remote=allow_remote)
 
   def __len__(self):
     raise NotImplementedError(
@@ -311,9 +316,12 @@ def open(
       perms: Optional[int],
       serializer: Optional[Callable[[Any], Union[bytes, str]]],
       deserializer: Optional[Callable[[Union[bytes, str]], Any]],
+      allow_remote: bool = False,
       **kwargs
   ) -> Sequence:
     """Opens a reader for a sequence."""
     del kwargs
-    return LineSequence(path, mode, perms, serializer, deserializer)
+    return LineSequence(
+        path, mode, perms, serializer, deserializer,
+        allow_remote=allow_remote)
 

pyglove/core/symbolic/base.py

@@ -2444,9 +2444,10 @@ def open_jsonl(
 def default_load_handler(
     path: str,
     file_format: Literal['json', 'txt'] = 'json',
+    allow_remote: bool = False,
     **kwargs) -> Any:
   """Default load handler from file."""
-  content = pg_io.readfile(path)
+  content = pg_io.readfile(path, allow_remote=allow_remote)
   if file_format == 'json':
     return from_json_str(content, allow_partial=True, **kwargs)
   elif file_format == 'txt':

pyglove/core/utils/json_conversion.py

@@ -21,6 +21,7 @@
 import importlib
 import inspect
 import marshal
+import msgpack
 import pickle
 import types
 import typing
@@ -380,9 +381,9 @@ def _type_name(
 class _OpaqueObject(JSONConvertible):
   """An JSON converter for opaque Python objects."""
 
-  def __init__(self, value: Any, encoded: bool = False):
+  def __init__(self, value: Any, encoded: bool = False, allow_pickle: bool = False):
     if encoded:
-      value = self.decode(value)
+      value = self.decode(value, allow_pickle=allow_pickle)
     self._value = value
 
   @property
@@ -392,17 +393,25 @@ def value(self) -> Any:
 
   def encode(self, value: Any) -> JSONValueType:
     try:
-      return base64.encodebytes(pickle.dumps(value)).decode('utf-8')
+      return base64.encodebytes(msgpack.packb(value)).decode('utf-8')
     except Exception as e:
       raise ValueError(
-          f'Cannot encode opaque object {value!r} with pickle.') from e
+          f'Cannot encode opaque object {value!r} with msgpack.') from e
 
-  def decode(self, json_value: JSONValueType) -> Any:
+  def decode(self, json_value: JSONValueType, allow_pickle: bool = False) -> Any:
     assert isinstance(json_value, str), json_value
+    data = base64.decodebytes(json_value.encode('utf-8'))
+    if data.startswith(b'\x80'):
+      if not allow_pickle:
+        raise RuntimeError(
+            'SECURITY ERROR: Insecure pickle detected in JSON. For security '
+            'reasons, loading is blocked. If you trust the source, use '
+            'allow_pickle=True.')
+      return pickle.loads(data)
     try:
-      return pickle.loads(base64.decodebytes(json_value.encode('utf-8')))
+      return msgpack.unpackb(data)
     except Exception as e:
-      raise ValueError('Cannot decode opaque object with pickle.') from e
+      raise ValueError('Cannot decode opaque object with msgpack.') from e
 
   def to_json(self, **kwargs) -> JSONValueType:
     return self.to_json_dict({
@@ -417,9 +426,9 @@ def from_json(
       context: Optional['JSONConversionContext'] = None,
       **kwargs
   ) -> Any:
-    del args, context, kwargs
     assert isinstance(json_value, dict) and 'value' in json_value, json_value
-    encoder = cls(json_value['value'], encoded=True)
+    allow_pickle = kwargs.get('allow_pickle', False)
+    encoder = cls(json_value['value'], encoded=True, allow_pickle=allow_pickle)
     return encoder.value
 
 

requirements.txt

@@ -1,4 +1,5 @@
 docstring-parser>=0.12
+msgpack>=1.0.0
 termcolor>=1.1.0
 
 # extras:io