CVE-2021-37662: Access of Uninitialized Pointer (CWE-824)
value: TensorFlow is an end-to-end open source platform for machine learning. In affected versions an attacker can generate undefined behavior via a reference binding to nullptr in `BoostedTreesCalculateBestGainsPerFeature` and similar attack can occur in `BoostedTreesCalculateBestFeatureSplitV2`. The [implementation](https://github.com/tensorflow/tensorflow/blob/84d053187cb80d975ef2b9684d4b61981bca0c41/tensorflow/core/kernels/boosted_trees/stats_ops.cc) does not validate the input values. We have patched the issue in GitHub commit 9c87c32c710d0b5b53dc6fd3bfde4046e1f7a5ad and in commit 429f009d2b2c09028647dd4bb7b3f6f414bbaad7. The fix will be included in TensorFlow 2.6.0. We will also cherrypick this commit on TensorFlow 2.5.1, TensorFlow 2.4.3, and TensorFlow 2.3.4, as these are also affected and still in supported range.
CWE description: The product accesses or uses a pointer that has not been initialized.
CVSS v3 base score: 7.8
Published: 2021-08-12T21:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/boosted_trees/stats_ops.cc
Commit date: 2021-07-28 13:43:55 -0700
Commit message: Add remaining missing validation to `BoostedTreesCalculateBestFeatureSplit`

PiperOrigin-RevId: 387423006
Change-Id: I8eaf30efb223011519e60707bfa751b275d3a443

Diff: diff --git a/tensorflow/core/kernels/boosted_trees/stats_ops.cc b/tensorflow/core/kernels/boosted_trees/stats_ops.cc
index 2636909855a38..60c1d191f5232 100644
--- a/tensorflow/core/kernels/boosted_trees/stats_ops.cc
+++ b/tensorflow/core/kernels/boosted_trees/stats_ops.cc
@@ -12,18 +12,20 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
 #include <limits>
+#include <string>
 #include <vector>
 
 #include "third_party/eigen3/Eigen/Core"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/kernels/boosted_trees/boosted_trees.pb.h"
 #include "tensorflow/core/kernels/boosted_trees/tree_helper.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/logging.h"
 
 namespace tensorflow {
 
 using Matrix =
@@ -252,16 +254,22 @@ class BoostedTreesCalculateBestFeatureSplitOp : public OpKernel {
 
   void Compute(OpKernelContext* const context) override {
     // node_id_range
     const Tensor* node_id_range_t;
     OP_REQUIRES_OK(context, context->input("node_id_range", &node_id_range_t));
+    OP_REQUIRES(
+        context, node_id_range_t->NumElements() == 2,
+        errors::InvalidArgument("node_id_range argument must have shape [2]"));
     const auto node_id_range = node_id_range_t->vec<int32>();
     const int32_t node_id_first = node_id_range(0);  // inclusive
     const int32_t node_id_last = node_id_range(1);   // exclusive
 
     const Tensor* stats_summary_t;
     OP_REQUIRES_OK(context, context->input("stats_summary", &stats_summary_t));
+    OP_REQUIRES(
+        context, stats_summary_t->shape().dims() == 4,
+        errors::InvalidArgument("stats_summary argument must have rank 4"));
     TTypes<float, 4>::ConstTensor stats_summary =
         stats_summary_t->tensor<float, 4>();
     const int32_t feature_dims = stats_summary_t->dim_size(1);
     // The last bucket is for default/missing value.
     const int32_t num_buckets = stats_summary_t->dim_size(2) - 1;
@@ -270,39 +278,49 @@ class BoostedTreesCalculateBestFeatureSplitOp : public OpKernel {
     DCHECK_GT(hessian_dim, 0);
     DCHECK_LE(hessian_dim, logits_dim * logits_dim);
 
     const Tensor* l1_t;
     OP_REQUIRES_OK(context, context->input("l1", &l1_t));
+    OP_REQUIRES(context, l1_t->NumElements() == 1,
+                errors::InvalidArgument("l1 argument must be a scalar"));
     const auto l1 = l1_t->scalar<float>()();
     DCHECK_GE(l1, 0);
     if (logits_dim_ > 1) {
       // Multi-class L1 regularization not supported yet.
       DCHECK_EQ(l1, 0);
     }
 
     const Tensor* l2_t;
     OP_REQUIRES_OK(context, context->input("l2", &l2_t));
+    OP_REQUIRES(context, l2_t->NumElements() == 1,
+                errors::InvalidArgument("l2 argument must be a scalar"));
     const auto l2 = l2_t->scalar<float>()();
     DCHECK_GE(l2, 0);
 
     const Tensor* tree_complexity_t;
     OP_REQUIRES_OK(context,
                    context->input("tree_complexity", &tree_complexity_t));
+    OP_REQUIRES(
+        context, tree_complexity_t->NumElements() == 1,
+        errors::InvalidArgument("tree_complexity argument must be a scalar"));
     const auto tree_complexity = tree_complexity_t->scalar<float>()();
 
     const Tensor* min_node_weight_t;
     OP_REQUIRES_OK(context,
                    context->input("min_node_weight", &min_node_weight_t));
+    OP_REQUIRES(
+        context, min_node_weight_t->NumElements() == 1,
+        errors::InvalidArgument("min_node_weight argument must be a scalar"));
     const auto min_node_weight = min_node_weight_t->scalar<float>()();
 
     std::vector<int32> output_node_ids;
     std::vector<float> output_gains;
     std::vector<int32> output_feature_dimensions;
     std::vector<int32> output_thresholds;
     std::vector<Eigen::VectorXf> output_left_node_contribs;
     std::vector<Eigen::VectorXf> output_right_node_contribs;
-    std::vector<string> output_split_types;
+    std::vector<std::string> output_split_types;
 
     // TODO(tanzheny) parallelize the computation.
     // Iterate each node and find the best gain per node.
     for (int32_t node_id = node_id_first; node_id < node_id_last; ++node_id) {
       float best_gain = std::numeric_limits<float>::lowest();

Vulnerable code:     std::vector<string> output_split_types;
Fixed code: #include <string>
#include "tensorflow/core/platform/errors.h"
    OP_REQUIRES(
        context, node_id_range_t->NumElements() == 2,
        errors::InvalidArgument("node_id_range argument must have shape [2]"));
    OP_REQUIRES(
        context, stats_summary_t->shape().dims() == 4,
        errors::InvalidArgument("stats_summary argument must have rank 4"));
    OP_REQUIRES(context, l1_t->NumElements() == 1,
                errors::InvalidArgument("l1 argument must be a scalar"));
    OP_REQUIRES(context, l2_t->NumElements() == 1,
                errors::InvalidArgument("l2 argument must be a scalar"));
    OP_REQUIRES(
        context, tree_complexity_t->NumElements() == 1,
        errors::InvalidArgument("tree_complexity argument must be a scalar"));
    OP_REQUIRES(
        context, min_node_weight_t->NumElements() == 1,
        errors::InvalidArgument("min_node_weight argument must be a scalar"));
    std::vector<std::string> output_split_types;