diff --git a/Dockerfile b/Dockerfile index 3cb28c0bcc..4259bcd3a9 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -FROM python:3.6 +FROM python:3.9 RUN apt-get -y update && apt-get -y install ffmpeg # RUN apt-get -y update && apt-get -y install git wget python-dev python3-dev libopenmpi-dev python-pip zlib1g-dev cmake python-opencv diff --git a/baselines/common/vec_env/vec_env.py b/baselines/common/vec_env/vec_env.py index fc6098e2ec..5214bda7cd 100644 --- a/baselines/common/vec_env/vec_env.py +++ b/baselines/common/vec_env/vec_env.py @@ -1,9 +1,30 @@ import contextlib import os from abc import ABC, abstractmethod - +import io +import builtins from baselines.common.tile_images import tile_images - +safe_builtins = { + 'range', + 'complex', + 'set', + 'frozenset', + 'slice', +} + +class RestrictedUnpickler(pickle.Unpickler): + + def find_class(self, module, name): + """Only allow safe classes from builtins""" + if module == "builtins" and name in safe_builtins: + return getattr(builtins, name) + """Forbid everything else""" + raise pickle.UnpicklingError("global '%s.%s' is forbidden" % + (module, name)) + +def restricted_loads(s): + """Helper function analogous to pickle.loads()""" + return RestrictedUnpickler(io.BytesIO(s)).load() class AlreadySteppingError(Exception): """ Raised when an asynchronous step is running while @@ -200,7 +221,8 @@ def __getstate__(self): return cloudpickle.dumps(self.x) def __setstate__(self, ob): - import pickle + #import pickle + restricted_loads(ob) self.x = pickle.loads(ob)