Skip to content

Commit

Permalink
[jit] Validate mobile module fields parsed by flatbuffer loader (pyto…
Browse files Browse the repository at this point in the history
…rch#127437)

Fixing error in `torch.jit.load` Python API function that cause crash in C-backend of PyTorch.
The mobile module is succesfully parsed from flatbuffer format, but its fields are used without any validation.

Fixes pytorch#127434

Pull Request resolved: pytorch#127437
Approved by: https://github.com/davidberard98
  • Loading branch information
apach301 authored and pytorchmergebot committed Jun 3, 2024
1 parent e017b56 commit f343f98
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions torch/csrc/jit/mobile/flatbuffer_loader.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,11 @@ mobile::Module FlatbufferLoader::parseModule(
"Parsing flatbuffer module: Corrupted ivalues/object_types field");
TORCH_CHECK(
reinterpret_cast<const char*>(ivalues) < end, "Corrupted ivalues field");
TORCH_CHECK(
module->storage_data_size() >= 0,
"Parsing flatbuffer module: illegal storage_data_size: ",
module->storage_data_size(),
", expected to be non negative");
all_ivalues_.resize(ivalues->size());
all_types_.resize(module->object_types()->size());
storages_.resize(module->storage_data_size());
Expand Down

0 comments on commit f343f98

Please sign in to comment.