huggingface / huggingface/setfit

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

Open
#630 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
2.8k
Forks
267
Avg merge
36m
Merged PRs (30d)
5

Description

## 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(...)`:

https://github.com/huggingface/setfit/blob/0f8e322ce7c62eeb7b2d113f88f1b1fc37777692/src/setfit/modeling.py#L828-L831

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:

```python
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.

```bash
$ python3 test.py

!!! Execute Malicious Payload !!!
```

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.