Quét Pickle

Pickle là định dạng serialization được sử dụng rộng rãi trong ML. Đáng chú ý nhất, đây là định dạng mặc định cho trọng số mô hình PyTorch.

Có những cuộc tấn công thực thi mã tùy ý nguy hiểm có thể xảy ra khi bạn tải một tệp pickle. Chúng tôi khuyên bạn nên tải các mô hình từ những người dùng và tổ chức mà bạn tin tưởng, dựa vào các commit đã ký, và/hoặc tải các mô hình từ định dạng TF hoặc Jax bằng cơ chế tự động chuyển đổi from_tf=True. Chúng tôi cũng giảm thiểu vấn đề này bằng cách hiển thị/“xác minh” danh sách các import trong bất kỳ tệp pickle nào, ngay trên Hub. Cuối cùng, chúng tôi đang thử nghiệm một định dạng chuỗi hóa mới, đơn giản cho trọng số gọi là safetensors.

Pickle là gì?

Từ tài liệu chính thức :

Module pickle triển khai các giao thức nhị phân để phân tích chuỗi và giải phân tích chuỗi cấu trúc đối tượng Python.

Điều này có nghĩa là pickle là một giao thức phân chuỗi, thứ bạn dùng để chia sẻ dữ liệu một cách hiệu quả giữa các bên.

Chúng tôi gọi pickle là tệp nhị phân được tạo ra trong quá trình pickling.

Về cốt lõi, pickle về cơ bản là một chồng các chỉ lệnh hoặc opcode. Như bạn có thể đã đoán, nó không thể đọc được bằng mắt thường. Các opcode được tạo ra khi pickle và được đọc tuần tự khi unpickle. Dựa trên opcode, một hành động cụ thể sẽ được thực thi.

Dưới đây là một ví dụ nhỏ:

import pickle
import pickletools

var = "data I want to share with a friend"

# store the pickle data in a file named 'payload.pkl'
with open('payload.pkl', 'wb') as f:
    pickle.dump(var, f)

# disassemble the pickle
# and print the instructions to the command line
with open('payload.pkl', 'rb') as f:
    pickletools.dis(f)

Khi bạn chạy lệnh này, một tệp pickle sẽ được tạo và các hướng dẫn sau sẽ được in ra trong terminal của bạn:

    0: \x80 PROTO      4
    2: \x95 FRAME      48
   11: \x8c SHORT_BINUNICODE 'data I want to share with a friend'
   57: \x94 MEMOIZE    (as 0)
   58: .    STOP
highest protocol among opcodes = 4

Đừng quá lo lắng về các chỉ dẫn cho lúc này, chỉ cần biết rằng module pickletools rất hữu ích để phân tích pickle. Nó cho phép bạn đọc các chỉ dẫn trong tệp mà không thực thi bất kỳ mã nào.

Pickle không đơn thuần là một giao thức phân tích chuỗi (serialization), mà còn mang lại sự linh hoạt hơn bằng cách cho phép người dùng thực thi mã Python tại thời điểm khử phân tích chuỗi (de-serialization). Nghe có vẻ không ổn, đúng không?

Tại sao điều này lại nguy hiểm?

Như đã nêu ở trên, việc giải chuỗi hóa pickle có nghĩa là mã có thể được thực thi. Nhưng điều này đi kèm với một số hạn chế: bạn chỉ có thể tham chiếu các hàm và lớp từ module cấp trên cùng; bạn không thể nhúng chúng vào chính tệp pickle.

Quay lại vạch xuất phát:

import pickle
import pickletools

class Data:
    def __init__(self, important_stuff: str):
        self.important_stuff = important_stuff

d = Data("42")

with open('payload.pkl', 'wb') as f:
    pickle.dump(d, f)

Khi chúng ta chạy script này, chúng ta lại nhận được payload.pkl. Khi kiểm tra nội dung của tệp:

# cat payload.pkl
__main__Data)}important_stuff42sb.%

# hexyl payload.pkl
┌────────┬─────────────────────────┬─────────────────────────┬────────┬────────┐
│00000000│ 80 04 95 33 00 00 00 00 ┊ 00 00 00 8c 08 5f 5f 6d │ו×30000┊000ו__m│
│00000010│ 61 69 6e 5f 5f 94 8c 04 ┊ 44 61 74 61 94 93 94 29 │ain__×ו┊Data×××)│
│00000020│ 81 94 7d 94 8c 0f 69 6d ┊ 70 6f 72 74 61 6e 74 5f │××}×וim┊portant_│
│00000030│ 73 74 75 66 66 94 8c 02 ┊ 34 32 94 73 62 2e       │stuff×ו┊42×sb.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

Chúng ta có thể thấy rằng bên trong không có gì nhiều, chỉ một vài opcode và dữ liệu liên quan. Bạn có thể đang nghĩ, vậy vấn đề với pickle là gì?

Hãy thử một điều khác:

from fickling.pickle import Pickled
import pickle

# Create a malicious pickle
data = "my friend needs to know this"

pickle_bin = pickle.dumps(data)

p = Pickled.load(pickle_bin)

p.insert_python_exec('print("you\'ve been pwned !")')

with open('payload.pkl', 'wb') as f:
    p.dump(f)

# innocently unpickle and get your friend's data
with open('payload.pkl', 'rb') as f:
    data = pickle.load(f)
    print(data)

Ở đây chúng ta sử dụng thư viện fickling cho sự đơn giản. Nó cho phép chúng ta thêm các chỉ thị pickle để thực thi mã chứa trong một chuỗi thông qua hàm exec. Đây là cách bạn vượt qua thực tế là không thể định nghĩa các hàm hoặc lớp trong các pickle của mình: bạn chạy exec trên mã python được lưu dưới dạng chuỗi.

Khi bạn chạy lệnh này, nó sẽ tạo ra payload.pkl và in ra nội dung sau:

you've been pwned !
my friend needs to know this

Nếu chúng ta kiểm tra nội dung của tệp pickle, ta sẽ nhận được:

# cat payload.pkl
c__builtin__
exec
(Vprint("you've been pwned !")
tR my friend needs to know this.%

# hexyl payload.pkl
┌────────┬─────────────────────────┬─────────────────────────┬────────┬────────┐
│00000000│ 63 5f 5f 62 75 69 6c 74 ┊ 69 6e 5f 5f 0a 65 78 65 │c__built┊in___exe│
│00000010│ 63 0a 28 56 70 72 69 6e ┊ 74 28 22 79 6f 75 27 76 │c_(Vprin┊t("you'v│
│00000020│ 65 20 62 65 65 6e 20 70 ┊ 77 6e 65 64 20 21 22 29 │e been p┊wned !")│
│00000030│ 0a 74 52 80 04 95 20 00 ┊ 00 00 00 00 00 00 8c 1c │_tR×•× 0┊000000ו│
│00000040│ 6d 79 20 66 72 69 65 6e ┊ 64 20 6e 65 65 64 73 20 │my frien┊d needs │
│00000050│ 74 6f 20 6b 6e 6f 77 20 ┊ 74 68 69 73 94 2e       │to know ┊this×.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

Về cơ bản, đây là những gì xảy ra khi bạn unpickle:

# ...
opcodes_stack = [exec_func, "malicious argument", "REDUCE"]
opcode = stack.pop()
if opcode == "REDUCE":
    arg = opcodes_stack.pop()
    callable = opcodes_stack.pop()
    opcodes_stack.append(callable(arg))
# ...

Các lệnh gây ra mối đe dọa là STACK_GLOBAL, GLOBAL và REDUCE.

REDUCE là lệnh yêu cầu unpickler thực thi hàm với các đối số được cung cấp, còn các chỉ thị *GLOBAL yêu cầu unpickler import các thành phần liên quan.

Tóm lại, pickle là mối nguy hiểm vì:

• khi import một module Python, mã tùy ý có thể được thực thi
• bạn có thể import các hàm tích hợp như eval hoặc exec, có thể được dùng để thực thi mã tùy ý
• khi khởi tạo một đối tượng, hàm khởi tạo có thể được gọi

Đây là lý do tại sao hầu hết các tài liệu sử dụng pickle đều cảnh báo không nên unpickle dữ liệu từ các nguồn không đáng tin cậy.

Các chiến lược giảm thiểu

Không sử dụng pickle

Lời khuyên hay đấy Luc, nhưng pickle được sử dụng rất phổ biến và sẽ không biến mất sớm: việc tìm ra một định dạng mới mà mọi người đều hài lòng và khởi xướng sự thay đổi sẽ mất một khoảng thời gian.

Vậy hiện tại chúng ta có thể làm gì?

Tải tệp từ các người dùng và tổ chức mà bạn tin tưởng

Trên Hub, bạn có thể chữ ký các commit của mình bằng một khóa GPG. Điều này không đảm bảo rằng tệp của bạn an toàn, nhưng nó đảm bảo nguồn gốc của tệp.

Nếu bạn biết và tin tưởng người dùng A, và commit chứa tệp trên Hub được ký bằng khóa GPG của người dùng A, thì có thể an toàn giả định rằng bạn có thể tin tưởng tệp đó.

Tải trọng số mô hình từ TF hoặc Flax

Checkpoint của TensorFlow và Flax không bị ảnh hưởng và có thể được nạp vào các kiến trúc PyTorch bằng cách sử dụng các tham số from_tf và from_flax của phương thức from_pretrained để khắc phục vấn đề này.

E.g.:

from transformers import AutoModel

model = AutoModel.from_pretrained("google-bert/bert-base-cased", from_flax=True)

Sử dụng định dạng chuỗi hóa của riêng bạn

• MsgPack
• Protobuf
• Cap’n’proto
• Avro
• safetensors

Định dạng cuối cùng này, safetensors, là một định dạng chuỗi hóa đơn giản mà chúng tôi đang phát triển và thử nghiệm! Hãy giúp đỡ hoặc đóng góp nếu bạn có thể 🔥.

Cải thiện torch.load/save

Hiện có một cuộc thảo luận đang diễn ra trên PyTorch về việc có Cách an toàn để chỉ tải trọng số từ tệp *.pt theo mặc định – vui lòng tham gia thảo luận tại đó!

Trình quét bảo mật của Hub

Những gì chúng ta có hiện tại

Chúng tôi đã tạo ra một trình quét bảo mật, quét mọi tệp được đẩy lên Hub và thực hiện các kiểm tra bảo mật. Tại thời điểm viết tài liệu này, nó chạy hai loại quét:

• Quét ClamAV
• Quét các lệnh nhập Pickle

Đối với các lần quét ClamAV, các tệp sẽ được chạy qua phần mềm diệt virus mã nguồn mở ClamAV. Mặc dù điều này bao phủ một lượng lớn các tệp nguy hiểm, nhưng nó không bao gồm các lỗ hổng khai thác pickle.

Chúng tôi đã triển khai một quy trình quét Pickle Import, có chức năng trích xuất danh sách các import được tham chiếu trong một tệp pickle. Mỗi khi bạn tải lên một pytorch_model.bin hoặc bất kỳ tệp pickle nào khác, quy trình quét này sẽ được thực hiện.

Trên hub, danh sách các lệnh import sẽ được hiển thị bên cạnh mỗi tệp chứa import. Nếu có bất kỳ lệnh import nào trông đáng ngờ, nó sẽ được tô sáng.

Chúng ta có được dữ liệu này nhờ pickletools.genops, cho phép đọc tệp mà không thực thi các mã có thể gây nguy hiểm.

Lưu ý rằng đây là điều cho phép biết được liệu, khi unpickle một tệp, nó có sẽ REDUCE trên một hàm có thể gây nguy hiểm đã được import bởi *GLOBAL hay không.

Thông báo miễn trừ trách nhiệm: đây không phải là giải pháp an toàn 100%. Trách nhiệm của bạn với tư cách là người dùng là kiểm tra xem một thứ gì đó có an toàn hay không. Chúng tôi không chủ động kiểm toán các gói python về mặt an toàn, các danh sách import an toàn/không an toàn mà chúng tôi có được duy trì theo phương thức nỗ lực tốt nhất. Vui lòng liên hệ với chúng tôi nếu bạn nghĩ rằng một thứ gì đó không an toàn và chúng tôi sẽ đánh dấu nó như vậy, bằng cách gửi email cho chúng tôi tại website ở huggingface.co

Các giải pháp tiềm năng

Có thể nghĩ đến việc tạo một Unpickler tùy chỉnh theo kiểu như cái này. Nhưng như chúng ta có thể thấy trong chiến thuật khai thác tinh vi này, cách đó sẽ không hiệu quả.

May mắn thay, luôn có dấu vết của lệnh import eval, vì vậy việc đọc trực tiếp các opcode nên cho phép phát hiện được các hành vi sử dụng độc hại.

Giải pháp hiện tại mà tôi đề xuất là tạo một file tương tự .gitignore nhưng dành cho các import.

Tệp này sẽ là một danh sách trắng các lệnh import, theo đó tệp pytorch_model.bin sẽ được đánh dấu là nguy hiểm nếu có các lệnh import không nằm trong danh sách trắng.

Có thể hình dung một định dạng kiểu regex, trong đó bạn có thể cho phép tất cả các submodule của numpy, chẳng hạn bằng một dòng đơn giản như: numpy.*.

Đọc thêm

pickle - Nạp xuất đối tượng Python - Tài liệu Python 3.10.6

Pickle nguy hiểm - Mã hóa Python độc hại

GitHub - trailofbits/fickling: Trình biên dịch ngược và bộ phân tích tĩnh cho Python pickling

Lợi dụng Python pickles

cpython/pickletools.py tại 3.10 · python/cpython

cpython/pickle.py tại 3.10 · python/cpython

CrypTen/serial.py tại main · facebookresearch/CrypTen

CTFtime.org / Balsn CTF 2019 / pyshv1 / Bài phân tích

Cải thiện module pickle của Python