commit 7c35874ad664e74c8e4252d67521f3986eadb0e6
author Andrew Calvano <calvano@fb.com> Fri Nov 17 17:29:04 2023 +0000
committer PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> Fri Nov 17 17:29:07 2023 +0000
tree f0b72063b9bda39857ab8d741960d77283aaba6b
parent 9f47580ad772c0da7e37a799711563534a51cba5

Fix for PyTorch mobile flatbuffer loader out of bounds reads (#110162)

Summary:
The mobile_ivalue_size field in the mobile_bytecode flatbuffer schema can be larger than the ivalues vector. This introduces potential for memory corruption when parsing the mobile_bytecode Module.

This diff fixes the issue by ensuring that  mobile_ivalue_size is less than the size of the ivalues vector.

Test Plan: contbuild & OSS CI

Differential Revision: D49687548

Pull Request resolved: https://github.com/pytorch/pytorch/pull/110162
Approved by: https://github.com/malfet

torch/csrc/jit/mobile/flatbuffer_loader.cpp

diff --git a/torch/csrc/jit/mobile/flatbuffer_loader.cpp b/torch/csrc/jit/mobile/flatbuffer_loader.cpp
index d8380d2..09b5e9a 100644
--- a/torch/csrc/jit/mobile/flatbuffer_loader.cpp
+++ b/torch/csrc/jit/mobile/flatbuffer_loader.cpp
@@ -302,7 +302,7 @@
   storage_loaded_.resize(module->storage_data_size(), false);
 
   mobile_ivalue_size_ = module_->mobile_ivalue_size();
-  if (mobile_ivalue_size_ == 0) {
+  if (mobile_ivalue_size_ == 0 || mobile_ivalue_size_ > ivalues->size()) {
     mobile_ivalue_size_ = ivalues->size();
   }