Skip to content

Commit e1ae2a4

Browse files
dbortfacebook-github-bot
authored andcommitted
Validate tensor sizes during method load
Summary: If a tensor's `sizes` array contains a negative value, the TensorImpl ctor will panic. Validate sizes before creating the TensorImpl to treat this as a non-fatal error when loading malformed PTE files. Differential Revision: D68180029
1 parent 9181e83 commit e1ae2a4

File tree

1 file changed

+13
-0
lines changed

1 file changed

+13
-0
lines changed

runtime/executor/tensor_parser_portable.cpp

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,19 @@ Result<Tensor> parseTensor(
101101
sizes = const_cast<exec_aten::SizesType*>(serialized_sizes);
102102
dim_order = const_cast<exec_aten::DimOrderType*>(serialized_dim_order);
103103
}
104+
// Validate sizes before using them in case the PTE data is bad. We can't
105+
// detect bad positive values, but we can reject negative values, which would
106+
// otherwise panic in the TensorImpl ctor. dim_order_to_stride() will validate
107+
// dim_order.
108+
for (int i = 0; i < dim; i++) {
109+
ET_CHECK_OR_RETURN_ERROR(
110+
sizes[i] >= 0,
111+
InvalidProgram,
112+
"Negative size[%d] %" PRId32,
113+
i,
114+
sizes[i]);
115+
}
116+
104117
// We will remove strides from schema.
105118
// Allocating strides buffer here and populating it.
106119
// In subsequent diffs we can remove strides accessor, however this

0 commit comments

Comments
 (0)