From 9519f81f610356517e2f99831a5cb65ff25f3aba Mon Sep 17 00:00:00 2001
From: iluem <57590186+Qhaoduoyu@users.noreply.github.com>
Date: Sun, 21 Apr 2024 19:04:00 +0800
Subject: [PATCH] Merge pull request from GHSA-mvrw-h7rc-22r8
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 注释调试if分支

* Improve objload security

* Update README.md

* support pdf url for latex translation

---------

Co-authored-by: binary-husky <96192199+binary-husky@users.noreply.github.com>
Co-authored-by: binary-husky <qingxu.fu@outlook.com>
---
 toolbox.py | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/toolbox.py b/toolbox.py
index 6cc62e7c3..aa6b9d099 100644
--- a/toolbox.py
+++ b/toolbox.py
@@ -9,6 +9,7 @@
 import glob
 import logging
 import uuid
+import pickle
 from functools import wraps
 from shared_utils.config_loader import get_conf
 from shared_utils.config_loader import set_conf
@@ -867,9 +868,20 @@ def __exit__(self, exc_type, exc_value, traceback):
             os.environ.pop("HTTPS_PROXY")
         return
 
+class SafeUnpickler(pickle.Unpickler):
+    # 定义允许的安全类
+    safe_classes = {
+        # 在这里添加其他安全的类
+    }
+
+    def find_class(self, module, name):
+        # 只允许特定的类进行反序列化
+        if f'{module}.{name}' in self.safe_classes:
+            return self.safe_classes[f'{module}.{name}']
+        # 如果尝试加载未授权的类，则抛出异常
+        raise pickle.UnpicklingError(f"Attempted to deserialize unauthorized class '{name}' from module '{module}'")
 
 def objdump(obj, file="objdump.tmp"):
-    import pickle
 
     with open(file, "wb+") as f:
         pickle.dump(obj, f)
@@ -877,13 +889,13 @@ def objdump(obj, file="objdump.tmp"):
 
 
 def objload(file="objdump.tmp"):
-    import pickle, os
+    import os
 
     if not os.path.exists(file):
         return
     with open(file, "rb") as f:
-        return pickle.load(f)
-
+        unpickler = SafeUnpickler(f)
+        return unpickler.load()
 
 def Singleton(cls):
     """