Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent CHECK-fail when building reference tensor.
The tensor constructor does not allow reference dtypes, as these should not show up explicitly. However, when passed these invalid types instead of building an invalid object the constructor crashes via a `CHECK`-fail. We have a static builder that properly handles this case but is not applicable given current usage.

Instead, before calling the constructor, we can check that the dtype is not a reference type and return an error otherwise, given that the dtype is user controlled so malicious users can trigger denial of service.

PiperOrigin-RevId: 409662503
Change-Id: I5892f831fde7f276cd7ab34519cf6b8061c71a59
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Nov 13, 2021
1 parent af2cab9 commit 6b5adc0
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions tensorflow/core/grappler/optimizers/constant_folding.cc
Expand Up @@ -1363,6 +1363,11 @@ Status ConstantFolding::EvaluateOneFoldable(const NodeDef& node,
input_tensor.ToString(),
" has a dtype of DT_INVALID."));
}
if (IsRefType(raw_val.dtype())) {
return errors::InvalidArgument(
"Not allowed to construct a tensor with reference dtype, got ",
DataTypeString(raw_val.dtype()));
}
Tensor* value = new Tensor(raw_val.dtype(), raw_val.tensor_shape());
if (!value->FromProto(raw_val)) {
delete (value);
Expand Down

0 comments on commit 6b5adc0

Please sign in to comment.