Varredura de Pickle

Pickle é um formato de serialização amplamente utilizado em ML. Mais notavelmente, é o formato padrão para os pesos de modelos do PyTorch.

Existem ataques perigosos de execução arbitrária de código que podem ser perpetrados ao carregar um arquivo pickle. Sugerimos carregar modelos de usuários e organizações em que você confia, confiar em commits assinados e/ou carregar modelos de formatos TF ou Jax com o mecanismo de conversão automática from_tf=True. Também mitigamos esse problema exibindo/“analisando” a lista de imports em qualquer arquivo pickle, diretamente no Hub. Por fim, estamos experimentando um novo formato de serialização simples para pesos chamado safetensors.

O que é um pickle?

Das documentações oficiais :

O módulo pickle implementa protocolos binários para serializar e desserializar uma estrutura de objetos Python.

Isso significa que pickle é um protocolo de serialização, algo que você usa para compartilhar dados de forma eficiente entre as partes.

Chamamos de pickle o arquivo binário gerado durante o processo de pickling.

Em sua essência, o pickle é basicamente uma pilha de instruções ou opcodes. Como você provavelmente já imaginou, ele não é legível por humanos. Os opcodes são gerados durante o pickling e lidos sequencialmente no unpickling. Com base no opcode, uma ação específica é executada.

Aqui está um pequeno exemplo:

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)

Quando você executa isso, será criado um arquivo pickle e as seguintes instruções serão exibidas no seu terminal:

    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

Não se preocupe demais com as instruções por enquanto, saiba apenas que o módulo pickletools é muito útil para analisar pickles. Ele permite que você leia as instruções no arquivo sem executar nenhum código.

Pickle não é simplesmente um protocolo de serialização; ele oferece mais flexibilidade ao permitir que os usuários executem código Python no momento da desserialização. Não parece bom, certo?

Por que isso é perigoso?

Conforme afirmado acima, desserializar pickle significa que o código pode ser executado. No entanto, isso vem com certas limitações: você só pode referenciar funções e classes do módulo de nível superior; não é possível incorporá-las no próprio arquivo pickle.

De volta à prancheta:

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)

Quando executamos este script, obtemos novamente o payload.pkl. Ao verificarmos o conteúdo do arquivo:

# 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.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

Vemos que não há muito conteúdo ali, apenas alguns opcodes e os dados associados. Você pode estar se perguntando: qual é o problema com o pickle?

Vamos tentar outra coisa:

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)

Aqui, estamos usando a biblioteca fickling pela simplicidade. Ela nos permite adicionar instruções de pickle para executar código contido em uma string por meio da função exec. É assim que você contorna o fato de não poder definir funções ou classes nos seus pickles: você executa código Python salvo como string.

Ao executar isso, ele cria um payload.pkl e imprime o seguinte:

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

Se verificarmos o conteúdo do arquivo pickle, obtemos:

# 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×.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

Basicamente, é isso que acontece quando você faz o 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))
# ...

As instruções que representam uma ameaça são STACK_GLOBAL, GLOBAL e REDUCE.

REDUCE é o que instrui o unpickler a executar a função com os argumentos fornecidos e as instruções *GLOBAL instruem o unpickler a import coisas.

Em resumo, o pickle é perigoso porque:

• ao importar um módulo Python, código arbitrário pode ser executado
• você pode importar funções embutidas como eval ou exec, que podem ser usadas para executar código arbitrário
• ao instanciar um objeto, o construtor pode ser chamado

É por isso que a maioria dos documentos que usam pickle afirma: não desfaça a serialização (unpickle) de dados de fontes não confiáveis.

Estratégias de Mitigação

Não use pickle

Conselho sensato, Luc, mas o pickle é usado profusamente e não vai desaparecer tão cedo: encontrar um novo formato que agrade a todos e iniciar a mudança levará algum tempo.

Então, o que podemos fazer por enquanto?

Carregue arquivos de usuários e organizações em que você confia

No Hub, você tem a capacidade de assinar seus commits com uma chave GPG. Isso não garante que seu arquivo seja seguro, mas garante a origem do arquivo.

Se você conhece e confia no usuário A e o commit que inclui o arquivo no Hub é assinado pela chave GPG do usuário A, é bastante seguro assumir que você pode confiar no arquivo.

Carregar pesos do modelo do TF ou Flax

Os checkpoints do TensorFlow e do Flax não são afetados e podem ser carregados em arquiteturas do PyTorch usando os argumentos from_tf e from_flax do método from_pretrained para contornar esse problema.

E.g.:

from transformers import AutoModel

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

Use seu próprio formato de serialização

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

Este último formato, safetensors, é um formato de serialização simples no qual estamos trabalhando e experimentando atualmente! Ajude ou contribua, se puder 🔥.

Melhorar torch.load/save

Há uma discussão aberta em andamento no PyTorch sobre ter uma forma segura de carregar apenas pesos de arquivos *.pt por padrão – por favor, participe lá!

Scanner de Segurança do Hub

O que temos agora

Criamos um scanner de segurança que examina todos os arquivos enviados ao Hub e executa verificações de segurança. No momento da redação, ele executa dois tipos de varreduras:

• Escaneamentos do ClamAV
• Varreduras de importação de Pickle

Para as verificações do ClamAV, os arquivos são processados pelo antivírus de código aberto ClamAV. Embora isso cubra uma boa quantidade de arquivos perigosos, não cobre exploits de pickle.

Implementamos um scan de Importação Pickle, que extrai a lista de imports referenciados em um arquivo pickle. Sempre que você faz upload de um pytorch_model.bin ou de qualquer outro arquivo pickle, esse scan é executado.

No hub, a lista de imports será exibida ao lado de cada arquivo que contém imports. Se algum import parecer suspeito, ele será destacado.

Obtemos esses dados graças ao pickletools.genops, que nos permite ler o arquivo sem executar código potencialmente perigoso.

Observe que é isso que permite saber se, ao descompactar um arquivo, ele executará REDUCE em uma função potencialmente perigosa que foi importada por *GLOBAL.

Aviso legal: isso não é 100% infalível. É sua responsabilidade, como usuário, verificar se algo é seguro ou não. Não auditamos ativamente pacotes Python quanto à segurança; as listas de importações seguras/insuras que mantemos são geridas da melhor forma possível. Entre em contato conosco se você achar que algo não é seguro e o marcarmos como tal, enviando um e-mail para website at huggingface.co

Soluções potenciais

Pode-se pensar em criar um Unpickler personalizado, como este. Mas, como podemos ver neste exploit sofisticado, isso não funcionará.

Felizmente, sempre há um rastro do import eval, portanto, a leitura direta dos opcodes deve permitir detectar o uso malicioso.

A solução que proponho no momento é criar um arquivo semelhante a um .gitignore, mas para imports.

Este arquivo seria uma lista de permissões de imports que marcaria um arquivo pytorch_model.bin como perigoso caso existam imports não incluídos na lista de permissões.

Pode-se imaginar um formato semelhante a regex no qual seria possível permitir todos os submódulos do numpy, por exemplo, por meio de uma linha simples como: numpy.*.

Leituras adicionais

pickle - serialização de objetos Python - documentação do Python 3.10.6

Picles perigosos - Serialização Python maliciosa

GitHub - trailofbits/fickling: Um descompilador e analisador estático de pickling em Python

Explorando pickles do Python

cpython/pickletools.py em 3.10 · python/cpython

cpython/pickle.py em 3.10 · python/cpython

CrypTen/serial.py em main · facebookresearch/CrypTen

CTFtime.org / Balsn CTF 2019 / pyshv1 / Writeup

Reabilitando o módulo pickle do Python