CVE-2022-35984: Reachable Assertion (CWE-617)
value: TensorFlow is an open source platform for machine learning. `ParameterizedTruncatedNormal` assumes `shape` is of type `int32`. A valid `shape` of type `int64` results in a mismatched type `CHECK` fail that can be used to trigger a denial of service attack. We have patched the issue in GitHub commit 72180be03447a10810edca700cbc9af690dfeb51. The fix will be included in TensorFlow 2.10.0. We will also cherrypick this commit on TensorFlow 2.9.1, TensorFlow 2.8.1, and TensorFlow 2.7.2, as these are also affected and still in supported range. There are no known workarounds for this issue.
CWE description: The product contains an assert() or similar statement that can be triggered by an attacker, which leads to an application exit or other behavior that is more severe than necessary.
CVSS v3 base score: 7.5
Published: 2022-09-16T22:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/parameterized_truncated_normal_op.cc, tensorflow/python/kernel_tests/random/parameterized_truncated_normal_op_test.py
Commit date: 2022-06-22 22:37:01 -0700
Commit message: Fix tensor shape dtype bug in parameterized_truncated_normal.

The original made the assumption that the dtype is int32 when it could
also be int64 - leading to a crash due to mismatched type.

PiperOrigin-RevId: 456685278

Diff: diff --git a/tensorflow/core/kernels/parameterized_truncated_normal_op.cc b/tensorflow/core/kernels/parameterized_truncated_normal_op.cc
index 24b7e3f4ebdd5..a007d37c4e290 100644
--- a/tensorflow/core/kernels/parameterized_truncated_normal_op.cc
+++ b/tensorflow/core/kernels/parameterized_truncated_normal_op.cc
@@ -30,10 +30,11 @@ limitations under the License.
 
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_util.h"
 #include "tensorflow/core/kernels/stateless_random_ops.h"
 #include "tensorflow/core/lib/random/random_distributions.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/util/guarded_philox_random.h"
 #include "tensorflow/core/util/work_sharder.h"
@@ -628,24 +629,22 @@ class ParameterizedTruncatedNormalOp : public OpKernel {
         errors::InvalidArgument("Input shape should be a vector, got shape: ",
                                 shape_tensor.shape().DebugString()));
     OP_REQUIRES(ctx, shape_tensor.NumElements() > 0,
                 errors::InvalidArgument("Shape tensor must not be empty, got ",
                                         shape_tensor.DebugString()));
-    int32_t num_batches = shape_tensor.flat<int32>()(0);
+    TensorShape tensor_shape;
+    OP_REQUIRES_OK(ctx, tensor::MakeShape(shape_tensor, &tensor_shape));
 
+    int32_t num_batches = tensor_shape.dim_size(0);
     int32_t samples_per_batch = 1;
-    const int32_t num_dims = shape_tensor.dim_size(0);
+    const int32_t num_dims = tensor_shape.dims();
     for (int32_t i = 1; i < num_dims; i++) {
-      samples_per_batch *= shape_tensor.flat<int32>()(i);
+      samples_per_batch *= tensor_shape.dim_size(i);
     }
     const int32_t num_elements = num_batches * samples_per_batch;
 
     // Allocate the output before fudging num_batches and samples_per_batch.
-    auto shape_vec = shape_tensor.flat<int32>();
-    TensorShape tensor_shape;
-    OP_REQUIRES_OK(ctx, TensorShapeUtils::MakeShape(
-                            shape_vec.data(), shape_vec.size(), &tensor_shape));
     Tensor* samples_tensor;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, tensor_shape, &samples_tensor));
 
     // Parameters must be 0-d or 1-d.
     OP_REQUIRES(ctx, means_tensor.dims() <= 1,
diff --git a/tensorflow/python/kernel_tests/random/parameterized_truncated_normal_op_test.py b/tensorflow/python/kernel_tests/random/parameterized_truncated_normal_op_test.py
index a1c6072abfda4..8ad859230edcc 100644
--- a/tensorflow/python/kernel_tests/random/parameterized_truncated_normal_op_test.py
+++ b/tensorflow/python/kernel_tests/random/parameterized_truncated_normal_op_test.py
@@ -301,10 +301,33 @@ class ParameterizedTruncatedNormalTest(test.TestCase):
       assert (~np.isnan(samples)).all()
       assert (~np.isnan(samples_stateless)).all()
       self.assertAllGreater(samples, 0.)
       self.assertAllGreater(samples_stateless, 0.)
 
+  def testShapeTypes(self):
+    for shape_dtype in [np.int32, np.int64]:
+      shape = np.array([1000], dtype=shape_dtype)
+      sample_op = random_ops.parameterized_truncated_normal(
+          shape=shape, means=0.0, stddevs=0.1, minvals=-1., maxvals=1.)
+      new_seed = random_ops.random_uniform([2],
+                                           seed=1234,
+                                           minval=0,
+                                           maxval=(2**31 - 1),
+                                           dtype=np.int32)
+      sample_op_stateless = stateless.stateless_parameterized_truncated_normal(
+          shape=shape,
+          seed=new_seed,
+          means=0.0,
+          stddevs=0.1,
+          minvals=-1.,
+          maxvals=1.)
+
+      samples = self.evaluate(sample_op)
+      stateless_samples = self.evaluate(sample_op_stateless)
+      self.assertAllEqual(samples.shape, shape)
+      self.assertAllEqual(stateless_samples.shape, shape)
+
   def testStatelessParameterizedTruncatedNormalHasGrads(self):
     mean = variables.Variable(0.01)
     stddev = variables.Variable(1.)
     minval = variables.Variable(-1.)
     maxval = variables.Variable(1.)

Vulnerable code:     int32_t num_batches = shape_tensor.flat<int32>()(0);
    const int32_t num_dims = shape_tensor.dim_size(0);
      samples_per_batch *= shape_tensor.flat<int32>()(i);
    auto shape_vec = shape_tensor.flat<int32>();
    TensorShape tensor_shape;
    OP_REQUIRES_OK(ctx, TensorShapeUtils::MakeShape(
                            shape_vec.data(), shape_vec.size(), &tensor_shape));
Fixed code: #include "tensorflow/core/framework/tensor_util.h"
    TensorShape tensor_shape;
    OP_REQUIRES_OK(ctx, tensor::MakeShape(shape_tensor, &tensor_shape));
    int32_t num_batches = tensor_shape.dim_size(0);
    const int32_t num_dims = tensor_shape.dims();
      samples_per_batch *= tensor_shape.dim_size(i);
  def testShapeTypes(self):
    for shape_dtype in [np.int32, np.int64]:
      shape = np.array([1000], dtype=shape_dtype)
      sample_op = random_ops.parameterized_truncated_normal(
          shape=shape, means=0.0, stddevs=0.1, minvals=-1., maxvals=1.)
      new_seed = random_ops.random_uniform([2],
                                           seed=1234,
                                           minval=0,
                                           maxval=(2**31 - 1),
                                           dtype=np.int32)
      sample_op_stateless = stateless.stateless_parameterized_truncated_normal(
          shape=shape,
          seed=new_seed,
          means=0.0,
          stddevs=0.1,
          minvals=-1.,
          maxvals=1.)

      samples = self.evaluate(sample_op)
      stateless_samples = self.evaluate(sample_op_stateless)
      self.assertAllEqual(samples.shape, shape)
      self.assertAllEqual(stateless_samples.shape, shape)