From 8b011e55ac4324a0250b338e95881025cec4a17f Mon Sep 17 00:00:00 2001
From: Tiffany Song <tiffanysong@google.com>
Date: Thu, 30 Jul 2026 10:14:52 -0700
Subject: [PATCH] [mac] UDIF: handle zlib/bzip2 chunks larger than uInt

UDIFBlockChunkReadStream::HandleZLib() and HandleBZ2() assigned size_t
buffer lengths directly to z_stream/bz_stream's 32-bit
avail_in/avail_out fields and called the decompressor once. When a
chunk's decompressed size (sector_count * 512) exceeded UINT_MAX the
assignment truncated and the chunk was rejected, causing the partition
to be skipped during DMG analysis.

Switch both functions to the standard streaming idiom: call
inflate()/BZ2_bzDecompress() in a loop, refilling avail_in / avail_out
via base::saturated_cast<uInt>() from the remaining 64-bit byte counts
whenever they reach zero. For chunks that fit in 32 bits the loop runs
once and behavior is unchanged.

Fixed: 517926950
Change-Id: Id7f4137786dea7635f7ae4525109daee41fa4a3e
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/8167125
Reviewed-by: Yaw Frempong <yawfrempong@google.com>
Commit-Queue: Tiffany Song <tiffanysong@google.com>
Cr-Commit-Position: refs/heads/main@{#1671186}
---

diff --git a/chrome/utility/safe_browsing/mac/udif.cc b/chrome/utility/safe_browsing/mac/udif.cc
index 7babcae2..48245fa7 100644
--- a/chrome/utility/safe_browsing/mac/udif.cc
+++ b/chrome/utility/safe_browsing/mac/udif.cc
@@ -11,6 +11,7 @@
 
 #include <algorithm>
 #include <array>
+#include <limits>
 #include <memory>
 #include <optional>
 #include <utility>
@@ -26,6 +27,7 @@
 #include "base/notimplemented.h"
 #include "base/notreached.h"
 #include "base/numerics/ostream_operators.h"
+#include "base/numerics/safe_conversions.h"
 #include "base/numerics/safe_math.h"
 #include "base/strings/sys_string_conversions.h"
 #include "chrome/utility/safe_browsing/mac/convert_big_endian.h"
@@ -363,6 +365,11 @@
 
 }  // namespace
 
+size_t& GetMaxDecompressChunkSize() {
+  static size_t max_chunk_size = std::numeric_limits<uint32_t>::max();
+  return max_chunk_size;
+}
+
 UDIFParser::UDIFParser(ReadStream* stream)
     : stream_(stream),
       partition_names_(),
@@ -744,12 +751,12 @@
                                                    const UDIFBlockChunk* chunk)
     : stream_(stream),
       chunk_(chunk),
-      length_in_bytes_(chunk->sector_count * block_size),
       offset_(0),
       decompress_buffer_(),
       did_decompress_(false) {
-  // Make sure the multiplication above did not overflow.
-  CHECK(length_in_bytes_ == 0 || length_in_bytes_ >= block_size);
+  base::CheckedNumeric<size_t> length =
+      base::CheckedNumeric<size_t>(chunk->sector_count) * block_size;
+  CHECK(length.AssignIfValid(&length_in_bytes_));
 }
 
 UDIFBlockChunkReadStream::~UDIFBlockChunkReadStream() = default;
@@ -849,12 +856,38 @@
     }
 
     decompress_buffer_.resize(length_in_bytes_);
-    zlib.next_in = compressed_data.data();
-    zlib.avail_in = compressed_data.size();
-    zlib.next_out = decompress_buffer_.data();
-    zlib.avail_out = decompress_buffer_.size();
+    base::span<const uint8_t> remaining_in = compressed_data;
+    base::span<uint8_t> remaining_out = decompress_buffer_;
 
-    int rv = inflate(&zlib, Z_FINISH);
+    int rv = Z_OK;
+    while (rv == Z_OK && (!remaining_in.empty() || !remaining_out.empty())) {
+      zlib.next_in = const_cast<uint8_t*>(remaining_in.data());
+      zlib.avail_in = base::checked_cast<uInt>(
+          std::min(remaining_in.size(), GetMaxDecompressChunkSize()));
+      zlib.next_out = remaining_out.data();
+      zlib.avail_out = base::checked_cast<uInt>(
+          std::min(remaining_out.size(), GetMaxDecompressChunkSize()));
+
+      uInt avail_in_before = zlib.avail_in;
+      uInt avail_out_before = zlib.avail_out;
+
+      int flush = (remaining_in.size() <= GetMaxDecompressChunkSize() &&
+                   remaining_out.size() <= GetMaxDecompressChunkSize())
+                      ? Z_FINISH
+                      : Z_NO_FLUSH;
+      rv = inflate(&zlib, flush);
+
+      size_t bytes_consumed = avail_in_before - zlib.avail_in;
+      size_t bytes_produced = avail_out_before - zlib.avail_out;
+
+      if (bytes_consumed == 0 && bytes_produced == 0 && rv == Z_OK) {
+        break;
+      }
+
+      remaining_in = remaining_in.subspan(bytes_consumed);
+      remaining_out = remaining_out.subspan(bytes_produced);
+    }
+
     inflateEnd(&zlib);
 
     if (rv != Z_STREAM_END) {
@@ -884,12 +917,35 @@
     }
 
     decompress_buffer_.resize(length_in_bytes_);
-    bz.next_in = reinterpret_cast<char*>(compressed_data.data());
-    bz.avail_in = compressed_data.size();
-    bz.next_out = reinterpret_cast<char*>(decompress_buffer_.data());
-    bz.avail_out = decompress_buffer_.size();
+    base::span<const uint8_t> remaining_in = compressed_data;
+    base::span<uint8_t> remaining_out = decompress_buffer_;
 
-    int rv = BZ2_bzDecompress(&bz);
+    int rv = BZ_OK;
+    while (rv == BZ_OK && (!remaining_in.empty() || !remaining_out.empty())) {
+      bz.next_in =
+          const_cast<char*>(reinterpret_cast<const char*>(remaining_in.data()));
+      bz.avail_in = base::checked_cast<unsigned int>(
+          std::min(remaining_in.size(), GetMaxDecompressChunkSize()));
+      bz.next_out = reinterpret_cast<char*>(remaining_out.data());
+      bz.avail_out = base::checked_cast<unsigned int>(
+          std::min(remaining_out.size(), GetMaxDecompressChunkSize()));
+
+      unsigned int avail_in_before = bz.avail_in;
+      unsigned int avail_out_before = bz.avail_out;
+
+      rv = BZ2_bzDecompress(&bz);
+
+      size_t bytes_consumed = avail_in_before - bz.avail_in;
+      size_t bytes_produced = avail_out_before - bz.avail_out;
+
+      if (bytes_consumed == 0 && bytes_produced == 0 && rv == BZ_OK) {
+        break;
+      }
+
+      remaining_in = remaining_in.subspan(bytes_consumed);
+      remaining_out = remaining_out.subspan(bytes_produced);
+    }
+
     BZ2_bzDecompressEnd(&bz);
 
     if (rv != BZ_STREAM_END) {
diff --git a/chrome/utility/safe_browsing/mac/udif.h b/chrome/utility/safe_browsing/mac/udif.h
index ee4d7de..932151d3 100644
--- a/chrome/utility/safe_browsing/mac/udif.h
+++ b/chrome/utility/safe_browsing/mac/udif.h
@@ -90,6 +90,9 @@
   std::vector<uint8_t> signature_blob_;  // DMG signature.
 };
 
+// Returns a reference to the maximum chunk size used for decompression.
+size_t& GetMaxDecompressChunkSize();
+
 }  // namespace dmg
 }  // namespace safe_browsing
 
diff --git a/chrome/utility/safe_browsing/mac/udif_unittest.cc b/chrome/utility/safe_browsing/mac/udif_unittest.cc
index 5927f5f..3950d91 100644
--- a/chrome/utility/safe_browsing/mac/udif_unittest.cc
+++ b/chrome/utility/safe_browsing/mac/udif_unittest.cc
@@ -11,6 +11,7 @@
 
 #include <array>
 
+#include "base/auto_reset.h"
 #include "base/compiler_specific.h"
 #include "base/containers/span.h"
 #include "base/files/file.h"
@@ -197,6 +198,16 @@
   RunReadAllTest(100000);
 }
 
+TEST_P(UDIFParserTest, ReadAll_SmallDecompressChunk_64) {
+  base::AutoReset<size_t> auto_reset(&GetMaxDecompressChunkSize(), 64);
+  RunReadAllTest(512);
+}
+
+TEST_P(UDIFParserTest, ReadAll_SmallDecompressChunk_1024) {
+  base::AutoReset<size_t> auto_reset(&GetMaxDecompressChunkSize(), 1024);
+  RunReadAllTest(512);
+}
+
 constexpr UDIFTestCase cases[] = {

diff --git a/chrome/utility/safe_browsing/mac/udif_unittest.cc b/chrome/utility/safe_browsing/mac/udif_unittest.cc
index 5927f5f..3950d91 100644
--- a/chrome/utility/safe_browsing/mac/udif_unittest.cc
+++ b/chrome/utility/safe_browsing/mac/udif_unittest.cc
@@ -11,6 +11,7 @@
 
 #include <array>
 
+#include "base/auto_reset.h"
 #include "base/compiler_specific.h"
 #include "base/containers/span.h"
 #include "base/files/file.h"
@@ -197,6 +198,16 @@
   RunReadAllTest(100000);
 }
 
+TEST_P(UDIFParserTest, ReadAll_SmallDecompressChunk_64) {
+  base::AutoReset<size_t> auto_reset(&GetMaxDecompressChunkSize(), 64);
+  RunReadAllTest(512);
+}
+
+TEST_P(UDIFParserTest, ReadAll_SmallDecompressChunk_1024) {
+  base::AutoReset<size_t> auto_reset(&GetMaxDecompressChunkSize(), 1024);
+  RunReadAllTest(512);
+}
+
 constexpr UDIFTestCase cases[] = {
     {"dmg_UDBZ_GPTSPUD.dmg", kGPTExpectedPartitions, UDIFTestCase::ALL_PASS},
     {"dmg_UDBZ_NONE.dmg", kNoPartitionMap, UDIFTestCase::ALL_PASS},