Vulnerability: Memory exhaustion in Tensorflow ### Impact 
The [implementation of `ThreadPoolHandle`](https://github.com/tensorflow/tensorflow/blob/5100e359aef5c8021f2e71c7b986420b85ce7b3d/tensorflow/core/kernels/data/experimental/threadpool_dataset_op.cc#L79-L135) can be used to trigger a denial of service attack by allocating too much memory:

```python
import tensorflow as tf
y = tf.raw_ops.ThreadPoolHandle(num_threads=0x60000000,display_name='tf')
```

This is because the `num_threads` argument is only checked to not be negative, but there is no upper bound on its value.
    
### Patches
We have patched the issue in GitHub commit [e3749a6d5d1e8d11806d4a2e9cc3123d1a90b75e](https://github.com/tensorflow/tensorflow/commit/e3749a6d5d1e8d11806d4a2e9cc3123d1a90b75e).

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.

### For more information
Please consult [our security guide](https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) for more information regarding the security model and how to contact us with issues and questions.

### Attribution
This vulnerability has been reported by Yu Tian of Qihoo 360 AIVul Team.
Vulnerability ID: GHSA-c582-c96p-r5cq
CWE: CWE-400
Severity score: 4.3
Project: https://github.com/tensorflow/tensorflow
Published: 2022-02-10
File: tensorflow/core/kernels/data/experimental/threadpool_dataset_op.cc
Commit message: [tf.data] Set limit on number of threads used in threadpool_dataset.

PiperOrigin-RevId: 410922677
Change-Id: Ib25814a99043ab10805b5d2d7088ae0e0b7b04fd
Commit author: Andrew Audibert

Vulnerable version (before fix):
 namespace experimental {     PrivateThreadPoolDatasetOp::kDatasetType; /* static */ constexpr const char* const PrivateThreadPoolDatasetOp::kDatasetOp;  class ThreadPoolResource : public ResourceBase {  public:   ThreadPoolResource(Env* env, const ThreadOptions& thread_options, class ThreadPoolHandleOp : public OpKernel {     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."));   }    // The resource is deleted from the resource manager only when it is private 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"));   *output = new Dataset(ctx,                         DatasetContext(DatasetContext::Params(                             {PrivateThreadPoolDatasetOp::kDatasetType, void PrivateThreadPoolDatasetOp::MakeDataset(OpKernelContext* ctx,   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"));   *output = new Dataset(ctx, input, num_threads); } 

Fixed version (after fix):
 namespace experimental {     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, class ThreadPoolHandleOp : public OpKernel {     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_OK(ctx, ValidateNumThreads(num_threads_));   }    // The resource is deleted from the resource manager only when it is private void PrivateThreadPoolDatasetOp::MakeDatasetFromOptions(OpKernelContext* ctx,                                                         DatasetBase* input,                                                         int32_t num_threads,                                                         DatasetBase** output) {   OP_REQUIRES_OK(ctx, ValidateNumThreads(num_threads));   *output = new Dataset(ctx,                         DatasetContext(DatasetContext::Params(                             {PrivateThreadPoolDatasetOp::kDatasetType, void PrivateThreadPoolDatasetOp::MakeDataset(OpKernelContext* ctx,   int64_t num_threads = 0;   OP_REQUIRES_OK(       ctx, ParseScalarArgument<int64_t>(ctx, "num_threads", &num_threads));   OP_REQUIRES_OK(ctx, ValidateNumThreads(num_threads));   *output = new Dataset(ctx, input, num_threads); } 

Patch:
patch: @@ -39,6 +39,22 @@ namespace experimental {
     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,
@@ -83,9 +99,7 @@ class ThreadPoolHandleOp : public OpKernel {
     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
@@ -531,8 +545,7 @@ 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,
@@ -546,8 +559,7 @@ void PrivateThreadPoolDatasetOp::MakeDataset(OpKernelContext* ctx,
   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);
 }
path: tensorflow/core/kernels/data/experimental/threadpool_dataset_op.cc