CVE-2022-21732: Allocation of Resources Without Limits or Throttling (CWE-770)
value: Tensorflow is an Open Source Machine Learning Framework. The implementation of `ThreadPoolHandle` can be used to trigger a denial of service attack by allocating too much memory. This is because the `num_threads` argument is only checked to not be negative, but there is no upper bound on its value. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, as these are also affected and still in supported range.
CWE description: The product allocates a reusable resource or group of resources on behalf of an actor without imposing any restrictions on the size or number of resources that can be allocated, in violation of the intended security policy for that actor.
CVSS v3 base score: 6.5
Published: 2022-02-03T12:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/data/experimental/threadpool_dataset_op.cc
Commit date: 2021-11-18 16:14:53 -0800
Commit message: [tf.data] Set limit on number of threads used in threadpool_dataset.

PiperOrigin-RevId: 410922677
Change-Id: Ib25814a99043ab10805b5d2d7088ae0e0b7b04fd

Diff: diff --git a/tensorflow/core/kernels/data/experimental/threadpool_dataset_op.cc b/tensorflow/core/kernels/data/experimental/threadpool_dataset_op.cc
index c58e5f68e6c6c..fefaf8249b76b 100644
--- a/tensorflow/core/kernels/data/experimental/threadpool_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/threadpool_dataset_op.cc
@@ -37,10 +37,26 @@ namespace experimental {
     MaxIntraOpParallelismDatasetOp::kDatasetOp;
 /* static */ constexpr const char* const
     PrivateThreadPoolDatasetOp::kDatasetType;
 /* static */ constexpr const char* const PrivateThreadPoolDatasetOp::kDatasetOp;
 
+namespace {
+// To prevent integer overflow issues when allocating threadpool memory for an
+// unreasonable number of threads.
+constexpr int kThreadLimit = 65536;
+
+Status ValidateNumThreads(int32_t num_threads) {
+  if (num_threads < 0) {
+    return errors::InvalidArgument("`num_threads` must be >= 0");
+  }
+  if (num_threads >= kThreadLimit) {
+    return errors::InvalidArgument("`num_threads` must be < ", kThreadLimit);
+  }
+  return Status::OK();
+}
+}  // namespace
+
 class ThreadPoolResource : public ResourceBase {
  public:
   ThreadPoolResource(Env* env, const ThreadOptions& thread_options,
                      const string& name, int num_threads, bool low_latency_hint,
                      int max_intra_op_parallelism)
@@ -81,13 +97,11 @@ class ThreadPoolHandleOp : public OpKernel {
   explicit ThreadPoolHandleOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("display_name", &display_name_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("num_threads", &num_threads_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("max_intra_op_parallelism",
                                      &max_intra_op_parallelism_));
-    OP_REQUIRES(
-        ctx, num_threads_ > 0,
-        errors::InvalidArgument("`num_threads` must be greater than zero."));
+    OP_REQUIRES_OK(ctx, ValidateNumThreads(num_threads_));
   }
 
   // The resource is deleted from the resource manager only when it is private
   // to kernel. Ideally the resource should be deleted when it is no longer held
   // by anyone, but it would break backward compatibility.
@@ -529,12 +543,11 @@ class PrivateThreadPoolDatasetOp::Dataset : public DatasetBase {
 /* static */
 void PrivateThreadPoolDatasetOp::MakeDatasetFromOptions(OpKernelContext* ctx,
                                                         DatasetBase* input,
                                                         int32_t num_threads,
                                                         DatasetBase** output) {
-  OP_REQUIRES(ctx, num_threads >= 0,
-              errors::InvalidArgument("`num_threads` must be >= 0"));
+  OP_REQUIRES_OK(ctx, ValidateNumThreads(num_threads));
   *output = new Dataset(ctx,
                         DatasetContext(DatasetContext::Params(
                             {PrivateThreadPoolDatasetOp::kDatasetType,
                              PrivateThreadPoolDatasetOp::kDatasetOp})),
                         input, num_threads);
@@ -544,12 +557,11 @@ void PrivateThreadPoolDatasetOp::MakeDataset(OpKernelContext* ctx,
                                              DatasetBase* input,
                                              DatasetBase** output) {
   int64_t num_threads = 0;
   OP_REQUIRES_OK(
       ctx, ParseScalarArgument<int64_t>(ctx, "num_threads", &num_threads));
-  OP_REQUIRES(ctx, num_threads >= 0,
-              errors::InvalidArgument("`num_threads` must be >= 0"));
+  OP_REQUIRES_OK(ctx, ValidateNumThreads(num_threads));
   *output = new Dataset(ctx, input, num_threads);
 }
 
 namespace {
 

Vulnerable code:     OP_REQUIRES(
        ctx, num_threads_ > 0,
        errors::InvalidArgument("`num_threads` must be greater than zero."));
  OP_REQUIRES(ctx, num_threads >= 0,
              errors::InvalidArgument("`num_threads` must be >= 0"));
  OP_REQUIRES(ctx, num_threads >= 0,
              errors::InvalidArgument("`num_threads` must be >= 0"));
Fixed code: namespace {
// To prevent integer overflow issues when allocating threadpool memory for an
// unreasonable number of threads.
constexpr int kThreadLimit = 65536;

Status ValidateNumThreads(int32_t num_threads) {
  if (num_threads < 0) {
    return errors::InvalidArgument("`num_threads` must be >= 0");
  }
  if (num_threads >= kThreadLimit) {
    return errors::InvalidArgument("`num_threads` must be < ", kThreadLimit);
  }
  return Status::OK();
}
}  // namespace

    OP_REQUIRES_OK(ctx, ValidateNumThreads(num_threads_));
  OP_REQUIRES_OK(ctx, ValidateNumThreads(num_threads));
  OP_REQUIRES_OK(ctx, ValidateNumThreads(num_threads));