huggingface / huggingface/setfit
Unsafe deserialization via joblib.load("model_head.pkl") leads to remote code execution
- 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.