File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff line change 88
99#include < executorch/runtime/executor/tensor_parser.h>
1010
11+ #include < climits>
12+
1113#include < executorch/runtime/core/exec_aten/exec_aten.h>
1214#include < executorch/runtime/core/exec_aten/util/dim_order_util.h>
1315#include < executorch/runtime/core/exec_aten/util/scalar_type_util.h>
@@ -118,17 +120,24 @@ Result<Tensor> parseTensor(
118120 dim_order =
119121 const_cast <executorch::aten::DimOrderType*>(serialized_dim_order);
120122 }
121- // Validate sizes before using them in case the PTE data is bad. We can't
122- // detect bad positive values, but we can reject negative values, which would
123- // otherwise panic in the TensorImpl ctor. dim_order_to_stride() will validate
124- // dim_order.
123+ // Validate sizes before using them in case the PTE data is bad. Reject
124+ // negative values and check that the product of all dimensions doesn't
125+ // overflow ssize_t, which would otherwise abort in the TensorImpl ctor.
126+ // dim_order_to_stride() will validate dim_order.
127+ ssize_t numel = 1 ;
125128 for (flatbuffers::uoffset_t i = 0 ; i < dim; i++) {
126129 ET_CHECK_OR_RETURN_ERROR (
127130 sizes[i] >= 0 ,
128131 InvalidProgram,
129132 " Negative size[%zu] %" PRId32,
130133 static_cast <size_t >(i),
131134 sizes[i]);
135+ ET_CHECK_OR_RETURN_ERROR (
136+ sizes[i] == 0 || numel <= SSIZE_MAX / sizes[i],
137+ InvalidProgram,
138+ " Overflow computing numel at dim %zu" ,
139+ static_cast <size_t >(i));
140+ numel *= sizes[i];
132141 }
133142
134143 // We will remove strides from schema.
You can’t perform that action at this time.
0 commit comments