# NVIDIA NeMo Model File Vulnerability PoC

This repository contains a harmless proof-of-concept for a NVIDIA NeMo `.nemo` model restoration issue.

## Summary

A crafted `.nemo` archive can include an attacker-controlled `model_weights.ckpt` file. During model restoration, NeMo loads this checkpoint using PyTorch deserialization.

The PoC demonstrates that attacker-controlled pickle content can execute during the public `restore_from()` loading workflow.

## Safety

The payload is intentionally harmless. It only creates this marker file:

```text
/tmp/nemo_mfv_public_restore_dict_marker.txt
```

No destructive action is performed.

## Reproduction

Create a clean Python environment and install dependencies:

```bash
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -U pip
pip install -r requirements.txt
```

Run:

```bash
python reproduce.py
```

Expected PoC result:

```text
[+] Marker after restore: True
[+] Marker content:
NeMo public restore dict payload marker
```

The restore process may raise an exception about an unexpected key in the state_dict. This happens after the payload has already executed, which demonstrates that unsafe deserialization occurs before state_dict validation.

## Affected Component

- Framework: NVIDIA NeMo
- Format: `.nemo`
- File inside archive: `model_weights.ckpt`
- Loading API: `restore_from()`
- Sink: `torch.load(..., weights_only=False)`

## Impact

If a user or automated ML pipeline restores a malicious `.nemo` model from an untrusted source, attacker-controlled code may execute during model loading.