# Verification Report: CVE-2026-24747 Fix

## Fix Summary

This fix addresses CVE-2026-24747, a high-severity vulnerability in PyTorch's `weights_only` unpickler (`torch/_weights_only_unpickler.py`) that allows memory corruption and potential arbitrary code execution via malicious `.pth` checkpoint files loaded with `torch.load(..., weights_only=True)`. The fix adds three defenses: (1) type checking on SETITEM/SETITEMS opcodes to restrict them to dict/OrderedDict/Counter types only, preventing `Tensor.__setitem__` from being invoked through the pickle stream; (2) storage provenance tracking via a `_safe_storages` set that records storages loaded from the checkpoint archive through `persistent_load`; and (3) validation in the BUILD opcode handler for Tensor and Parameter types that ensures any `UntypedStorage` passed to `set_()` or `__setstate__()` originated from the archive rather than being constructed via REDUCE with attacker-controlled bytes.

## Changes Made

### File Modified: `torch/_weights_only_unpickler.py`

1. **Added `_safe_storages` set to `Unpickler.__init__`**: Tracks the identity (`id()`) of UntypedStorage objects that were legitimately loaded from the checkpoint archive via `persistent_load`. This provides a provenance check for storages used in BUILD operations.

2. **Added `_check_set_item_target` method**: New validation method that checks whether the top-of-stack object is a dict, OrderedDict, or Counter before allowing SETITEM/SETITEMS operations. Raises `UnpicklingError` if the target is any other type (e.g., Tensor).

3. **SETITEM handler (original line 465)**: Added call to `self._check_set_item_target("SETITEM")` before executing the assignment. This blocks `Tensor.__setitem__` and any other non-dict `__setitem__` from being triggered.

4. **SETITEMS handler (original line 468)**: Added call to `self._check_set_item_target("SETITEMS")` before the assignment loop. Same protection as SETITEM but for batch operations.

5. **BUILD handler for Tensor (original line 421)**: Added validation that when `state` is a tuple containing an `UntypedStorage` as its first element, that storage must exist in `_safe_storages` (i.e., came from `persistent_load`). This blocks the BUILD+REDUCE bypass where an attacker constructs a malicious UntypedStorage via `REDUCE(_codecs.encode, ...) + REDUCE(UntypedStorage, ...)` and uses BUILD to rebind a tensor to it.

6. **BUILD handler for Parameter (original line 422)**: Same storage provenance validation applied to `Parameter.__setstate__`, which internally also calls `set_()` when the state has 4+ elements.

7. **BINPERSID handler (original line 535)**: Modified to track storages loaded via `persistent_load` by adding their `id()` to `self._safe_storages` before appending to the stack. This is the source of truth for legitimate archive-sourced storages.

## Verification Steps

### Step 1: Patch applies cleanly
```
$ cd /home/vscode/.local/lib/python3.12/site-packages && patch -p1 < proposed_fix.diff
patching file torch/_weights_only_unpickler.py
```

### Step 2: Run verification script
```
$ bash coding/verify_fix.sh
```

### Step 3: Verification output
```
=== Test 1: SETITEM/SETITEMS exploit should be BLOCKED ===
[+] PASS: SETITEMS exploit blocked correctly

=== Test 2: BUILD+REDUCE UntypedStorage bypass should be BLOCKED ===
[+] PASS: BUILD+REDUCE bypass blocked correctly

=== Test 3: Normal model loading should still work ===
[+] PASS: Normal model loaded successfully
    Keys: ['weight', 'bias']
    Weight values match: True
    Bias values match: True

=== Test 4: Complex model with OrderedDict should still work ===
[+] PASS: Complex model loaded successfully
    Keys: ['0.weight', '0.bias', '2.weight', '2.bias']
    Type: OrderedDict (correct)

=== Test 5: SETITEM on dict should still work ===
[+] PASS: SETITEM on dict works correctly

=== Results ===
Passed: 5
Failed: 0

FIX_VERIFIED
```

## Test Results

| # | Test | Description | Result |
|---|------|-------------|--------|
| 1 | SETITEMS exploit | Malicious pickle uses SETITEMS to call `tensor[idx] = value` | **BLOCKED** - raises UnpicklingError |
| 2 | BUILD+REDUCE bypass | Malicious pickle uses REDUCE(UntypedStorage) + BUILD to rebind tensor storage | **BLOCKED** - raises UnpicklingError |
| 3 | Normal Linear model | `torch.nn.Linear(10, 5)` state dict save/load round-trip | **PASS** - values match |
| 4 | Complex Sequential model | Multi-layer Sequential model with OrderedDict state dict | **PASS** - all keys load correctly |
| 5 | Manual SETITEM on dict | Handcrafted pickle that uses SETITEM correctly on a dict | **PASS** - normal usage unaffected |

### Edge cases tested:
- **SETITEMS on Tensor**: Correctly blocked with descriptive error message mentioning allowed types
- **BUILD with REDUCE-constructed UntypedStorage**: Correctly blocked because storage ID not in `_safe_storages`
- **Normal dict SETITEM**: Still works (regression check)
- **Complex model with nested OrderedDict**: Still works (regression check)
- **Idempotency**: Script runs identically on repeated executions

## Remaining Concerns

1. **`_codecs.encode` in allowed globals**: While this fix prevents the BUILD+REDUCE bypass, `_codecs.encode` remains in the allowed globals and can still construct arbitrary byte sequences. Consider removing it or restricting allowed encodings in a future hardening pass.

2. **BUILD on user-allowed globals**: The BUILD handler for user-allowed globals still permits `__setstate__`, `__dict__.update`, and `setattr` operations. If users allowlist types with dangerous `__setstate__` methods, they could still be exploited.

3. **`id()`-based tracking**: The `_safe_storages` tracking uses Python `id()` values. While Python objects loaded from the archive persist for the duration of unpickling (so id reuse is not a concern within a single `load()` call), this is worth noting as a design decision.

4. **REDUCE for UntypedStorage still permitted**: `UntypedStorage` can still be constructed via REDUCE, which allocates uninitialized heap memory (information leak). A more comprehensive fix would restrict UntypedStorage construction to only `persistent_load`, but this may break legitimate use cases where users have allowlisted custom types that create storages.

5. **Recommended additional hardening**: Track all REDUCE-created objects and validate that BUILD never operates on REDUCE-created storages, rather than just checking the first element of the state tuple.