43 lines
1.7 KiB
Diff
43 lines
1.7 KiB
Diff
|
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
|
||
|
|
|