Removed file-extension-based arbitrary code execution attack vector #2946
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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
andconvert_ckpt_to_diffusers.py
there are if-statements that decide whichload
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: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
.ckpt
extension to some variation of.Safetensors
, ensuring there is a capital letter anywhere in the extension (eg.malicious_pickle.SAFETENSORS
)torch.load
and the payload will execute.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 fortorch.load
, as this significantly reduces the amount of Python methods that an attacker has to work with when crafting malicious payloadsinside a pickle file.
Automatic1111 already uses this with some success. (https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/safe.py)