diff --git a/Main.py b/Main.py index 8bbbd74b..fd7d8d9a 100644 --- a/Main.py +++ b/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 diff --git a/MultiServer.py b/MultiServer.py index 1f421aaa..59960a2a 100644 --- a/MultiServer.py +++ b/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)) diff --git a/Utils.py b/Utils.py index 9c117109..abf359f4 100644 --- a/Utils.py +++ b/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. diff --git a/WebHostLib/api/generate.py b/WebHostLib/api/generate.py index 5a66d1e6..7bcbdbcf 100644 --- a/WebHostLib/api/generate.py +++ b/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"]) diff --git a/WebHostLib/customserver.py b/WebHostLib/customserver.py index 2ebb40d6..156c1252 100644 --- a/WebHostLib/customserver.py +++ b/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 diff --git a/WebHostLib/generate.py b/WebHostLib/generate.py index a84b17a8..02f5a037 100644 --- a/WebHostLib/generate.py +++ b/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)) diff --git a/WebHostLib/upload.py b/WebHostLib/upload.py index ee4ba6a5..4c5d411d 100644 --- a/WebHostLib/upload.py +++ b/WebHostLib/upload.py @@ -1,4 +1,3 @@ -import base64 import json import pickle import typing diff --git a/test/general/test_options.py b/test/general/test_options.py index 7a3743e5..d8ce7017 100644 --- a/test/general/test_options.py +++ b/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]))