Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix tf.compat.v1.placeholder_with_default vulnerability with quantize…
…d types.

When iterating through the tensor to extract shape values, an underlying missing kernel
(`StridedSlice` for quantized types) causes an error, which then results in a `nullptr`
being passed to `ParseDimensionValue()`, causing a segfault.

The `nullptr` check allows the missing kernel error to propagate.
Adding the missing kernel registrations allows the shape values
to be extracted successfully.

PiperOrigin-RevId: 445045957
  • Loading branch information
cantonios authored and tensorflower-gardener committed Apr 28, 2022
1 parent 527a612 commit 237822b
Show file tree
Hide file tree
Showing 3 changed files with 6 additions and 2 deletions.
1 change: 1 addition & 0 deletions tensorflow/core/kernels/strided_slice_op.cc
Expand Up @@ -431,6 +431,7 @@ class StridedSliceAssignOp : public OpKernel {
StridedSliceAssignOp<CPUDevice, type, true>)

TF_CALL_ALL_TYPES(REGISTER_STRIDED_SLICE);
TF_CALL_QUANTIZED_TYPES(REGISTER_STRIDED_SLICE);

#undef REGISTER_STRIDED_SLICE

Expand Down
2 changes: 1 addition & 1 deletion tensorflow/core/kernels/strided_slice_op_impl.h
Expand Up @@ -292,7 +292,7 @@ TF_CALL_GPU_ALL_TYPES(DECLARE_FOR_N_GPU);
#endif // END GOOGLE_CUDA || TENSORFLOW_USE_ROCM

TF_CALL_ALL_TYPES(DECLARE_FOR_N_CPU);

TF_CALL_QUANTIZED_TYPES(DECLARE_FOR_N_CPU);

#undef INSTANTIATE
#undef DECLARE_FOR_N_CPU
Expand Down
5 changes: 4 additions & 1 deletion tensorflow/python/eager/pywrap_tfe_src.cc
Expand Up @@ -688,9 +688,12 @@ bool SetOpAttrScalar(TFE_Context* ctx, TFE_Op* op, const char* key,
for (int i = 0; i < num_dims; ++i) {
tensorflow::Safe_PyObjectPtr inner_py_value(
PySequence_ITEM(py_value, i));
// If an error is generated when iterating through object, we can
// sometimes get a nullptr.
if (inner_py_value.get() == Py_None) {
dims[i] = -1;
} else if (!ParseDimensionValue(key, inner_py_value.get(), status,
} else if (inner_py_value.get() == nullptr ||
!ParseDimensionValue(key, inner_py_value.get(), status,
&dims[i])) {
return false;
}
Expand Down

0 comments on commit 237822b

Please sign in to comment.