et check

Author: Github Executorch

runtime/core/portable_type/tensor_impl.cpp

@@ -12,6 +12,7 @@
 #include <climits>
 #include <cstdint>
 
+
 #include <c10/util/irange.h>
 
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
@@ -40,7 +41,7 @@ ssize_t compute_numel(const TensorImpl::SizesType* sizes, ssize_t dim) {
         static_cast<ssize_t>(sizes[i]),
         i);
     ET_CHECK_MSG(
-        sizes[i] == 0 || numel <= SSIZE_MAX / sizes[i],
+        sizes[i] == 0 || numel <= (ssize_t)(SIZE_MAX / 2) / sizes[i],
         "Overflow computing numel: %zd * %zd would overflow ssize_t at dimension %zd",
         numel,
         static_cast<ssize_t>(sizes[i]),

runtime/executor/tensor_parser_portable.cpp

@@ -8,6 +8,8 @@
 
 #include <executorch/runtime/executor/tensor_parser.h>
 
+#include <climits>
+
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
@@ -118,17 +120,24 @@ Result<Tensor> parseTensor(
     dim_order =
         const_cast<executorch::aten::DimOrderType*>(serialized_dim_order);
   }
-  // Validate sizes before using them in case the PTE data is bad. We can't
-  // detect bad positive values, but we can reject negative values, which would
-  // otherwise panic in the TensorImpl ctor. dim_order_to_stride() will validate
-  // dim_order.
+  // Validate sizes before using them in case the PTE data is bad. Reject
+  // negative values and check that the product of all dimensions doesn't
+  // overflow ssize_t, which would otherwise abort in the TensorImpl ctor.
+  // dim_order_to_stride() will validate dim_order.
+  ssize_t numel = 1;
   for (flatbuffers::uoffset_t i = 0; i < dim; i++) {
     ET_CHECK_OR_RETURN_ERROR(
         sizes[i] >= 0,
         InvalidProgram,
         "Negative size[%zu] %" PRId32,
         static_cast<size_t>(i),
         sizes[i]);
+    ET_CHECK_OR_RETURN_ERROR(
+        sizes[i] == 0 || numel <= (ssize_t)(SIZE_MAX / 2) / sizes[i],
+        InvalidProgram,
+        "Overflow computing numel at dim %zu",
+        static_cast<size_t>(i));
+    numel *= sizes[i];
   }
 
   // We will remove strides from schema.