From cb280d556bc4569839803e5eed2e8b935a636cc2 Mon Sep 17 00:00:00 2001 From: destro4evr-rgb Date: Fri, 16 Oct 2026 08:58:01 -0400 Subject: [PATCH] kernels: validate labels_indices batch indices in CTCLossOpGPU to prevent heap OOB write The GPU path of CTCLossV2 (CTCLossOpGPU) calls DoHistogram which reads each labels_indices[i,0] (a user-controlled int64_t batch index) and increments labels_lengths[key] without any bounds check. labels_lengths is a std::vector of size batch_size; if key >= batch_size or key < 0, std::vector::operator[] writes beyond the allocated heap storage. The CPU path (CTCLossOp) prevents this via SparseTensor::IndicesValid() and a FastBoundsCheck per batch group. The GPU path had neither guard. Add two checks to CTCLossOpGPU::Compute: 1. labels_indices->dim_size(1) > 1 (matching the CPU-path guard at line 106) so that DoHistogram's hard-coded stride of 2 is always valid. 2. A pre-flight loop using FastBoundsCheck(batch_idx, batch_size) on every row of labels_indices before DoHistogram is called, rejecting any out-of-range batch index with InvalidArgumentError. Add testGPULabelsIndicesOOBBatchIndexRaisesError to CTCLossTestV2 covering the batch-index-out-of-range path on GPU (skipped when no GPU is present). --- tensorflow/core/kernels/ctc_loss_op.cc | 32 +++++------ .../kernel_tests/nn_ops/ctc_loss_op_test.py | 55 +++++++++++++++++++ 2 files changed, 71 insertions(+), 16 deletions(-) diff --git a/tensorflow/core/kernels/ctc_loss_op.cc b/tensorflow/core/kernels/ctc_loss_op.cc index b7b0967427e..ffccb39ef8e 100644 --- a/tensorflow/core/kernels/ctc_loss_op.cc +++ b/tensorflow/core/kernels/ctc_loss_op.cc @@ -54,17 +54,6 @@ using se::StreamExecutor; using se::dnn::RnnStateTensorDescriptor; using se::dnn::ToDataType; -template -void DoHistogram(OpKernelContext* ctx, const Tensor* labels_indices, - int num_indices, int batch_size, - std::vector* labels_lengths) { - const T* h_in = labels_indices->flat().data(); - for (int i = 0; i < num_indices; i++) { - const T& key = h_in[i * 2]; - (*labels_lengths)[key]++; - } -} - } // end namespace #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM @@ -283,6 +272,10 @@ class CTCLossOpGPU : public OpKernel { absl::InvalidArgumentError("labels_indices is not a matrix")); OP_REQUIRES(ctx, TensorShapeUtils::IsVector(labels_values->shape()), absl::InvalidArgumentError("labels_values is not a vector")); + OP_REQUIRES(ctx, labels_indices->dim_size(1) == 2, + absl::InvalidArgumentError(absl::StrCat( + "labels_indices second dimension must be 2. Received ", + labels_indices->dim_size(1)))); const TensorShape& inputs_shape = inputs->shape(); const int64_t max_time_raw = inputs_shape.dim_size(0); @@ -318,6 +311,18 @@ class CTCLossOpGPU : public OpKernel { OP_REQUIRES(ctx, batch_size != 0, absl::InvalidArgumentError("batch_size must not be 0")); + // Validate batch indices and convert labels_indices to labels_lengths. + std::vector labels_lengths(batch_size, 0); + const auto labels_mat = labels_indices->matrix(); + for (int64_t i = 0; i < num_indices; ++i) { + const int64_t batch_idx = labels_mat(i, 0); + OP_REQUIRES(ctx, FastBoundsCheck(batch_idx, batch_size), + absl::InvalidArgumentError(absl::StrCat( + "labels_indices batch index ", batch_idx, + " out of range [0, ", batch_size, ")"))); + ++labels_lengths[batch_idx]; + } + Tensor* loss = nullptr; OP_REQUIRES_OK(ctx, ctx->allocate_output("loss", seq_len->shape(), &loss)); @@ -325,11 +330,6 @@ class CTCLossOpGPU : public OpKernel { OP_REQUIRES_OK(ctx, ctx->allocate_output("gradient", inputs_shape, &gradient)); - // Convert the labels_indices to labels_lengths. - std::vector labels_lengths(batch_size, 0); - DoHistogram(ctx, labels_indices, num_indices, batch_size, - &labels_lengths); - StreamExecutor* executor = ctx->op_device_context()->stream()->parent(); se::dnn::DataType data_type = ToDataType::value; diff --git a/tensorflow/python/kernel_tests/nn_ops/ctc_loss_op_test.py b/tensorflow/python/kernel_tests/nn_ops/ctc_loss_op_test.py index 2dd2e9a811c..7e9ec3c1ad8 100644 --- a/tensorflow/python/kernel_tests/nn_ops/ctc_loss_op_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/ctc_loss_op_test.py @@ -1236,6 +1236,61 @@ class CTCLossTestV3(test.TestCase, parameterized.TestCase): self.assertAllClose(loss, ref_loss, atol=1e-6) self.assertAllClose(grad, ref_grad, atol=2e-6) + @test_util.run_v2_only + def testGPULabelsIndicesOOBBatchIndexRaisesError(self): + """GPU CTCLossV2 rejects out-of-range labels_indices batch indices.""" + if not test.is_gpu_available(): + self.skipTest("Need GPU for testing.") + batch_size = 2 + num_frames = 4 + num_classes = 3 + labels_values = constant_op.constant([1, 1], dtype=dtypes.int32) + logits = constant_op.constant( + np.zeros([num_frames, batch_size, num_classes], dtype=np.float32)) + sequence_length = constant_op.constant( + [num_frames] * batch_size, dtype=dtypes.int32) + for bad_batch in (99, -1): + labels_indices = constant_op.constant( + [[bad_batch, 0], [0, 1]], dtype=dtypes.int64) + with ops.device("/GPU:0"): + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, + rf"labels_indices batch index {bad_batch} out of range"): + self.evaluate( + gen_ctc_ops.ctc_loss_v2( + logits, + labels_indices, + labels_values, + sequence_length, + preprocess_collapse_repeated=False, + ctc_merge_repeated=True, + ignore_longer_outputs_than_inputs=False)) + + @test_util.run_v2_only + def testGPULabelsIndicesWrongColumnCountRaisesError(self): + """GPU CTCLossV2 requires labels_indices to have 2 columns.""" + if not test.is_gpu_available(): + self.skipTest("Need GPU for testing.") + logits = constant_op.constant(np.zeros([4, 2, 3], dtype=np.float32)) + labels_values = constant_op.constant([1, 1], dtype=dtypes.int32) + sequence_length = constant_op.constant([4, 4], dtype=dtypes.int32) + for bad_indices in ([[0], [1]], [[0, 0, 99], [1, 0, 0]]): + labels_indices = constant_op.constant( + bad_indices, dtype=dtypes.int64) + with ops.device("/GPU:0"): + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, + r"labels_indices second dimension must be 2"): + self.evaluate( + gen_ctc_ops.ctc_loss_v2( + logits, + labels_indices, + labels_values, + sequence_length, + preprocess_collapse_repeated=False, + ctc_merge_repeated=True, + ignore_longer_outputs_than_inputs=False)) + @test_util.run_all_in_graph_and_eager_modes class CTCLossDeterministicTest(test.TestCase, parameterized.TestCase):