From b5a549d748eb69033afeca354bbc87e098402272 Mon Sep 17 00:00:00 2001 From: destro4evr-rgb Date: Sat, 10 Oct 2026 04:21:09 -0400 Subject: [PATCH] security: fix heap OOB read in KmeansPlusPlusInitializationOp When all input points are identical the cumulative distance sum is zero. draw_one_sample() then multiplies RandFloat() * 0.0f = 0.0f and calls std::upper_bound on an all-zeros array, which returns the past-the-end pointer. This makes draw_one_sample() return num_points (one past the last valid row), which is immediately used as a row index into the points and points_half_squared_norm Eigen objects, causing a heap OOB read of point_dimensions * sizeof(float) bytes beyond the tensor allocation. Fix: detect total == 0.0f before calling upper_bound and fall back to uniform sampling (matching the first-iteration path). Add a defensive std::min clamp for floating-point rounding edge cases. --- tensorflow/core/kernels/clustering_ops.cc | 46 +++++++++++++------- tensorflow/python/ops/clustering_ops_test.py | 41 +++++++++++++++++ 2 files changed, 72 insertions(+), 15 deletions(-) diff --git a/tensorflow/core/kernels/clustering_ops.cc b/tensorflow/core/kernels/clustering_ops.cc index 3fe17a43285..f91a57e5045 100644 --- a/tensorflow/core/kernels/clustering_ops.cc +++ b/tensorflow/core/kernels/clustering_ops.cc @@ -108,11 +108,6 @@ class KmeansPlusPlusInitializationOp : public OpKernel { const int64_t point_dimensions = points_tensor.dim_size(1); const int64_t num_to_sample = num_to_sample_tensor.scalar()(); const int64_t seed = seed_tensor.scalar()(); - const int64_t num_retries_per_sample = [&]() { - const int64_t value = num_retries_per_sample_tensor.scalar()(); - return value >= 0 ? value - : 2 + static_cast(std::log(num_to_sample)); - }(); OP_REQUIRES(context, num_points > 0, absl::InvalidArgumentError("Expected points.rows() > 0.")); @@ -123,6 +118,12 @@ class KmeansPlusPlusInitializationOp : public OpKernel { absl::StrCat("Expected num_to_sample <= points.rows(). ", num_to_sample, " vs ", num_points, "."))); + const int64_t num_retries_per_sample = [&]() { + const int64_t value = num_retries_per_sample_tensor.scalar()(); + return value >= 0 ? value + : 2 + static_cast(std::log(num_to_sample)); + }(); + Tensor* output_sampled_points_tensor; OP_REQUIRES_OK(context, context->allocate_output( @@ -156,16 +157,28 @@ class KmeansPlusPlusInitializationOp : public OpKernel { auto draw_one_sample = [&]() -> int64_t { if (sampled_indices.empty()) return rng.Uniform64(num_points); + const float total = min_distances_cumsum(num_points - 1); int64_t index = 0; do { - // If v is drawn from Uniform[0, distances.sum()), then - // Prob[cumsum(distances)(i - 1) <= v < cumsum(distances)(i)] is - // proportional to distances(i). - index = std::upper_bound( - min_distances_cumsum.data(), - min_distances_cumsum.data() + num_points, - rng.RandFloat() * min_distances_cumsum(num_points - 1)) - - min_distances_cumsum.data(); + if (!(total > 0.0f) || std::isinf(total)) { + // total is zero, negative, NaN, or infinity (e.g. all points + // identical, NaN distances, or overflow). Fall back to uniform + // sampling to avoid upper_bound returning end-of-range and an + // out-of-bounds index. + index = rng.Uniform64(num_points); + } else { + // If v is drawn from Uniform[0, distances.sum()), then + // Prob[cumsum(distances)(i - 1) <= v < cumsum(distances)(i)] is + // proportional to distances(i). + index = std::upper_bound( + min_distances_cumsum.data(), + min_distances_cumsum.data() + num_points, + rng.RandFloat() * total) - + min_distances_cumsum.data(); + // upper_bound guarantees index < num_points when total > 0, but + // clamp defensively against floating-point rounding edge cases. + index = std::min(index, num_points - 1); + } } while (sampled_indices.find(index) != sampled_indices.end()); return index; }; @@ -175,6 +188,7 @@ class KmeansPlusPlusInitializationOp : public OpKernel { min_distances = min_distances.cwiseMin(GetHalfSquaredDistancesToY( points, points_half_squared_norm, points.row(sampled_index), points_half_squared_norm(sampled_index))); + min_distances(sampled_index) = 0.0f; return sampled_index; }; @@ -182,14 +196,15 @@ class KmeansPlusPlusInitializationOp : public OpKernel { Eigen::VectorXf best_new_min_distances(num_points); float best_potential = std::numeric_limits::infinity(); int64_t best_sampled_index = 0; - for (int i = 1 + num_retries_per_sample; i > 0; --i) { + for (int64_t i = 0; i <= num_retries_per_sample; ++i) { const int64_t sampled_index = draw_one_sample(); Eigen::VectorXf new_min_distances = min_distances.cwiseMin(GetHalfSquaredDistancesToY( points, points_half_squared_norm, points.row(sampled_index), points_half_squared_norm(sampled_index))); + new_min_distances(sampled_index) = 0.0f; const float potential = new_min_distances.sum(); - if (potential < best_potential) { + if (i == 0 || potential < best_potential) { best_potential = potential; best_sampled_index = sampled_index; best_new_min_distances.swap(new_min_distances); @@ -201,6 +216,7 @@ class KmeansPlusPlusInitializationOp : public OpKernel { for (int64_t i = 0; i < num_to_sample; ++i) { if (i > 0) { + min_distances = min_distances.cwiseMax(0.0f); std::partial_sum(min_distances.data(), min_distances.data() + num_points, min_distances_cumsum.data()); diff --git a/tensorflow/python/ops/clustering_ops_test.py b/tensorflow/python/ops/clustering_ops_test.py index ab513910644..9f2b92543eb 100644 --- a/tensorflow/python/ops/clustering_ops_test.py +++ b/tensorflow/python/ops/clustering_ops_test.py @@ -52,6 +52,47 @@ class KmeansPlusPlusInitializationTest(test.TestCase): for seed in range(100): self.runTestWithSeed(seed) + @test_util.run_in_graph_and_eager_modes + def testIdenticalPoints(self): + points = np.array([[1.0, 2.0], [1.0, 2.0]], dtype=np.float32) + with self.cached_session(): + for retries in (0, 3): + sampled = clustering_ops.kmeans_plus_plus_initialization( + points, num_to_sample=2, seed=12345, num_retries_per_sample=retries + ) + result = self.evaluate(sampled) + self.assertEqual(result.shape, (2, 2)) + self.assertAllClose(result, points) + + @test_util.run_in_graph_and_eager_modes + def testNaNPoints(self): + points = np.full((2, 2), float("nan"), dtype=np.float32) + with self.cached_session(): + for retries in (0, 3): + sampled = clustering_ops.kmeans_plus_plus_initialization( + points, num_to_sample=2, seed=12345, num_retries_per_sample=retries + ) + result = self.evaluate(sampled) + self.assertEqual(result.shape, (2, 2)) + self.assertEqual(result.dtype, np.float32) + self.assertTrue(np.isnan(result).all()) + + @test_util.run_in_graph_and_eager_modes + def testOverflowPoints(self): + # Values large enough that squared distance overflows float32 to +inf + # without individual norms overflowing to +inf (which would yield NaN), + # exercising the std::isinf(total) fallback in draw_one_sample. + points = np.array([[1e19, 1e19], [-1e19, -1e19]], dtype=np.float32) + with self.cached_session(): + for retries in (0, 3): + sampled = clustering_ops.kmeans_plus_plus_initialization( + points, num_to_sample=2, seed=12345, num_retries_per_sample=retries + ) + result = self.evaluate(sampled) + self.assertEqual(result.shape, (2, 2)) + self.assertEqual(result.dtype, np.float32) + self.assertCountEqual(result.tolist(), points.tolist()) + @test_util.run_all_in_graph_and_eager_modes class KMC2InitializationTest(test.TestCase):