From a7e4df924652a151677b3e0b95609db27cb04631 Mon Sep 17 00:00:00 2001
From: Daniel Nguyen <nguyendaniel@google.com>
Date: Tue, 11 Aug 2020 21:01:09 +0000
Subject: [PATCH] switched from unsigned char to TF_Bool. Added input checking
 for TF_AllocatorAttributes

---
 tensorflow/c/BUILD           |  1 +
 tensorflow/c/c_api_macros.h  |  6 ++++++
 tensorflow/c/kernels.cc      | 14 ++++++++++----
 tensorflow/c/kernels.h       |  2 +-
 tensorflow/c/kernels_test.cc |  7 +++----
 tensorflow/c/tf_tensor.h     |  2 +-
 6 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/tensorflow/c/BUILD b/tensorflow/c/BUILD
index c4c0420aaec..7b78fdb4cf6 100644
--- a/tensorflow/c/BUILD
+++ b/tensorflow/c/BUILD
@@ -81,6 +81,7 @@ tf_cuda_library(
     hdrs = [
         "c_api.h",
         "c_api_internal.h",
+        "c_api_macros.h",
         "tf_datatype.h",
         "tf_tensor.h",
         "tf_tstring.h",
diff --git a/tensorflow/c/c_api_macros.h b/tensorflow/c/c_api_macros.h
index ce24e4d8cbd..58547a57b6e 100644
--- a/tensorflow/c/c_api_macros.h
+++ b/tensorflow/c/c_api_macros.h
@@ -30,6 +30,12 @@ limitations under the License.
 #endif  // _WIN32
 #endif  // SWIG
 
+// TF_Bool is the C API typedef for unsigned char, while TF_BOOL is 
+// the datatype for boolean tensors. 
+#ifndef TF_BOOL_DEFINED 
+#define TF_Bool unsigned char 
+#endif // TF_BOOL_DEFINED
+
 // Macro used to calculate struct size for maintaining ABI stability across 
 // different struct implementations. 
 #ifndef TF_OFFSET_OF_END
diff --git a/tensorflow/c/kernels.cc b/tensorflow/c/kernels.cc
index 6d9d7d57517..cce6f2c494d 100644
--- a/tensorflow/c/kernels.cc
+++ b/tensorflow/c/kernels.cc
@@ -261,7 +261,6 @@ TF_Tensor* TF_AllocateOutput(TF_OpKernelContext* context, int index,
                              size_t len, TF_Status* status) {
   TF_SetStatus(status, TF_OK, "");
   auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(context);
-
   static_assert(sizeof(int64_t) == sizeof(tensorflow::int64),
                 "64-bit int types should match in size");
   tensorflow::gtl::ArraySlice<tensorflow::int64> dimarray(
@@ -286,22 +285,29 @@ TF_Tensor* TF_AllocateTemp(TF_OpKernelContext* context, TF_DataType dtype,
                            TF_AllocatorAttributes* attributes,
                            TF_Status* status) {
   auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(context);
-  TF_SetStatus(status, TF_OK, ""); 
+  TF_SetStatus(status, TF_OK, "");
+  static_assert(sizeof(int64_t) == sizeof(tensorflow::int64),
+                "64-bit int types should match in size");
   tensorflow::gtl::ArraySlice<tensorflow::int64> dimarray(
       reinterpret_cast<tensorflow::int64*>(dims), num_dims);
+  if (attributes && !attributes->struct_size) {
+    TF_SetStatus(status, TF_INVALID_ARGUMENT, "TF_AllocatorAttributes struct "
+        "size member must be set to TF_ALLOCATOR_ATTRIBUTES_STRUCT_SIZE"); 
+    return nullptr;  
+  }
   tensorflow::AllocatorAttributes allocator_attr; 
-  if (attributes->on_host) { 
+  if (attributes && attributes->on_host) { 
     allocator_attr.set_on_host(true); 
   }
   tensorflow::Status s;
   tensorflow::Tensor tensor;
-  TF_Tensor* tf_tensor;
   s = cc_ctx->allocate_temp(static_cast<tensorflow::DataType>(dtype), 
       tensorflow::TensorShape(dimarray), &tensor, allocator_attr);
   if (!s.ok()) {
     ::tensorflow::Set_TF_Status_from_Status(status, s);
     return nullptr;
   }
+  TF_Tensor* tf_tensor;
   tf_tensor = TF_TensorFromTensor(tensor, &s);
   if (!s.ok()) {
     ::tensorflow::Set_TF_Status_from_Status(status, s);
diff --git a/tensorflow/c/kernels.h b/tensorflow/c/kernels.h
index ee865613b6e..0be31b76839 100644
--- a/tensorflow/c/kernels.h
+++ b/tensorflow/c/kernels.h
@@ -203,7 +203,7 @@ TF_CAPI_EXPORT TF_Tensor* TF_AllocateOutput(TF_OpKernelContext* context,
 // Allocates a temporary Tensor of the specified type and shape. The
 // Tensor must not be used after kernel construction is
 // complete. 
-
+//
 // num_dims must equal the size of array dims 
 TF_CAPI_EXPORT extern TF_Tensor* TF_AllocateTemp(TF_OpKernelContext* context,
     TF_DataType dtype, int64_t* dims, int num_dims, TF_AllocatorAttributes* 
diff --git a/tensorflow/c/kernels_test.cc b/tensorflow/c/kernels_test.cc
index e216b85c13f..7663a3570eb 100644
--- a/tensorflow/c/kernels_test.cc
+++ b/tensorflow/c/kernels_test.cc
@@ -368,13 +368,12 @@ class DeviceKernelOpTest : public OpsTestBase {
 #endif
 };
 
-// Helper function for tests that validates that the tensor has 
-// shape and type corresponding to dims and dtype. 
+// Validates that the tensor has shape and type corresponding to 
+// dims and dtype. 
 void validate_tensor(TF_Tensor* tensor, int64_t* dims, int64_t num_dims, 
                      TF_DataType dtype);
 
-// Helper function for tests that copies data of length 
-// tensor_size_bytes from values to tensor 
+// Copies data of length tensor_size_bytes from values to tensor.
 template <typename T> 
 void set_tensor_data(TF_Tensor* tensor, T* values, size_t tensor_size_bytes, 
                      TF_OpKernelContext* ctx);
diff --git a/tensorflow/c/tf_tensor.h b/tensorflow/c/tf_tensor.h
index ced57df77d4..1faa368f25a 100644
--- a/tensorflow/c/tf_tensor.h
+++ b/tensorflow/c/tf_tensor.h
@@ -50,7 +50,7 @@ extern "C" {
 typedef struct TF_AllocatorAttributes { 
   size_t struct_size; 
   // Set boolean to 1 for CPU allocation, else 0. 
-  unsigned char on_host; 
+  TF_Bool on_host; 
 } TF_AllocatorAttributes;