The implementation of shape inference for ConcatV2 can be used to trigger a denial of service attack via a segfault caused by a type confusion:
import tensorflow as tf
@tf.function
def test():
y = tf.raw_ops.ConcatV2(
values=[[1,2,3],[4,5,6]],
axis = 0xb500005b)
return y
test()
The axis argument is translated into concat_dim in the ConcatShapeHelper helper function. Then, a value for min_rank is computed based on concat_dim. This is then used to validate that the values tensor has at least the required rank:
int64_t concat_dim;
if (concat_dim_t->dtype() == DT_INT32) {
concat_dim = static_cast<int64_t>(concat_dim_t->flat<int32>()(0));
} else {
concat_dim = concat_dim_t->flat<int64_t>()(0);
}
// Minimum required number of dimensions.
const int min_rank = concat_dim < 0 ? -concat_dim : concat_dim + 1;
// ...
ShapeHandle input = c->input(end_value_index - 1);
TF_RETURN_IF_ERROR(c->WithRankAtLeast(input, min_rank, &input));
However, WithRankAtLeast receives the lower bound as a 64-bits value and then compares it against the maximum 32-bits integer value that could be represented:
Status InferenceContext::WithRankAtLeast(ShapeHandle shape, int64_t rank,
ShapeHandle* out) {
if (rank > kint32max) {
return errors::InvalidArgument("Rank cannot exceed kint32max");
}
// ...
}
Due to the fact that min_rank is a 32-bits value and the value of axis, the rank argument is a negative value, so the error check is bypassed.
We have patched the issue in GitHub commit...
2.5.32.6.32.7.02.7.12.5.32.6.32.7.02.7.12.5.32.6.32.7.02.7.1Exploitability
AV:NAC:LAT:NPR:LUI:NVulnerable System
VC:NVI:NVA:HSubsequent System
SC:NSI:NSA:N7.1/CVSS:4.0/AV:N/AC:L/AT:N/PR:L/UI:N/VC:N/VI:N/VA:H/SC:N/SI:N/SA:N