AI Tensor Engine for ROCm Unauthenticated Remote Code Execution Vulnerability

A remote code execution vulnerability has been identified in AI Tensor Engine for ROCm (AITER) versions through 0.1.14. The issue arises in the MessageQueue.recv() function within shm_broadcast.py, where the application deserializes untrusted data from a ZMQ SUB socket without authentication or validation. This vulnerability allows remote attackers to execute arbitrary code by sending a malicious pickle payload. Exploitation requires access to the writer's XPUB endpoint on the cluster network or the ability to supply a forged Handle with an attacker-controlled remote_subscribe_addr, targeting all remote reader workers simultaneously.

Impact

Exploitation of this vulnerability allows for unauthenticated remote code execution on every affected inference worker in the cluster, with the executed code running in the context of the worker process. This could lead to unauthorized access to model weights stored in GPU memory, exfiltration of those weights, and potential lateral movement within the cluster using worker credentials.

Reproduction

The vulnerability can be reproduced by binding a ZMQ XPUB socket to a TCP address that the AITER writer will connect to. After ensuring the writer subscribes to this socket, the attacker can send a pickle payload that, when deserialized by the AITER worker, executes arbitrary code. This can be automated with a script that mimics the attacker's actions, such as the published proof-of-concept.

Remediation

Users are advised to update to a version of AITER that addresses this vulnerability. If updating is not possible, consider implementing a workaround by modifying the broadcast mechanism to use a safer serialization format, such as msgpack or a custom solution that validates data before deserialization.