Skip to content

Unsafe deserialization via joblib.load("model_head.pkl") leads to remote code execution #630

Description

@Vancir

Description

SetFit loads model_head.pkl from a model repository using joblib.load(...). This is unsafe because joblib.load deserializes arbitrary Python objects and can execute attacker-controlled code during loading. As a result, if an attacker publishes a malicious SetFit model repository containing a crafted model_head.pkl, any user who loads that repository with SetFitModel.from_pretrained(...) can suffer arbitrary code execution on their machine.

This issue is especially dangerous because it is triggered during a normal model-loading workflow. A user may believe they are simply loading a remote model, but the deserialization of model_head.pkl silently executes code before the model is fully initialized. The trust_remote_code=False option does not help here, because the attack does not rely on Hugging Face remote code loading; it relies on unsafe pickle-based deserialization.

Root Cause

The vulnerable code directly loads model_head.pkl with joblib.load(...):

if model_head_file is not None:
model_head = joblib.load(model_head_file)
if isinstance(model_head, torch.nn.Module):
model_head.to(device)

The problem is that joblib.load uses Python pickle semantics under the hood. Loading a malicious pickle file can execute arbitrary code embedded by the attacker. Since model_head.pkl is fetched from the remote model repository and then deserialized without validation, a malicious repository can turn model loading into code execution.

Proof of Concept

An attacker can publish a malicious SetFit model repository such as XManFromXlab/setfit-ModelHead-RCE and place a crafted model_head.pkl inside it. The pickle payload can execute arbitrary Python code when deserialized.

A victim only needs to run the normal SetFit loading code:

from setfit import SetFitModel

model_id = "XManFromXlab/setfit-ModelHead-RCE"
model = SetFitModel.from_pretrained(model_id, trust_remote_code=False)

During from_pretrained(...), SetFit downloads model_head.pkl and calls joblib.load(model_head_file). At that moment, the malicious pickle payload executes on the victim host, leading to arbitrary code execution.

The key point is that this works even when trust_remote_code=False, because the exploit does not depend on custom remote Python modules. It abuses unsafe deserialization of a model artifact.

In this example, it will print the warning mesages.

$ python3 test.py

!!! Execute Malicious Payload !!!

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions