Core: Add restricted_dumps helper (#5117)

* Add pickling helper that check unpicklability

* Add test condition and generation error handling

* Fix incorrect call and make imports consistent

* Fix newline padding

* Change PicklingError to directly caused by UnpicklingError

Co-authored-by: Doug Hoskisson <beauxq@users.noreply.github.com>

* Revert to `pickle.dumps` for decompressed multidata

* Fix import order

* Restore pickle import in main

* Re-add for multidata in Main

* Remove multisave checks

* Update MultiServer.py

* Update customserver.py

---------

Co-authored-by: Doug Hoskisson <beauxq@users.noreply.github.com>
Co-authored-by: NewSoupVi <57900059+NewSoupVi@users.noreply.github.com>

Main.py

@@ -2,7 +2,6 @@ import collections
 import concurrent.futures
 import logging
 import os
-import pickle
 import tempfile
 import time
 import zipfile
@@ -14,7 +13,7 @@ from Fill import FillError, balance_multiworld_progression, distribute_items_res
     parse_planned_blocks, distribute_planned_blocks, resolve_early_locations_for_planned
 from NetUtils import convert_to_base_types
 from Options import StartInventoryPool
-from Utils import __version__, output_path, version_tuple
+from Utils import __version__, output_path, restricted_dumps, version_tuple
 from settings import get_settings
 from worlds import AutoWorld
 from worlds.generic.Rules import exclusion_rules, locality_rules
@@ -339,7 +338,7 @@ def main(args, seed=None, baked_server_options: dict[str, object] | None = None)
                 for key in ("slot_data", "er_hint_data"):
                     multidata[key] = convert_to_base_types(multidata[key])
 
-                multidata = zlib.compress(pickle.dumps(multidata), 9)
+                multidata = zlib.compress(restricted_dumps(multidata), 9)
 
                 with open(os.path.join(temp_dir, f'{outfilebase}.archipelago'), 'wb') as f:
                     f.write(bytes([3]))  # version of format

MultiServer.py

@@ -546,6 +546,7 @@ class Context:
 
     def _save(self, exit_save: bool = False) -> bool:
         try:
+            # Does not use Utils.restricted_dumps because we'd rather make a save than not make one
             encoded_save = pickle.dumps(self.get_save())
             with open(self.save_filename, "wb") as f:
                 f.write(zlib.compress(encoded_save))

Utils.py

@@ -483,6 +483,18 @@ def restricted_loads(s: bytes) -> Any:
     return RestrictedUnpickler(io.BytesIO(s)).load()
 
 
+def restricted_dumps(obj: Any) -> bytes:
+    """Helper function analogous to pickle.dumps()."""
+    s = pickle.dumps(obj)
+    # Assert that the string can be successfully loaded by restricted_loads
+    try:
+        restricted_loads(s)
+    except pickle.UnpicklingError as e:
+        raise pickle.PicklingError(e) from e
+
+    return s
+
+
 class ByValue:
     """
     Mixin for enums to pickle value instead of name (restores pre-3.11 behavior). Use as left-most parent.

WebHostLib/api/generate.py

@@ -1,11 +1,11 @@
 import json
-import pickle
 from uuid import UUID
 
 from flask import request, session, url_for
 from markupsafe import Markup
 from pony.orm import commit
 
+from Utils import restricted_dumps
 from WebHostLib import app
 from WebHostLib.check import get_yaml_data, roll_options
 from WebHostLib.generate import get_meta
@@ -56,7 +56,7 @@ def generate_api():
                     "detail": results}, 400
         else:
             gen = Generation(
-                options=pickle.dumps({name: vars(options) for name, options in gen_options.items()}),
+                options=restricted_dumps({name: vars(options) for name, options in gen_options.items()}),
                 # convert to json compatible
                 meta=json.dumps(meta), state=STATE_QUEUED,
                 owner=session["_id"])

WebHostLib/customserver.py

@@ -129,7 +129,7 @@ class WebHostContext(Context):
                 else:
                     row = GameDataPackage.get(checksum=game_data["checksum"])
                     if row:  # None if rolled on >= 0.3.9 but uploaded to <= 0.3.8. multidata should be complete
-                        game_data_packages[game] = Utils.restricted_loads(row.data)
+                        game_data_packages[game] = restricted_loads(row.data)
                         continue
                     else:
                         self.logger.warning(f"Did not find game_data_package for {game}: {game_data['checksum']}")
@@ -159,6 +159,7 @@ class WebHostContext(Context):
     @db_session
     def _save(self, exit_save: bool = False) -> bool:
         room = Room.get(id=self.room_id)
+        # Does not use Utils.restricted_dumps because we'd rather make a save than not make one
         room.multisave = pickle.dumps(self.get_save())
         # saving only occurs on activity, so we can "abuse" this information to mark this as last_activity
         if not exit_save:  # we don't want to count a shutdown as activity, which would restart the server again

WebHostLib/generate.py

@@ -1,11 +1,11 @@
 import concurrent.futures
 import json
 import os
-import pickle
 import random
 import tempfile
 import zipfile
 from collections import Counter
+from pickle import PicklingError
 from typing import Any
 
 from flask import flash, redirect, render_template, request, session, url_for
@@ -14,7 +14,7 @@ from pony.orm import commit, db_session
 from BaseClasses import get_seed, seeddigits
 from Generate import PlandoOptions, handle_name
 from Main import main as ERmain
-from Utils import __version__
+from Utils import __version__, restricted_dumps
 from WebHostLib import app
 from settings import ServerOptions, GeneratorOptions
 from worlds.alttp.EntranceRandomizer import parse_arguments
@@ -83,12 +83,18 @@ def start_generation(options: dict[str, dict | str], meta: dict[str, Any]):
               f"If you have a larger group, please generate it yourself and upload it.")
         return redirect(url_for(request.endpoint, **(request.view_args or {})))
     elif len(gen_options) >= app.config["JOB_THRESHOLD"]:
-        gen = Generation(
-            options=pickle.dumps({name: vars(options) for name, options in gen_options.items()}),
-            # convert to json compatible
-            meta=json.dumps(meta),
-            state=STATE_QUEUED,
-            owner=session["_id"])
+        try:
+            gen = Generation(
+                options=restricted_dumps({name: vars(options) for name, options in gen_options.items()}),
+                # convert to json compatible
+                meta=json.dumps(meta),
+                state=STATE_QUEUED,
+                owner=session["_id"])
+        except PicklingError as e:
+            from .autolauncher import handle_generation_failure
+            handle_generation_failure(e)
+            return render_template("seedError.html", seed_error=("PicklingError: " + str(e)))
+
         commit()
 
         return redirect(url_for("wait_seed", seed=gen.id))

WebHostLib/upload.py

@@ -1,4 +1,3 @@
-import base64
 import json
 import pickle
 import typing

test/general/test_options.py

@@ -1,7 +1,8 @@
 import unittest
 
-from BaseClasses import MultiWorld, PlandoOptions
-from Options import ItemLinks
+from BaseClasses import PlandoOptions
+from Options import ItemLinks, Choice
+from Utils import restricted_dumps
 from worlds.AutoWorld import AutoWorldRegister
 
 
@@ -73,9 +74,10 @@ class TestOptions(unittest.TestCase):
 
     def test_pickle_dumps(self):
         """Test options can be pickled into database for WebHost generation"""
-        import pickle
         for gamename, world_type in AutoWorldRegister.world_types.items():
             if not world_type.hidden:
                 for option_key, option in world_type.options_dataclass.type_hints.items():
                     with self.subTest(game=gamename, option=option_key):
-                        pickle.dumps(option.from_any(option.default))
+                        restricted_dumps(option.from_any(option.default))
+                        if issubclass(option, Choice) and option.default in option.name_lookup:
+                            restricted_dumps(option.from_text(option.name_lookup[option.default]))