openvinotoolkit/stable-diffusion-webui
1
1
Fork 0
mirror of https://github.com/openvinotoolkit/stable-diffusion-webui.git synced 2024-12-15 07:03:06 +03:00
Code Issues Packages Projects Releases Wiki Activity

become even stricter with pickles

Browse Source 
no pickle shall pass
thank you again, RyotaK
This commit is contained in:
AUTOMATIC 2022-10-11 17:03:00 +03:00
parent a05c824384
commit 66b7d7584f
1 changed files with 17 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
modules/safe.py
View File

@ -10,6 +10,7 @@ import torch
import numpy
import _codecs
import zipfile
import re
# PyTorch 1.13 and later have _TypedStorage renamed to TypedStorage
@ -54,11 +55,27 @@ class RestrictedUnpickler(pickle.Unpickler):
raise pickle.UnpicklingError(f"global '{module}/{name}' is forbidden")
allowed_zip_names = ["archive/data.pkl", "archive/version"]
allowed_zip_names_re = re.compile(r"^archive/data/\d+$")
def check_zip_filenames(filename, names):
for name in names:
if name in allowed_zip_names:
continue
if allowed_zip_names_re.match(name):
continue
raise Exception(f"bad file inside {filename}: {name}")
def check_pt(filename):
try:
# new pytorch format is a zip file
with zipfile.ZipFile(filename) as z:
check_zip_filenames(filename, z.namelist())
with z.open('archive/data.pkl') as file:
unpickler = RestrictedUnpickler(file)
unpickler.load()