OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-06 12:20:11 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Merge pull request #49333 from geetachavan1/cherrypicks_9XDBG

Browse Source 
Fix `tf.raw_ops.GetSessionTensor` and `tf.raw_ops.DeleteSessionTensor` null pointer dereferences.
This commit is contained in:
Mihai Maruseac 2021-05-20 08:25:53 -07:00 committed by GitHub
commit ad8dfb5d10
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tensorflow/core/kernels/session_ops.cc
View File

@ -91,7 +91,6 @@ TF_CALL_NUMBER_TYPES(REGISTER_GPU_KERNEL);
REGISTER_GPU_KERNEL(bool);
#undef REGISTER_GPU_KERNEL
class GetSessionTensorOp : public OpKernel {
public:
explicit GetSessionTensorOp(OpKernelConstruction* context)
@ -101,7 +100,11 @@ class GetSessionTensorOp : public OpKernel {
const Tensor& handle = ctx->input(0);
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);
}
@ -122,7 +125,6 @@ TF_CALL_NUMBER_TYPES(REGISTER_GPU_KERNEL);
REGISTER_GPU_KERNEL(bool);
#undef REGISTER_GPU_KERNEL
class DeleteSessionTensorOp : public OpKernel {
public:
explicit DeleteSessionTensorOp(OpKernelConstruction* context)
@ -131,7 +133,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);