Validate that DataType is valid when parsing from proto using

proto's DataType_IsValid check in AttrValueUtil, with tests.
Change: 142318102

1 files changed, 10 insertions, 1 deletions

diff --git a/tensorflow/core/framework/attr_value_util.cc b/tensorflow/core/framework/attr_value_util.cc
index 36c42db800..486ebb7b9b 100644
--- a/tensorflow/core/framework/attr_value_util.cc
+++ b/tensorflow/core/framework/attr_value_util.cc
@@ -192,8 +192,13 @@ Status AttrValueHasType(const AttrValue& attr_value, StringPiece type) {
         "AttrValue missing value with expected type '", type, "'");
   }
 
-  // Ref types and DT_INVALID are illegal.
+  // Ref types and DT_INVALID are illegal, and DataTypes must
+  // be a valid enum type.
   if (type == "type") {
+    if (!DataType_IsValid(attr_value.type())) {
+      return errors::InvalidArgument("AttrValue has invalid DataType enum: ",
+                                     attr_value.type());
+    }
     if (IsRefType(attr_value.type())) {
       return errors::InvalidArgument(
           "AttrValue must not have reference type value of ",
@@ -205,6 +210,10 @@ Status AttrValueHasType(const AttrValue& attr_value, StringPiece type) {
   } else if (type == "list(type)") {
     for (auto as_int : attr_value.list().type()) {
       const DataType dtype = static_cast<DataType>(as_int);
+      if (!DataType_IsValid(dtype)) {
+        return errors::InvalidArgument("AttrValue has invalid DataType enum: ",
+                                       as_int);
+      }
       if (IsRefType(dtype)) {
         return errors::InvalidArgument(
             "AttrValue must not have reference type value of ",