Skip to content

Commit cb48bbd

Browse files
Removed file-extension-based arbitrary code execution attack vector (invoke-ai#2946)
# The Problem Pickle files (.pkl, .ckpt, etc) are extremely unsafe as they can be trivially crafted to execute arbitrary code when parsed using `torch.load` Right now the conventional wisdom among ML researchers and users is to simply `not run untrusted pickle files ever` and instead only use Safetensor files, which cannot be injected with arbitrary code. This is very good advice. Unfortunately, **I have discovered a vulnerability inside of InvokeAI that allows an attacker to disguise a pickle file as a safetensor and have the payload execute within InvokeAI.** # How It Works Within `model_manager.py` and `convert_ckpt_to_diffusers.py` there are if-statements that decide which `load` method to use based on the file extension of the model file. The logic (written in a slightly more readable format than it exists in the codebase) is as follows: ``` if Path(file).suffix == '.safetensors': safetensor_load(file) else: unsafe_pickle_load(file) ``` A malicious actor would only need to create an infected .ckpt file, and then rename the extension to something that does not pass the `== '.safetensors'` check, but still appears to a user to be a safetensors file. For example, this might be something like `.Safetensors`, `.SAFETENSORS`, `SafeTensors`, etc. InvokeAI will happily import the file in the Model Manager and execute the payload. # Proof of Concept 1. Create a malicious pickle file. (https://gist.github.com/CodeZombie/27baa20710d976f45fb93928cbcfe368) 2. Rename the `.ckpt` extension to some variation of `.Safetensors`, ensuring there is a capital letter anywhere in the extension (eg. `malicious_pickle.SAFETENSORS`) 3. Import the 'model' like you would normally with any other safetensors file with the Model Manager. 4. Upon trying to select the model in the web ui, it will be loaded (or attempt to be converted to a Diffuser) with `torch.load` and the payload will execute. ![image](https://user-images.githubusercontent.com/466103/224835490-4cf97ff3-41b3-4a31-85df-922cc99042d2.png) # The Fix This pull request changes the logic InvokeAI uses to decide which model loader to use so that the safe behavior is the default. Instead of loading as a pickle if the extension is not exactly `.safetensors`, it will now **always** load as a safetensors file unless the extension is **exactly** `.ckpt`. # Notes: I think support for pickle files should be totally dropped ASAP as a matter of security, but I understand that there are reasons this would be difficult. In the meantime, I think `RestrictedUnpickler` or something similar should be implemented as a replacement for `torch.load`, as this significantly reduces the amount of Python methods that an attacker has to work with when crafting malicious payloads inside a pickle file. Automatic1111 already uses this with some success. (https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/safe.py)
2 parents f9abc6f + a0f47aa commit cb48bbd

2 files changed

Lines changed: 7 additions & 6 deletions

File tree

invokeai/backend/model_management/convert_ckpt_to_diffusers.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1075,9 +1075,10 @@ def load_pipeline_from_original_stable_diffusion_ckpt(
10751075
dlogging.set_verbosity_error()
10761076

10771077
checkpoint = (
1078-
load_file(checkpoint_path)
1079-
if Path(checkpoint_path).suffix == ".safetensors"
1080-
else torch.load(checkpoint_path)
1078+
torch.load(checkpoint_path)
1079+
if Path(checkpoint_path).suffix == ".ckpt"
1080+
else load_file(checkpoint_path)
1081+
10811082
)
10821083
cache_dir = global_cache_dir("hub")
10831084
pipeline_class = (

invokeai/backend/model_management/model_manager.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -732,9 +732,9 @@ def heuristic_import(
732732

733733
# another round of heuristics to guess the correct config file.
734734
checkpoint = (
735-
safetensors.torch.load_file(model_path)
736-
if model_path.suffix == ".safetensors"
737-
else torch.load(model_path)
735+
torch.load(model_path)
736+
if model_path.suffix == ".ckpt"
737+
else safetensors.torch.load_file(model_path)
738738
)
739739

740740
# additional probing needed if no config file provided

0 commit comments

Comments
 (0)