From c5f8695d1e77bd36abd66f56ce10eac24d6517ee Mon Sep 17 00:00:00 2001
From: Sagun Bajra <sagunb@google.com>
Date: Tue, 27 Sep 2022 14:50:36 -0700
Subject: [PATCH] Fix security vulnerability with calling ParseInt64Value with
 an invalid py_value.

PiperOrigin-RevId: 477281263
---
 tensorflow/python/eager/pywrap_tfe_src.cc | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensorflow/python/eager/pywrap_tfe_src.cc b/tensorflow/python/eager/pywrap_tfe_src.cc
index 74bc7cd7b7644a..be7205bb63e01e 100644
--- a/tensorflow/python/eager/pywrap_tfe_src.cc
+++ b/tensorflow/python/eager/pywrap_tfe_src.cc
@@ -256,6 +256,13 @@ PARSE_VALUE(ParseFloatValue, float, PyFloat_Check, PyFloat_AsDouble)
 #if PY_MAJOR_VERSION < 3
 bool ParseInt64Value(const string& key, PyObject* py_value, TF_Status* status,
                      int64_t* value) {
+  if (py_value == nullptr) {
+    TF_SetStatus(status, TF_INVALID_ARGUMENT,
+                 tensorflow::strings::StrCat(
+                     "Expecting int or long value for attr ", key, "."))
+        .c_str();
+    return false;
+  }
   if (PyInt_Check(py_value)) {
     *value = static_cast<int64_t>(PyInt_AsLong(py_value));
     return true;