CVE-2021-29548 (C++, published 2021-05-14T20:15Z)
Description: value: TensorFlow is an end-to-end open source platform for machine learning. An attacker can cause a runtime division by zero error and denial of service in `tf.raw_ops.QuantizedBatchNormWithGlobalNormalization`. This is because the implementation(https://github.com/tensorflow/tensorflow/blob/55a97caa9e99c7f37a0bbbeb414dc55553d3ae7f/tensorflow/core/kernels/quantized_batch_norm_op.cc) does not validate all constraints specified in the op's contract(https://www.tensorflow.org/api_docs/python/tf/raw_ops/QuantizedBatchNormWithGlobalNormalization). The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.4, as these are also affected and still in supported range.
Severity: LOW (CVSS v3 base score: 5.5)
Weakness: CWE-369 \u2014 Divide By Zero
The product divides a value by zero.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/core/kernels/quantized_batch_norm_op.cc
Commit message: Add missing validation in `QuantizedBatchNormWithGlobalNormalization`

PiperOrigin-RevId: 370123451
Change-Id: Id234d6dab1ec21230bb8e503dba30f899af87f33


Vulnerable code:
    const float input_min = context->input(1).flat<float>()(0);
    const float input_max = context->input(2).flat<float>()(0);
    const float mean_min = context->input(4).flat<float>()(0);
    const float mean_max = context->input(5).flat<float>()(0);
    const float var_min = context->input(7).flat<float>()(0);
    const float var_max = context->input(8).flat<float>()(0);
    const float beta_min = context->input(10).flat<float>()(0);
    const float beta_max = context->input(11).flat<float>()(0);
    const float gamma_min = context->input(13).flat<float>()(0);
    const float gamma_max = context->input(14).flat<float>()(0);

Fixed code:
    const auto& input_min_tensor = context->input(1);
    OP_REQUIRES(context, input_min_tensor.NumElements() == 1,
                errors::InvalidArgument("input_min must have 1 element"));
    const float input_min = input_min_tensor.flat<float>()(0);
    const auto& input_max_tensor = context->input(2);
    OP_REQUIRES(context, input_max_tensor.NumElements() == 1,
                errors::InvalidArgument("input_max must have 1 element"));
    const float input_max = input_max_tensor.flat<float>()(0);
    const auto& mean_min_tensor = context->input(4);
    OP_REQUIRES(context, mean_min_tensor.NumElements() == 1,
                errors::InvalidArgument("mean_min must have 1 element"));
    const float mean_min = mean_min_tensor.flat<float>()(0);
    const auto& mean_max_tensor = context->input(5);
    OP_REQUIRES(context, mean_max_tensor.NumElements() == 1,
                errors::InvalidArgument("mean_max must have 1 element"));
    const float mean_max = mean_max_tensor.flat<float>()(0);
    const auto& var_min_tensor = context->input(7);
    OP_REQUIRES(context, var_min_tensor.NumElements() == 1,
                errors::InvalidArgument("var_min must have 1 element"));
    const float var_min = var_min_tensor.flat<float>()(0);
    const auto& var_max_tensor = context->input(8);
    OP_REQUIRES(context, var_max_tensor.NumElements() == 1,
                errors::InvalidArgument("var_max must have 1 element"));
    const float var_max = var_max_tensor.flat<float>()(0);
    const auto& beta_min_tensor = context->input(10);
    OP_REQUIRES(context, beta_min_tensor.NumElements() == 1,
                errors::InvalidArgument("beta_min must have 1 element"));
    const float beta_min = beta_min_tensor.flat<float>()(0);
    const auto& beta_max_tensor = context->input(11);
    OP_REQUIRES(context, beta_max_tensor.NumElements() == 1,
                errors::InvalidArgument("beta_max must have 1 element"));
    const float beta_max = beta_max_tensor.flat<float>()(0);
    const auto& gamma_min_tensor = context->input(13);
    OP_REQUIRES(context, gamma_min_tensor.NumElements() == 1,
                errors::InvalidArgument("gamma_min must have 1 element"));
    const float gamma_min = gamma_min_tensor.flat<float>()(0);
    const auto& gamma_max_tensor = context->input(14);
    OP_REQUIRES(context, gamma_max_tensor.NumElements() == 1,
                errors::InvalidArgument("gamma_max must have 1 element"));
    const float gamma_max = gamma_max_tensor.flat<float>()(0);
    OP_REQUIRES(context, mean.NumElements() > 1,
                errors::InvalidArgument("Must have at least a mean value",
                                        gamma.shape().DebugString()));
    OP_REQUIRES(context, mean.NumElements() > 1,
                errors::InvalidArgument("Must have at least a mean value"));
    const auto last_dim = input.shape().dims() - 1;
    OP_REQUIRES(context,
                mean.shape().dim_size(0) == input.shape().dim_size(last_dim),
                errors::InvalidArgument("Must provide as many means as the "
                                        "last dimension of the input tensor: ",
                                        mean.shape().DebugString(), " vs. ",
                                        input.shape().DebugString()));
    OP_REQUIRES(
        context, mean.shape().dim_size(0) == var.shape().dim_size(0),
        errors::InvalidArgument(
            "Mean and variance tensors must have the same shape: ",
            mean.shape().DebugString(), " vs. ", var.shape().DebugString()));
    OP_REQUIRES(
        context, mean.shape().dim_size(0) == beta.shape().dim_size(0),
        errors::InvalidArgument(
            "Mean and beta tensors must have the same shape: ",
            mean.shape().DebugString(), " vs. ", beta.shape().DebugString()));
    OP_REQUIRES(
        context, mean.shape().dim_size(0) == gamma.shape().dim_size(0),
        errors::InvalidArgument(
            "Mean and gamma tensors must have the same shape: ",
            mean.shape().DebugString(), " vs. ", gamma.shape().DebugString()));