Out of bounds segmentation fault due to unequal op inputs in Tensorflow

Impact

tf.raw_ops.DynamicStitch specifies input sizes when it is registered.

REGISTER_OP("DynamicStitch")
    .Input("indices: N * int32")
    .Input("data: N * T")
    .Output("merged: T")
    .Attr("N : int >= 1")
    .Attr("T : type")
    .SetShapeFn(DynamicStitchShapeFunction);

When it receives a differing number of inputs, such as when it is called with an indices size 1 and a data size 2, it will crash.

import tensorflow as tf

# indices = 1*[tf.random.uniform([1,2], dtype=tf.dtypes.int32, maxval=100)]
indices = [tf.constant([[0, 1]]),]

# data = 2*[tf.random.uniform([1,2], dtype=tf.dtypes.float32, maxval=100)]
data = [tf.constant([[5, 6]]), tf.constant([[7, 8]])]

tf.raw_ops.DynamicStitch(
    indices=indices, 
    data=data)

Patches

We have patched the issue in GitHub commit f5381e0e10b5a61344109c1b7c174c68110f7629.

The fix will be included in TensorFlow 2.11. We will also cherrypick this commit on TensorFlow 2.10.1 as this is also affected.

For more information

Please consult our security guide for more information regarding the security model and how to contact us with issues and questions.

Attribution

This vulnerability has been reported by Zizhuang Deng of IIE, UCAS