Verify flatbuffer module fields are initialized (#109794)

apach301 · pytorchmergebot · commit 175ccfc4c844 · 2023-09-21T23:19:17.000Z
Fixes #109793 Add validation on flatbuffer module field to prevent segfault Pull Request resolved: #109794 Approved by: https://github.com/malfet
diff --git a/torch/csrc/jit/mobile/flatbuffer_loader.cpp b/torch/csrc/jit/mobile/flatbuffer_loader.cpp
@@ -291,9 +291,11 @@ mobile::Module FlatbufferLoader::parseModule(
   module_parsed_ = false;
 
   const auto* ivalues = module->ivalues();
-  TORCH_CHECK(ivalues != nullptr, "Corrupted ivalues field")
   TORCH_CHECK(
-      reinterpret_cast<const char*>(ivalues) < end, "Corrupted ivalues field")
+      ivalues && module->object_types(),
+      "Parsing flatbuffer module: Corrupted ivalues/object_types field");
+  TORCH_CHECK(
+      reinterpret_cast<const char*>(ivalues) < end, "Corrupted ivalues field");
   all_ivalues_.resize(ivalues->size());
   all_types_.resize(module->object_types()->size());
   storages_.resize(module->storage_data_size());
