From fa5bf7b37103fc9b619e31414a352ce858686bfe Mon Sep 17 00:00:00 2001
From: Subhan1794 <subhanashfaq66@gmail.com>
Date: Thu, 8 Oct 2026 17:46:17 +0500
Subject: [PATCH] Fix heap-buffer-overflow in NuqCodec::DecompressAndZeroPad on
 num == 0

When num == 0 with an aligned packed_ofs, num_groups computes to 0 and
num_groups - 1 wraps to SIZE_MAX, causing DecWholeGroup to write 256
floats out of bounds (ASan heap-buffer-overflow WRITE).

Return early on num == 0, mirroring the existing guard in
IntCodec::DecompressAndZeroPad (compression/int-inl.h).

Add TestZeroNum regression test: calls DecompressAndZeroPad with num=0
over a sentinel-filled buffer and verifies it is untouched. The test
crashes under ASan without the fix and passes with it.
---
 compression/nuq-inl.h   |  2 ++
 compression/nuq_test.cc | 27 +++++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/compression/nuq-inl.h b/compression/nuq-inl.h
index ce387f98..a4ad1561 100644
--- a/compression/nuq-inl.h
+++ b/compression/nuq-inl.h
@@ -756,6 +756,8 @@ class NuqCodec {
   static HWY_INLINE void DecompressAndZeroPad(
       D d, const PackedSpan<const NuqStream>& packed, size_t packed_ofs,
       Raw* HWY_RESTRICT raw, size_t num) {
+    // num == 0 would make `num_groups - 1` below wrap to SIZE_MAX.
+    if (num == 0) return;
     // If unaligned, load elements from the first group and update the args,
     // from which we compute new tables/indices below.
     size_t current_offset = packed_ofs;
diff --git a/compression/nuq_test.cc b/compression/nuq_test.cc
index 7ddee9ee..f42b456d 100644
--- a/compression/nuq_test.cc
+++ b/compression/nuq_test.cc
@@ -507,6 +507,31 @@ struct TestEncDec {
 void TestEncDecBF16() { hn::ForGEVectors<128, TestEncDec>()(BF16()); }
 void TestEncDecF32() { hn::ForGEVectors<128, TestEncDec>()(float()); }
 
+// Regression test: num == 0 must return immediately without any
+// out-of-bounds access (previously `num_groups - 1` wrapped to SIZE_MAX).
+struct TestZeroNum {
+  template <typename T, class D>
+  HWY_INLINE void operator()(T /*unused*/, D d) {
+    const hn::Repartition<float, D> df;
+    auto out = hwy::AllocateAligned<T>(kGroupSize);
+    auto nuq = hwy::AllocateAligned<NuqStream>(
+        hwy::RoundUpTo(NuqStream::PackedEnd(kGroupSize), hwy::VectorBytes()));
+    HWY_ASSERT(out && nuq);
+    // Sentinel-fill the output; the call must not touch it.
+    for (size_t i = 0; i < kGroupSize; ++i) {
+      out[i] = hwy::ConvertScalarTo<T>(-1.0f);
+    }
+    const auto nuq_span = MakeSpan(nuq.get(), kGroupSize);
+    NuqCodec::DecompressAndZeroPad(d, MakeConst(nuq_span), 0, out.get(), 0);
+    for (size_t i = 0; i < kGroupSize; ++i) {
+      HWY_ASSERT_EQ(hwy::ConvertScalarTo<T>(-1.0f), out[i]);
+    }
+  }
+};
+
+void TestZeroNumBF16() { hn::ForGEVectors<128, TestZeroNum>()(BF16()); }
+void TestZeroNumF32() { hn::ForGEVectors<128, TestZeroNum>()(float()); }
+
 // NOLINTNEXTLINE(google-readability-namespace-comments)
 }  // namespace HWY_NAMESPACE
 }  // namespace gcpp
@@ -529,6 +554,8 @@ HWY_EXPORT_AND_TEST_P(NuqTest, TestUnalignedOffsetF32);
 HWY_EXPORT_AND_TEST_P(NuqTest, TestAllNibble);
 HWY_EXPORT_AND_TEST_P(NuqTest, TestEncDecBF16);
 HWY_EXPORT_AND_TEST_P(NuqTest, TestEncDecF32);
+HWY_EXPORT_AND_TEST_P(NuqTest, TestZeroNumBF16);
+HWY_EXPORT_AND_TEST_P(NuqTest, TestZeroNumF32);
 #else
 GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(NuqTest);
 #endif  // GEMMA_ENABLE_NUQ