From ff70c47a396ef1e3cb73c90513da4f5cb71bebba Mon Sep 17 00:00:00 2001 From: Amit Patankar 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()(); 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()(); - 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