Segmentation fault during torch::jit::load

[apach301]

🐛 Describe the bug

Hi,
I found a bugt during testing with libFuzzer.

It is located at /pytorch_fuzz/torch/csrc/jit/frontend/source_range.cpp:93 in torch::jit::SourceRange::print_with_context().
There are a segmentation fault and heap buffer overflow discovered at this function when processing malformed model with torch::jit::load function.

The bug is triggered when begin_line obtained from another function is used to access an element of str array without checking its length.

    90       // determine CONTEXT line range
    91       size_t begin_line = start(); // beginning of lines to highlight
    92       size_t end_line = range_end;
--->93       while (begin_line > 0 && str[begin_line - 1] != '\\n')
    94         --begin_line;

The bug could be reproduced with aot_model_compiler binary:

aot_model_compiler --model=crash-file --model_name=name --model_version=1 --input_dims='1,3,224,224;2,2' --input_types='float;float'

Segmentation fault (crash.zip) is reproduced only with sanitizers, on clean binary it is std::out_of_range exception:

  "AsanReport": [
    "==155192==ERROR: AddressSanitizer: SEGV on unknown address 0x6160009b0f1b (pc 0x000000bf9b1e bp 0x7ffe2b2a5e50 sp 0x7ffe2b2a5860 T0)\n",
    "==155192==The signal is caused by a READ memory access.\n",
    "[Detaching after fork from child process 155586]\n",
    "    #0 0xbf9b1e in torch::jit::SourceRange::print_with_context(std::ostream&, unsigned long, bool, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) const /pytorch_fuzz/torch/csrc/jit/frontend/source_range.cpp:93:28\n",
    "    #1 0xbfbd85 in torch::jit::format_stack_trace(std::ostream&, std::vector<torch::jit::StackEntry, std::allocator<torch::jit::StackEntry> > const&) /pytorch_fuzz/torch/csrc/jit/frontend/source_range.cpp:60:13\n",
    "    #2 0x9b3cc92 in torch::jit::InterpreterStateImpl::formatStackTrace(std::ostream&) /pytorch_fuzz/torch/csrc/jit/runtime/interpreter.cpp:818:5\n",
    "    #3 0x9b34a4a in torch::jit::InterpreterStateImpl::handleError(torch::jit::ExceptionMessage const&, bool, c10::NotImplementedError*, c10::optional<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >) /pytorch_fuzz/torch/csrc/jit/runtime/interpreter.cpp:830:5\n",
    "    #4 0x9b3048e in torch::jit::InterpreterStateImpl::runImpl(std::vector<c10::IValue, std::allocator<c10::IValue> >&) /pytorch_fuzz/torch/csrc/jit/runtime/interpreter.cpp:801:7\n",
    "    #5 0x9ae1c7b in torch::jit::InterpreterStateImpl::run(std::vector<c10::IValue, std::allocator<c10::IValue> >&) /pytorch_fuzz/torch/csrc/jit/runtime/interpreter.cpp:1008:9\n",
    "    #6 0x9aaee38 in torch::jit::GraphExecutorImplBase::run(std::vector<c10::IValue, std::allocator<c10::IValue> >&) /pytorch_fuzz/torch/csrc/jit/runtime/graph_executor.cpp:544:31\n",
    "    #7 0xab135c in torch::jit::Function::operator()(std::vector<c10::IValue, std::allocator<c10::IValue> >, std::unordered_map<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, c10::IValue, std::hash<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::equal_to<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::allocator<std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const, c10::IValue> > > const&) /pytorch_fuzz/aten/src/ATen/core/function.h:61:5\n",
    "    #8 0x9e0533f in torch::jit::(anonymous namespace)::ScriptModuleDeserializer::readArchive(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)::$_1::operator()(c10::StrongTypePtr const&, c10::IValue) const /pytorch_fuzz/torch/csrc/jit/serialization/import.cpp:159:7\n",
    "    #9 0x9e0533f in c10::intrusive_ptr<c10::ivalue::Object, c10::detail::intrusive_target_default_null_type<c10::ivalue::Object> > std::__invoke_impl<c10::intrusive_ptr<c10::ivalue::Object, c10::detail::intrusive_target_default_null_type<c10::ivalue::Object> >, torch::jit::(anonymous namespace)::ScriptModuleDeserializer::readArchive(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)::$_1&, c10::StrongTypePtr, c10::IValue>(std::__invoke_other, torch::jit::(anonymous namespace)::ScriptModuleDeserializer::readArchive(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)::$_1&, c10::StrongTypePtr&&, c10::IValue&&) /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/invoke.h:60:14\n",
    "    #10 0x9e044b0 in std::enable_if<__and_<std::__not_<std::is_void<c10::intrusive_ptr<c10::ivalue::Object, c10::detail::intrusive_target_default_null_type<c10::ivalue::Object> > > >, std::is_convertible<std::__invoke_result<torch::jit::(anonymous namespace)::ScriptModuleDeserializer::readArchive(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)::$_1&, c10::StrongTypePtr, c10::IValue>::type, c10::intrusive_ptr<c10::ivalue::Object, c10::detail::intrusive_target_default_null_type<c10::ivalue::Object> > > >::value, c10::intrusive_ptr<c10::ivalue::Object, c10::detail::intrusive_target_default_null_type<c10::ivalue::Object> > >::type std::__invoke_r<c10::intrusive_ptr<c10::ivalue::Object, c10::detail::intrusive_target_default_null_type<c10::ivalue::Object> >, torch::jit::(anonymous namespace)::ScriptModuleDeserializer::readArchive(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)::$_1&, c10::StrongTypePtr, c10::IValue>(torch::jit::(anonymous namespace)::ScriptModuleDeserializer::readArchive(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)::$_1&, c10::StrongTypePtr&&, c10::IValue&&) /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/invoke.h:141:14\n",
    "    #11 0x9e0431b in std::_Function_handler<c10::intrusive_ptr<c10::ivalue::Object, c10::detail::intrusive_target_default_null_type<c10::ivalue::Object> > (c10::StrongTypePtr, c10::IValue), torch::jit::(anonymous namespace)::ScriptModuleDeserializer::readArchive(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)::$_1>::_M_invoke(std::_Any_data const&, c10::StrongTypePtr&&, c10::IValue&&) /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/std_function.h:291:9\n",
    "    #12 0x9eed44d in std::function<c10::intrusive_ptr<c10::ivalue::Object, c10::detail::intrusive_target_default_null_type<c10::ivalue::Object> > (c10::StrongTypePtr, c10::IValue)>::operator()(c10::StrongTypePtr, c10::IValue) const /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/std_function.h:622:14\n",
    "    #13 0x9eed0ee in torch::jit::Unpickler::readGlobal(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)::$_9::operator()() const /pytorch_fuzz/torch/csrc/jit/serialization/unpickler.cpp:704:20\n",
    "    #14 0x9eed0ee in void std::__invoke_impl<void, torch::jit::Unpickler::readGlobal(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)::$_9&>(std::__invoke_other, torch::jit::Unpickler::readGlobal(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)::$_9&) /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/invoke.h:60:14\n",
    "    #15 0x9ecb6d4 in torch::jit::Unpickler::readInstruction() /pytorch_fuzz/torch/csrc/jit/serialization/unpickler.cpp:444:7\n",
    "    #16 0x9ec9ac7 in torch::jit::Unpickler::run() /pytorch_fuzz/torch/csrc/jit/serialization/unpickler.cpp:226:27\n",
    "    #17 0x9ec9762 in torch::jit::Unpickler::parse_ivalue() /pytorch_fuzz/torch/csrc/jit/serialization/unpickler.cpp:183:3\n",
    "    #18 0x9e0c84e in torch::jit::readArchiveAndTensors(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, c10::optional<std::function<c10::StrongTypePtr (c10::QualifiedName const&)> >, c10::optional<std::function<c10::intrusive_ptr<c10::ivalue::Object, c10::detail::intrusive_target_default_null_type<c10::ivalue::Object> > (c10::StrongTypePtr, c10::IValue)> >, c10::optional<c10::Device>, caffe2::serialize::PyTorchStreamReader&, c10::Type::SingletonOrSharedTypePtr<c10::Type> (*)(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&), std::shared_ptr<torch::jit::DeserializationStorageContext>) /pytorch_fuzz/torch/csrc/jit/serialization/import_read.cpp:54:20\n",
    "    #19 0x9e02779 in torch::jit::(anonymous namespace)::ScriptModuleDeserializer::readArchive(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) /pytorch_fuzz/torch/csrc/jit/serialization/import.cpp:171:10\n",
    "    #20 0x9dfe60b in torch::jit::(anonymous namespace)::ScriptModuleDeserializer::deserialize(c10::optional<c10::Device>, std::unordered_map<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::hash<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::equal_to<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::allocator<std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > > >&) /pytorch_fuzz/torch/csrc/jit/serialization/import.cpp:274:19\n",
    "    #21 0x9e013c4 in torch::jit::load(std::shared_ptr<caffe2::serialize::ReadAdapterInterface>, c10::optional<c10::Device>, std::unordered_map<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::hash<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::equal_to<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::allocator<std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > > >&) /pytorch_fuzz/torch/csrc/jit/serialization/import.cpp:403:23\n",
    "    #22 0x9e01977 in torch::jit::load(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, c10::optional<c10::Device>, std::unordered_map<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::hash<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::equal_to<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::allocator<std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > > >&) /pytorch_fuzz/torch/csrc/jit/serialization/import.cpp:375:17\n",
    "    #23 0x9e0176c in torch::jit::load(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, c10::optional<c10::Device>) /pytorch_fuzz/torch/csrc/jit/serialization/import.cpp:367:10\n",
    "    #24 0x5cdfd2 in LLVMFuzzerTestOneInput /load_fuzz.cc:45:18\n",
    "    #25 0x4fac81 in fuzzer::Fuzzer::ExecuteCallback(unsigned char const*, unsigned long) /llvm-project/compiler-rt/lib/fuzzer/FuzzerLoop.cpp:611:15\n",
    "    #26 0x4e4b9c in fuzzer::RunOneTest(fuzzer::Fuzzer*, char const*, unsigned long) /llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:324:6\n",
    "    #27 0x4ea8eb in fuzzer::FuzzerDriver(int*, char***, int (*)(unsigned char const*, unsigned long)) /llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:860:9\n",
    "    #28 0x513e82 in main /llvm-project/compiler-rt/lib/fuzzer/FuzzerMain.cpp:20:10\n",
    "    #29 0x7f672e91a0b2 in __libc_start_main /build/glibc-eX1tMB/glibc-2.31/csu/../csu/libc-start.c:308:16\n",
    "    #30 0x4df4bd in _start (/fuzz/load_fuzz+0x4df4bd)\n",

Also there are Asan reports on heap buffer overflow (crash.zip) and heap use after free (crash.zip) for the same place.

Another segmentation fault (crash.zip) is occured in different place but has stacktrace very similar to the segfault described above.
In /pytorch_fuzz/c10/util/intrusive_ptr.h:269:

    267        if (target_ != NullType::singleton()) {
    268          size_t new_refcount =
--->269              detail::atomic_refcount_increment(target_->refcount_);
    270          TORCH_INTERNAL_ASSERT_DEBUG_ONLY(

Versions

PyTorch version: 1.12.0a0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: 14.0.0 (https://github.com/llvm/llvm-project.git de5b16d8ca2d14ff0d9b6be9cf40566bc7eb5a01)
CMake version: version 3.22.4
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 26 2021, 20:14:08) [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-4.15.0-175-generic-x86_64-with-glibc2.29
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: N/A

Versions of relevant libraries:
[pip3] numpy==1.22.2
[conda] Could not collect

[Gamrix]

torch::jit::load should only be used with trusted inputs. This is a Won't Fix

You are welcome to supply a PR, and I will review it.