The Danger in Model Files

A model file looks like harmless data. But some formats are really programs. Loading one can run code, on your machine, with your permissions. This is one of the most overlooked risks in machine learning, because sharing models from the internet is so common. In this lesson, we demonstrate the problem safely, and see the cure.

• Model files can be programs
• Loading can run code
• Safer formats exist

Code that runs while loading

Here is the problem, in a harmless form. A class defines a special method, that tells Python how to rebuild the object. Instead of rebuilding data, it says, call this function, with these arguments. Here the function is just print, so nothing bad happens. We save the object into bytes, as a model file would be. Then we load it. The message prints, during loading. An attacker would not call print. They would call a function that steals files, or installs malware. And all they need is for you to load their file.

import pickle

class Harmless:
    def __reduce__(self):
        return (print, ("this code ran just by loading the file",))

blob = pickle.dumps(Harmless())
print("loading a model file...")
pickle.loads(blob)
print("...and it already executed code.")

loading a model file...
this code ran just by loading the file
...and it already executed code.

Safer ways to store models

What should you do instead? Prefer formats that store only data, never code. Safe tensors, for neural network weights. O N N X, for portable models, which contains a graph of operations, and no arbitrary code. Plain text formats, like J S O N, for simple models, such as linear coefficients. And, if you must use pickle, for scikit learn models, only load files that you created yourself, or that come from a source you fully trust, and verify them with a checksum, or a signature, before loading.

Verify what you load

Treat every model file like downloaded software. Get it from a trusted registry, not from a random link. Check its checksum against the published one. Where possible, verify a digital signature. Scan it, with tools that detect suspicious model files. Load it first in an isolated sandbox, with no secrets, and no network. And keep a record, in your registry, of where each model came from, and who approved it. Provenance is your defence.

Summary