1 files changed, 27 insertions, 18 deletions

diff --git a/lib/jxl/toc.cc b/lib/jxl/toc.cc
index 3a2193e..24cdd02 100644
--- a/lib/jxl/toc.cc
+++ b/lib/jxl/toc.cc
@@ -20,10 +20,9 @@ size_t MaxBits(const size_t num_sizes) {
   return 1 + kBitsPerByte + entry_bits + kBitsPerByte;
 }
 
-Status ReadGroupOffsets(size_t toc_entries, BitReader* JXL_RESTRICT reader,
-                        std::vector<uint64_t>* JXL_RESTRICT offsets,
-                        std::vector<uint32_t>* JXL_RESTRICT sizes,
-                        uint64_t* total_size) {
+Status ReadToc(size_t toc_entries, BitReader* JXL_RESTRICT reader,
+               std::vector<uint32_t>* JXL_RESTRICT sizes,
+               std::vector<coeff_order_t>* JXL_RESTRICT permutation) {
   if (toc_entries > 65536) {
     // Prevent out of memory if invalid JXL codestream causes a bogus amount
     // of toc_entries such as 2720436919446 to be computed.
@@ -31,6 +30,11 @@ Status ReadGroupOffsets(size_t toc_entries, BitReader* JXL_RESTRICT reader,
     return JXL_FAILURE("too many toc entries");
   }
 
+  sizes->clear();
+  sizes->resize(toc_entries);
+  if (reader->TotalBitsConsumed() >= reader->TotalBytes() * kBitsPerByte) {
+    return JXL_STATUS(StatusCode::kNotEnoughBytes, "Not enough bytes for TOC");
+  }
   const auto check_bit_budget = [&](size_t num_entries) -> Status {
     // U32Coder reads 2 bits to recognize variant and kTocDist cheapest variant
     // is Bits(10), this way at least 12 bits are required per toc-entry.
@@ -44,35 +48,40 @@ Status ReadGroupOffsets(size_t toc_entries, BitReader* JXL_RESTRICT reader,
     return JXL_STATUS(StatusCode::kNotEnoughBytes, "Not enough bytes for TOC");
   };
 
-  JXL_DASSERT(offsets != nullptr && sizes != nullptr);
-  std::vector<coeff_order_t> permutation;
-  if (reader->ReadFixedBits<1>() == 1 && toc_entries > 0) {
-    // Skip permutation description if the toc_entries is 0.
+  JXL_DASSERT(toc_entries > 0);
+  if (reader->ReadFixedBits<1>() == 1) {
     JXL_RETURN_IF_ERROR(check_bit_budget(toc_entries));
-    permutation.resize(toc_entries);
-    JXL_RETURN_IF_ERROR(
-        DecodePermutation(/*skip=*/0, toc_entries, permutation.data(), reader));
+    permutation->resize(toc_entries);
+    JXL_RETURN_IF_ERROR(DecodePermutation(/*skip=*/0, toc_entries,
+                                          permutation->data(), reader));
   }
-
   JXL_RETURN_IF_ERROR(reader->JumpToByteBoundary());
   JXL_RETURN_IF_ERROR(check_bit_budget(toc_entries));
-  sizes->clear();
-  sizes->reserve(toc_entries);
   for (size_t i = 0; i < toc_entries; ++i) {
-    sizes->push_back(U32Coder::Read(kTocDist, reader));
+    (*sizes)[i] = U32Coder::Read(kTocDist, reader);
   }
   JXL_RETURN_IF_ERROR(reader->JumpToByteBoundary());
   JXL_RETURN_IF_ERROR(check_bit_budget(0));
+  return true;
+}
+
+Status ReadGroupOffsets(size_t toc_entries, BitReader* JXL_RESTRICT reader,
+                        std::vector<uint64_t>* JXL_RESTRICT offsets,
+                        std::vector<uint32_t>* JXL_RESTRICT sizes,
+                        uint64_t* total_size) {
+  std::vector<coeff_order_t> permutation;
+  JXL_RETURN_IF_ERROR(ReadToc(toc_entries, reader, sizes, &permutation));
 
-  // Prefix sum starting with 0 and ending with the offset of the last group
   offsets->clear();
-  offsets->reserve(toc_entries);
+  offsets->resize(toc_entries);
+
+  // Prefix sum starting with 0 and ending with the offset of the last group
   uint64_t offset = 0;
   for (size_t i = 0; i < toc_entries; ++i) {
     if (offset + (*sizes)[i] < offset) {
       return JXL_FAILURE("group offset overflow");
     }
-    offsets->push_back(offset);
+    (*offsets)[i] = offset;
     offset += (*sizes)[i];
   }
   if (total_size) {