Bug #8 — Potential Eval Injection in create_state_spec_from_checkpoint (axlearn)

[YuvalFradkin1]

Title:
Bug #8 — Potential Eval Injection in create_state_spec_from_checkpoint (axlearn)

## Bug #8 — Potential Eval Injection (CWE-95)

I identified a potential code execution issue in:

axlearn/cloud/gcp/examples/colocated_python_benchmark.py

Specifically in:
create_state_spec_from_checkpoint()

### Summary
The function processes checkpoint metadata and evaluates the `shape` field when it is provided as a string:

```python
# pylint: disable=eval-used
shape = eval(value["shape"]) if isinstance(value["shape"], str) else value["shape"]

If the checkpoint index file is not trusted, this may allow execution of arbitrary Python code during checkpoint loading.

Context

• The shape value originates from a checkpoint index file
• The checkpoint path may come from user input or external sources
• No validation or sanitization appears to be applied before evaluation

Impact

If a malicious checkpoint is loaded, this could lead to arbitrary code execution in the context of the running process.

This may be relevant in workflows where checkpoints are shared via external sources (e.g. research artifacts, storage buckets, etc.).

Responsible Disclosure

I have a working proof-of-concept and full technical details, but I prefer to share them privately.

Please let me know the appropriate channel for secure disclosure.