CVE-2020-15202: Other (NVD-CWE-Other)
value: In Tensorflow before versions 1.15.4, 2.0.3, 2.1.2, 2.2.1 and 2.3.1, the `Shard` API in TensorFlow expects the last argument to be a function taking two `int64` (i.e., `long long`) arguments. However, there are several places in TensorFlow where a lambda taking `int` or `int32` arguments is being used. In these cases, if the amount of work to be parallelized is large enough, integer truncation occurs. Depending on how the two arguments of the lambda are used, this can result in segfaults, read/write outside of heap allocated arrays, stack overflows, or data corruption. The issue is patched in commits 27b417360cbd671ef55915e4bb6bb06af8b8a832 and ca8c013b5e97b1373b3bb1c97ea655e69f31a575, and is released in TensorFlow versions 1.15.4, 2.0.3, 2.1.2, 2.2.1, or 2.3.1.
CWE description: NVD is only using a subset of CWE for mapping instead of the entire CWE, and the weakness type is not covered by that subset.
CVSS v3 base score: 9
Published: 2020-09-25T19:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/boosted_trees/prediction_ops.cc, tensorflow/core/kernels/image/crop_and_resize_op.cc, tensorflow/core/kernels/linalg/banded_triangular_solve_op.cc, tensorflow/core/kernels/nth_element_op.cc, tensorflow/core/kernels/parameterized_truncated_normal_op.cc, tensorflow/core/kernels/random_binomial_op.cc, tensorflow/core/kernels/random_poisson_op.cc, tensorflow/core/kernels/stateless_random_ops.cc, tensorflow/core/kernels/topk_op.cc
Commit date: 2020-09-18 17:56:30 -0700
Commit message: Prevent integer truncation from 64 to 32 bits.

The `tensorflow::Shard` functions last argument must be a 2 argument function where both arguments are `int64` (`long long`, 64 bits). However, there are usages where code passes in a function where arguments are `int` or `int32` (32 bits). In these cases, it is possible that the integer truncation would later cause a segfault or other unexpected behavior.

PiperOrigin-RevId: 332560414
Change-Id: Ief649406babc8d4f60b3e7a9d573cbcc5ce5b767

Diff: diff --git a/tensorflow/core/kernels/boosted_trees/prediction_ops.cc b/tensorflow/core/kernels/boosted_trees/prediction_ops.cc
index 19be606f184..e3a908d1b6b 100644
--- a/tensorflow/core/kernels/boosted_trees/prediction_ops.cc
+++ b/tensorflow/core/kernels/boosted_trees/prediction_ops.cc
@@ -119,11 +119,11 @@ class BoostedTreesTrainingPredictOp : public OpKernel {
     } else {
       output_tree_ids.setConstant(latest_tree);
       auto do_work = [&resource, &bucketized_features, &cached_tree_ids,
                       &cached_node_ids, &output_partial_logits,
                       &output_node_ids, latest_tree,
-                      this](int32 start, int32 end) {
+                      this](int64 start, int64 end) {
         for (int32 i = start; i < end; ++i) {
           int32 tree_id = cached_tree_ids(i);
           int32 node_id = cached_node_ids(i);
           std::vector<float> partial_tree_logits(logits_dimension_, 0.0);
 
@@ -235,11 +235,11 @@ class BoostedTreesPredictOp : public OpKernel {
       return;
     }
 
     const int32 last_tree = resource->num_trees() - 1;
     auto do_work = [&resource, &bucketized_features, &output_logits, last_tree,
-                    this](int32 start, int32 end) {
+                    this](int64 start, int64 end) {
       for (int32 i = start; i < end; ++i) {
         std::vector<float> tree_logits(logits_dimension_, 0.0);
         int32 tree_id = 0;
         int32 node_id = 0;
         while (true) {
@@ -338,11 +338,11 @@ class BoostedTreesExampleDebugOutputsOp : public OpKernel {
     // For each given example, traverse through all trees keeping track of the
     // features used to split and the associated logits at each point along the
     // path. Note: feature_ids has one less value than logits_path because the
     // first value of each logit path will be the bias.
     auto do_work = [&resource, &bucketized_features, &output_debug_info,
-                    last_tree](int32 start, int32 end) {
+                    last_tree](int64 start, int64 end) {
       for (int32 i = start; i < end; ++i) {
         // Proto to store debug outputs, per example.
         boosted_trees::DebugOutput example_debug_info;
         // Initial bias prediction. E.g., prediction based off training mean.
         const auto& tree_logits = resource->node_value(0, 0);
diff --git a/tensorflow/core/kernels/image/crop_and_resize_op.cc b/tensorflow/core/kernels/image/crop_and_resize_op.cc
index 1979b0514c6..5c196df9cfe 100644
--- a/tensorflow/core/kernels/image/crop_and_resize_op.cc
+++ b/tensorflow/core/kernels/image/crop_and_resize_op.cc
@@ -221,11 +221,11 @@ struct CropAndResize<CPUDevice, T> {
     const int crop_height = crops.dimension(1);
     const int crop_width = crops.dimension(2);
     const int depth = crops.dimension(3);
 
     // Sharding across boxes.
-    auto CropAndResizePerBox = [&](int start_box, int limit_box) {
+    auto CropAndResizePerBox = [&](int64 start_box, int64 limit_box) {
       for (int b = start_box; b < limit_box; ++b) {
         const float y1 = boxes(b, 0);
         const float x1 = boxes(b, 1);
         const float y2 = boxes(b, 2);
         const float x2 = boxes(b, 3);
@@ -447,11 +447,11 @@ struct CropAndResizeBackpropImage<CPUDevice, T> {
     const int crop_width = grads.dimension(2);
     const int depth = grads.dimension(3);
 
     grads_image.setZero();
 
-    auto CropAndResizeBackImgPerBox = [&](int start_box, int limit_box) {
+    auto CropAndResizeBackImgPerBox = [&](int64 start_box, int64 limit_box) {
       for (int b = start_box; b < limit_box; ++b) {
         const float y1 = boxes(b, 0);
         const float x1 = boxes(b, 1);
         const float y2 = boxes(b, 2);
         const float x2 = boxes(b, 3);
diff --git a/tensorflow/core/kernels/linalg/banded_triangular_solve_op.cc b/tensorflow/core/kernels/linalg/banded_triangular_solve_op.cc
index 6758dcf5b8b..bec28088ad1 100644
--- a/tensorflow/core/kernels/linalg/banded_triangular_solve_op.cc
+++ b/tensorflow/core/kernels/linalg/banded_triangular_solve_op.cc
@@ -191,11 +191,12 @@ struct LaunchBatchBandedTriangularSolve {
     OP_REQUIRES(context, min_abs_pivot > RealScalar(0),
                 errors::InvalidArgument("Input matrix is not invertible."));
 
     Shard(worker_threads.num_threads, worker_threads.workers, batch_size,
           cost_per_unit,
-          [&in_x, &in_y, adjoint, lower, &bcast, out](int start, int limit) {
+          [&in_x, &in_y, adjoint, lower, &bcast, out](int64 start,
+                                                      int64 limit) {
             SequentialBandedTriangularSolveKernel<Scalar>::Run(
                 in_x, in_y, lower, adjoint, bcast, out, start, limit);
           });
   }
 };
diff --git a/tensorflow/core/kernels/nth_element_op.cc b/tensorflow/core/kernels/nth_element_op.cc
index dced32ef7df..95328b3e469 100644
--- a/tensorflow/core/kernels/nth_element_op.cc
+++ b/tensorflow/core/kernels/nth_element_op.cc
@@ -93,11 +93,12 @@ struct NthElementFunctor<CPUDevice, T> {
     // then num_rows = d1*d2...dk-1, last_dim = dk.
     const int num_rows = output_tensor.NumElements();
     const int last_dim = input_tensor.dim_size(input_tensor.dims() - 1);
 
     // Allocate each row to different shard.
-    auto SubNthElement = [&, input, output, last_dim, n](int start, int limit) {
+    auto SubNthElement = [&, input, output, last_dim, n](int64 start,
+                                                         int64 limit) {
       // std::nth_element would rearrange the array, so we need a new buffer.
       std::vector<T> buf(last_dim);
 
       for (int b = start; b < limit; ++b) {
         // Copy from one row of elements to buffer
diff --git a/tensorflow/core/kernels/parameterized_truncated_normal_op.cc b/tensorflow/core/kernels/parameterized_truncated_normal_op.cc
index ba1fd280ce7..a63457551ac 100644
--- a/tensorflow/core/kernels/parameterized_truncated_normal_op.cc
+++ b/tensorflow/core/kernels/parameterized_truncated_normal_op.cc
@@ -68,12 +68,12 @@ struct TruncatedNormalFunctor<CPUDevice, T> {
     const T kStdDevsInsideBoundsToUseRandnSampler = T(1.3);
     auto worker_threads = *(ctx->device()->tensorflow_cpu_worker_threads());
 
     auto do_work = [samples_per_batch, num_elements, &ctx, &means, &stddevs,
                     &minvals, &maxvals, &gen, &output,
-                    kStdDevsInsideBoundsToUseRandnSampler](int start_batch,
-                                                           int limit_batch) {
+                    kStdDevsInsideBoundsToUseRandnSampler](int64 start_batch,
+                                                           int64 limit_batch) {
       // Capturing "gen" by-value would only make a copy for the _shared_
       // lambda.  Since we want to let each worker have its own copy, we pass
       // "gen" by reference and explicitly do a copy assignment here.
       random::PhiloxRandom gen_copy = gen;
       // Skip takes units of 128 bytes.  +3 is so rounding doesn't lead to
@@ -331,12 +331,12 @@ struct TruncatedNormalFunctorV2<CPUDevice, T> {
     const T kStdDevsInsideBoundsToUseRandnSampler = T(1.3);
     auto worker_threads = *(ctx->device()->tensorflow_cpu_worker_threads());
 
     auto do_work = [num_batches, samples_per_batch, &ctx, &bcast, &means,
                     &stddevs, &minvals, &maxvals, &gen, &output,
-                    kStdDevsInsideBoundsToUseRandnSampler](int start_output,
-                                                           int limit_output) {
+                    kStdDevsInsideBoundsToUseRandnSampler](int64 start_output,
+                                                           int64 limit_output) {
       // Capturing "gen" by-value would only make a copy for the _shared_
       // lambda.  Since we want to let each worker have its own copy, we pass
       // "gen" by reference and explicitly do a copy assignment here.
       random::PhiloxRandom gen_copy = gen;
       using Uniform = random::UniformDistribution<random::PhiloxRandom, T>;
diff --git a/tensorflow/core/kernels/random_binomial_op.cc b/tensorflow/core/kernels/random_binomial_op.cc
index d040cd93d00..e2487d68be8 100644
--- a/tensorflow/core/kernels/random_binomial_op.cc
+++ b/tensorflow/core/kernels/random_binomial_op.cc
@@ -182,11 +182,11 @@ struct RandomBinomialFunctor<CPUDevice, T, U> {
 
     // The output layout is [B1, ... Bk, H1, ... Hm]. We have [B1, ... Bk] for
     // the sample shape and [H1, ... Hm] for the batch shape of the samples.
     // We have B1 * ... * Bk samples per batch member we need.
     auto DoWork = [num_batches, samples_per_batch, &bcast, &counts, &probs,
-                   &gen, &output](int start_output, int limit_output) {
+                   &gen, &output](int64 start_output, int64 limit_output) {
       // Vectorized intermediate calculations for uniform rejection sampling.
       // We always generate at most 4 samples.
       Eigen::array<T, 4> z;
       Eigen::array<T, 4> g;
       const bool should_bcast = bcast.IsBroadcastingRequired();
diff --git a/tensorflow/core/kernels/random_poisson_op.cc b/tensorflow/core/kernels/random_poisson_op.cc
index aa9a0bfe214..dcb7d6b0f0e 100644
--- a/tensorflow/core/kernels/random_poisson_op.cc
+++ b/tensorflow/core/kernels/random_poisson_op.cc
@@ -95,11 +95,11 @@ struct PoissonFunctor<CPUDevice, T, U> {
     // avoid a couple flops which can be done on a per-rate basis.
 
     typedef random::UniformDistribution<random::PhiloxRandom, CT> Uniform;
 
     auto DoWork = [num_samples, num_rate, &rng, samples_flat, rate_flat](
-                      int start_output, int limit_output) {
+                      int64 start_output, int64 limit_output) {
       // Capturing "rng" by value would only make a copy for the _shared_
       // lambda.  Since we want to let each worker have its own copy, we pass
       // "rng" by reference and explicitly do a copy assignment.
 
       Uniform uniform;
diff --git a/tensorflow/core/kernels/stateless_random_ops.cc b/tensorflow/core/kernels/stateless_random_ops.cc
index 91944db1f62..0d4488b538f 100644
--- a/tensorflow/core/kernels/stateless_random_ops.cc
+++ b/tensorflow/core/kernels/stateless_random_ops.cc
@@ -250,11 +250,11 @@ class StatelessRandomGammaOp : public StatelessRandomOpBase {
 
     // We partition work first across alphas then across samples-per-alpha to
     // avoid a couple flops which can be done on a per-alpha basis.
 
     auto DoWork = [samples_per_alpha, num_alphas, &random, samples_flat,
-                   alpha_flat](int start_output, int limit_output) {
+                   alpha_flat](int64 start_output, int64 limit_output) {
       // Capturing "random" by-value would only make a copy for the _shared_
       // lambda.  Since we want to let each worker have its own copy, we pass
       // "random" by reference and explicitly do a copy assignment.
 
       using Eigen::numext::exp;
diff --git a/tensorflow/core/kernels/topk_op.cc b/tensorflow/core/kernels/topk_op.cc
index 50325b7bcfe..3390bd07308 100644
--- a/tensorflow/core/kernels/topk_op.cc
+++ b/tensorflow/core/kernels/topk_op.cc
@@ -134,11 +134,11 @@ struct TopKFunctor<CPUDevice, T> {
       }
 
       return Status::OK();
     }
 
-    auto SortIndices = [&](int start_batch, int limit_batch) {
+    auto SortIndices = [&](int64 start_batch, int64 limit_batch) {
       for (int32 b = start_batch; b < limit_batch; ++b) {
         const T* input_data = &input(b, 0);
         const auto stable_comp = [input_data](const int32 a, const int32 b) {
           if (input_data[b] < input_data[a]) {
             return true;

Vulnerable code:                       this](int32 start, int32 end) {
                    this](int32 start, int32 end) {
                    last_tree](int32 start, int32 end) {
    auto CropAndResizePerBox = [&](int start_box, int limit_box) {
    auto CropAndResizeBackImgPerBox = [&](int start_box, int limit_box) {
          [&in_x, &in_y, adjoint, lower, &bcast, out](int start, int limit) {
    auto SubNthElement = [&, input, output, last_dim, n](int start, int limit) {
                    kStdDevsInsideBoundsToUseRandnSampler](int start_batch,
                                                           int limit_batch) {
                    kStdDevsInsideBoundsToUseRandnSampler](int start_output,
                                                           int limit_output) {
                   &gen, &output](int start_output, int limit_output) {
                      int start_output, int limit_output) {
                   alpha_flat](int start_output, int limit_output) {
    auto SortIndices = [&](int start_batch, int limit_batch) {
Fixed code:                       this](int64 start, int64 end) {
                    this](int64 start, int64 end) {
                    last_tree](int64 start, int64 end) {
    auto CropAndResizePerBox = [&](int64 start_box, int64 limit_box) {
    auto CropAndResizeBackImgPerBox = [&](int64 start_box, int64 limit_box) {
          [&in_x, &in_y, adjoint, lower, &bcast, out](int64 start,
                                                      int64 limit) {
    auto SubNthElement = [&, input, output, last_dim, n](int64 start,
                                                         int64 limit) {
                    kStdDevsInsideBoundsToUseRandnSampler](int64 start_batch,
                                                           int64 limit_batch) {
                    kStdDevsInsideBoundsToUseRandnSampler](int64 start_output,
                                                           int64 limit_output) {
                   &gen, &output](int64 start_output, int64 limit_output) {
                      int64 start_output, int64 limit_output) {
                   alpha_flat](int64 start_output, int64 limit_output) {
    auto SortIndices = [&](int64 start_batch, int64 limit_batch) {