pytorch - 💡(How to fix) Fix [JIT] Use-after-free in Source::_M_dispose() during TorchScript error-path cleanup

Summary

A crafted .pt file triggers a use-after-free crash during error cleanup in SourceImporterImpl::parseSourceIfNeeded(). When TorchScript parsing encounters a malformed source, the exception path unwinds Source objects. Because ConcreteSourceRangeUnpickler holds a SourceRangeDeserializer whose cached_sources map contains shared_ptr<Source> references, the destruction order creates a circular ownership cycle. One entry's Source::_M_dispose() runs after its internal state is already partially freed, reading from address 0x10 (null + offset). Crashes on 2.7.0 and 2.11.0.

Root cause

cached_sources in SourceRangeDeserializer creates a shared ownership cycle:

Source → ConcreteSourceRangeUnpickler → SourceRangeDeserializer
       → cached_sources (map of shared_ptr<Source>) → Source  [cycle]

When the error path triggers stack unwinding, _Hashtable::~_Hashtable() destroys the cached_sources entries. An entry's Source::_M_dispose() fires after the Source's internal state is invalid, causing a null-offset read.

Crash

UndefinedBehaviorSanitizer: SEGV on unknown address 0x000000000010
The signal is caused by a READ memory access.
Hint: address points to the zero page.
  #0 std::_Sp_counted_ptr_inplace<torch::jit::Source, ...>::_M_dispose()
  #1 std::_Hashtable<...>::~_Hashtable()     // cached_sources destruction
  #2 std::_Sp_counted_base::_M_release()
  #3 torch::jit::ConcreteSourceRangeUnpickler::~ConcreteSourceRangeUnpickler()
  #4 std::_Sp_counted_ptr_inplace<torch::jit::Source, ...>::_M_dispose()
  #5 std::_Sp_counted_base::_M_release()
  #6 torch::jit::SourceImporterImpl::parseSourceIfNeeded() (.cold)
  ...
  #21 torch::jit::load()

Reproduction

import torch
torch.jit.load("poc-049-source-dispose.pt")
# → Segmentation fault

PoC file available on request (3,306 bytes).

Suggested fix

Two options:

Option A — clear cached_sources before destruction to break the cycle:

ConcreteSourceRangeUnpickler::~ConcreteSourceRangeUnpickler() {
    if (deserializer) {
        deserializer->cached_sources.clear();
    }
}

Option B — use weak_ptr in cached_sources to avoid the ownership cycle:

std::unordered_map<
    c10::intrusive_ptr<c10::ivalue::Tuple>,
    std::weak_ptr<Source>> cached_sources;

Environment

• PyTorch 2.11.0+cpu, 2.7.0+cpu (both affected)
• Linux x86_64
• Crash is deterministic

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel