From 1cf45b831eeb0cab8655c9c7c5d06ec6f45fc41b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 13 May 2022 20:14:36 -0700
Subject: [PATCH] [tfg][functiondef_import] Emit error on empty function
 attributes

Disallow empty function attributes. Emit an error when one is encountered.

PiperOrigin-RevId: 448629968
---
 .../ir/importexport/functiondef_import.cc     |  2 +
 .../invalid_generic_function_attr_name.pbtxt  | 52 +++++++++++++++++++
 2 files changed, 54 insertions(+)
 create mode 100644 tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_function_attr_name.pbtxt

diff --git a/tensorflow/core/ir/importexport/functiondef_import.cc b/tensorflow/core/ir/importexport/functiondef_import.cc
index c2f7bfb1fe8034..5118550c7607fc 100644
--- a/tensorflow/core/ir/importexport/functiondef_import.cc
+++ b/tensorflow/core/ir/importexport/functiondef_import.cc
@@ -333,6 +333,8 @@ Status ImportGenericFunction(
   // Import the function attributes with a `tf.` prefix to match the current
   // infrastructure expectations.
   for (const auto& namedAttr : func.attr()) {
+    if (namedAttr.first.empty())
+      return InvalidArgument("Invalid function attribute name");
     const std::string& name = "tf." + namedAttr.first;
     const AttrValue& tf_attr = namedAttr.second;
     TF_ASSIGN_OR_RETURN(Attribute attr,
diff --git a/tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_function_attr_name.pbtxt b/tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_function_attr_name.pbtxt
new file mode 100644
index 00000000000000..7a0f18f6732027
--- /dev/null
+++ b/tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_function_attr_name.pbtxt
@@ -0,0 +1,52 @@
+# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
+
+# CHECK: Invalid function attribute name
+
+library {
+  function {
+    signature {
+      name: "foo"
+      input_arg {
+        name: "a"
+      }
+      output_arg {
+        name: "d"
+      }
+    }
+    node_def {
+      op: "Const"
+      attr {
+        key: "_b"
+        value {
+          placeholder: "T"
+        }
+      }
+      attr {
+        key: "dtype"
+        value {
+          type: DT_INT32
+        }
+      }
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_INT32
+            tensor_shape {
+            }
+          }
+        }
+      }
+    }
+    ret {
+      key: "d"
+      value: "a"
+    }
+    attr {
+      key: ""
+      value {
+        s: "a"
+      }
+    }
+  }
+}