Remove very old pickle compatibility code modifying old atg package names

petastorm/etl/dataset_metadata.py

@@ -25,7 +25,7 @@
 from six.moves import cPickle as pickle
 
 from petastorm import utils
-from petastorm.etl.legacy import depickle_legacy_package_name_compatible
+from petastorm.etl.safe_pickle import safe_loads
 from petastorm.fs_utils import FilesystemResolver, get_filesystem_and_path_or_paths, get_dataset_path
 from petastorm.unischema import Unischema
 
@@ -380,7 +380,7 @@ def get_schema(dataset):
     # Since we have moved the unischema class around few times, unpickling old schemas will not work. In this case we
     # override the old import path to get backwards compatibility
 
-    schema = depickle_legacy_package_name_compatible(ser_schema)
+    schema = safe_loads(ser_schema)
 
     return schema
 

petastorm/etl/rowgroup_indexing.py

@@ -22,7 +22,7 @@
 
 from petastorm import utils
 from petastorm.etl import dataset_metadata
-from petastorm.etl.legacy import depickle_legacy_package_name_compatible
+from petastorm.etl.safe_pickle import safe_loads
 from petastorm.fs_utils import FilesystemResolver
 
 logger = logging.getLogger(__name__)
@@ -154,5 +154,5 @@ def get_row_group_indexes(dataset):
 
     serialized_indexes = dataset_metadata_dict[ROWGROUPS_INDEX_KEY]
 
-    index_dict = depickle_legacy_package_name_compatible(serialized_indexes)
+    index_dict = safe_loads(serialized_indexes)
     return index_dict

petastorm/etl/safe_pickle.py

@@ -0,0 +1,45 @@
+#  Copyright (c) 2017-2018 Uber Technologies, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import io
+import pickle
+
+# List of packages which symbols are allowed to be loaded
+safe_modules = {
+    "petastorm",
+    "collections",
+    "numpy",
+    "pyspark",
+    "decimal",
+    "builtins",
+    "copy_reg",
+    "__builtin__",
+}
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+
+    def find_class(self, module, name):
+        """Only allow safe classes from builtins"""
+        package_name = module.split(".")[0]
+        if package_name in safe_modules:
+            return super().find_class(module, name)
+        else:
+            # Forbid everything else
+            raise pickle.UnpicklingError(f"global {module, name} is forbidden")
+
+
+def safe_loads(s):
+    """Helper function analogous to pickle.loads()"""
+    return RestrictedUnpickler(io.BytesIO(s)).load()