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>

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