From ff70c47a396ef1e3cb73c90513da4f5cb71bebba Mon Sep 17 00:00:00 2001
From: Amit Patankar <amitpatankar@google.com>
Date: Tue, 13 Apr 2021 14:24:00 -0700
Subject: [PATCH] Fix `tf.raw_ops.GetSessionTensor` and
 `tf.raw_ops.DeleteSessionTensor` null pointer dereferences.

---
 tensorflow/core/kernels/session_ops.cc | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/session_ops.cc b/tensorflow/core/kernels/session_ops.cc
index e7e73549..dab59e70 100644
--- a/tensorflow/core/kernels/session_ops.cc
+++ b/tensorflow/core/kernels/session_ops.cc
@@ -119,6 +119,11 @@ class GetSessionTensorOp : public OpKernel {
     const string& name = handle.scalar<tstring>()();
     Tensor val;
     OP_REQUIRES_OK(ctx, ctx->session_state()->GetTensor(name, &val));
+    auto session_state = ctx->session_state();
+    OP_REQUIRES(ctx, session_state != nullptr,
+		errors::FailedPrecondition(
+		    "GetSessionTensor called on null session state"));
+    OP_REQUIRES_OK(ctx, session_state->GetTensor(name, &val));
     ctx->set_output(0, val);
   }
 
@@ -160,7 +165,11 @@ class DeleteSessionTensorOp : public OpKernel {
   void Compute(OpKernelContext* ctx) override {
     const Tensor& handle = ctx->input(0);
     const string& name = handle.scalar<tstring>()();
-    OP_REQUIRES_OK(ctx, ctx->session_state()->DeleteTensor(name));
+    auto session_state = ctx->session_state();
+    OP_REQUIRES(ctx, session_state != nullptr,
+		errors::FailedPrecondition(
+	            "DeleteSessionTensor called on null session state"));
+    OP_REQUIRES_OK(ctx, session_state->DeleteTensor(name));
   }
 
   TF_DISALLOW_COPY_AND_ASSIGN(DeleteSessionTensorOp);
-- 
2.23.0