gh-152718: Reject oversized table counts in the profiling binary reader

Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py

@@ -980,6 +980,9 @@ class TestBinaryFormatValidation(BinaryFormatTestBase):
     HDR_OFF_STR_TABLE = 36
     HDR_OFF_FRAME_TABLE = 44
     FILE_HEADER_PLACEHOLDER_SIZE = 64
+    FILE_FOOTER_SIZE = 32
+    FTR_OFF_STRINGS = 0
+    FTR_OFF_FRAMES = 4
 
     def test_replay_rejects_more_threads_than_declared(self):
         """Replay rejects files with more unique threads than the header declares."""
@@ -1041,6 +1044,88 @@ def test_replay_rejects_trailing_partial_sample_header(self):
                 reader.replay_samples(RawCollector())
             self.assertEqual(str(cm.exception), "Truncated sample data: 1 trailing bytes")
 
+    # Minimum on-disk size of one table entry (see binary_io.h).
+    MIN_STRING_ENTRY_SIZE = 1
+    MIN_FRAME_ENTRY_SIZE = 7
+
+    def _read_offset(self, filename, hdr_off):
+        with open(filename, "rb") as raw:
+            raw.seek(hdr_off)
+            return struct.unpack("=Q", raw.read(8))[0]
+
+    def _patch_footer_count(self, filename, ftr_off, value):
+        size = os.path.getsize(filename)
+        with open(filename, "r+b") as raw:
+            raw.seek(size - self.FILE_FOOTER_SIZE + ftr_off)
+            raw.write(struct.pack("=I", value))
+
+    def test_open_rejects_string_count_larger_than_file(self):
+        """Open rejects a footer string count larger than the file."""
+        samples = [[make_interpreter(0, [
+            make_thread(1, [make_frame("s.py", 10, "s")])
+        ])]]
+        filename = self.create_binary_file(samples, compression="none")
+        size = os.path.getsize(filename)
+        str_off = self._read_offset(filename, self.HDR_OFF_STR_TABLE)
+        max_strings = (size - str_off) // self.MIN_STRING_ENTRY_SIZE
+        self._patch_footer_count(filename, self.FTR_OFF_STRINGS, 0xFFFFFFFF)
+
+        with self.assertRaises(ValueError) as cm:
+            with BinaryReader(filename):
+                pass
+        self.assertEqual(
+            str(cm.exception),
+            f"Invalid string count 4294967295 exceeds maximum "
+            f"possible {max_strings}",
+        )
+
+    def test_open_rejects_frame_count_larger_than_file(self):
+        """Open rejects a footer frame count larger than the file."""
+        samples = [[make_interpreter(0, [
+            make_thread(1, [make_frame("f.py", 10, "f")])
+        ])]]
+        filename = self.create_binary_file(samples, compression="none")
+        size = os.path.getsize(filename)
+        frame_off = self._read_offset(filename, self.HDR_OFF_FRAME_TABLE)
+        max_frames = (size - frame_off) // self.MIN_FRAME_ENTRY_SIZE
+        self._patch_footer_count(filename, self.FTR_OFF_FRAMES, 0xFFFFFFFF)
+
+        # On a 32-bit build this count overflows the allocation size, so the
+        # size_t overflow guard (OverflowError) fires before the count check.
+        with self.assertRaises((ValueError, OverflowError)) as cm:
+            with BinaryReader(filename):
+                pass
+        if isinstance(cm.exception, ValueError):
+            self.assertEqual(
+                str(cm.exception),
+                f"Invalid frame count 4294967295 exceeds maximum "
+                f"possible {max_frames}",
+            )
+
+    def test_open_accepts_frame_count_at_capacity_boundary(self):
+        """A frame count at the file-size cap opens; one more is rejected."""
+        samples = [[make_interpreter(0, [
+            make_thread(1, [make_frame("f.py", 10, "f")])
+        ])]]
+        filename = self.create_binary_file(samples, compression="none")
+        size = os.path.getsize(filename)
+        frame_off = self._read_offset(filename, self.HDR_OFF_FRAME_TABLE)
+        max_frames = (size - frame_off) // self.MIN_FRAME_ENTRY_SIZE
+
+        self._patch_footer_count(filename, self.FTR_OFF_FRAMES, max_frames)
+        with BinaryReader(filename):
+            pass
+
+        self._patch_footer_count(filename, self.FTR_OFF_FRAMES, max_frames + 1)
+        with self.assertRaises(ValueError) as cm:
+            with BinaryReader(filename):
+                pass
+        self.assertEqual(
+            str(cm.exception),
+            f"Invalid frame count {max_frames + 1} exceeds maximum "
+            f"possible {max_frames}",
+        )
+
 
 class TestBinaryEncodings(BinaryFormatTestBase):
     """Tests specifically targeting different stack encodings."""

Misc/NEWS.d/next/Library/2026-07-01-12-00-00.gh-issue-152718.pr0f1l.rst

@@ -0,0 +1,2 @@
+Fix unbounded memory allocation in the :mod:`profiling.sampling` binary profile
+reader when a file declares more string or frame entries than it contains.

Modules/_remote_debugging/binary_io.h

@@ -91,6 +91,10 @@ static_assert(SAMPLE_HEADER_FIXED_SIZE == 13,
 static_assert(FILE_FOOTER_SIZE == 32,
              "FILE_FOOTER_SIZE must remain 32");
 
+/* Minimum on-disk bytes of a string (1) and frame (7) table entry. */
+#define MIN_STRING_ENTRY_SIZE 1
+#define MIN_FRAME_ENTRY_SIZE  7
+
 /* Buffer sizes: 512KB balances syscall amortization against memory use,
  * and aligns well with filesystem block sizes and zstd dictionary windows */
 #define WRITE_BUFFER_SIZE       (512 * 1024)

Modules/_remote_debugging/binary_io_reader.c

@@ -237,9 +237,32 @@ reader_decompress_samples(BinaryReader *reader, const uint8_t *data)
 }
 #endif
 
+/* Reject a table/run count whose entries cannot fit in the bytes still
+ * available; a malicious file could otherwise drive a huge allocation.
+ * Each entry occupies at least min_entry_size bytes. */
+static int
+reader_validate_count(const char *what, uint32_t count,
+                      size_t available_bytes, size_t min_entry_size)
+{
+    size_t max_possible = available_bytes / min_entry_size;
+    if (count > max_possible) {
+        PyErr_Format(PyExc_ValueError,
+            "Invalid %s count %u exceeds maximum possible %zu",
+            what, count, max_possible);
+        return -1;
+    }
+    return 0;
+}
+
 static inline int
 reader_parse_string_table(BinaryReader *reader, const uint8_t *data, size_t file_size)
 {
+    if (reader_validate_count("string", reader->strings_count,
+                              file_size - reader->string_table_offset,
+                              MIN_STRING_ENTRY_SIZE) < 0) {
+        return -1;
+    }
+
     reader->strings = PyMem_Calloc(reader->strings_count, sizeof(PyObject *));
     if (!reader->strings && reader->strings_count > 0) {
         PyErr_NoMemory();
@@ -280,6 +303,12 @@ reader_parse_frame_table(BinaryReader *reader, const uint8_t *data, size_t file_
     }
 #endif
 
+    if (reader_validate_count("frame", reader->frames_count,
+                              file_size - reader->frame_table_offset,
+                              MIN_FRAME_ENTRY_SIZE) < 0) {
+        return -1;
+    }
+
     size_t alloc_size = (size_t)reader->frames_count * sizeof(FrameEntry);
     reader->frames = PyMem_Malloc(alloc_size);
     if (!reader->frames && reader->frames_count > 0) {
@@ -1053,15 +1082,10 @@ binary_reader_replay(BinaryReader *reader, PyObject *collector, PyObject *progre
                 return -1;
             }
 
-            /* Validate RLE count to prevent DoS from malicious files.
-             * Each RLE sample needs at least 2 bytes (1 byte min varint + 1 status byte).
-             * Also reject absurdly large counts that would exhaust memory. */
-            size_t remaining_data = reader->sample_data_size - offset;
-            size_t max_possible_samples = remaining_data / 2;
-            if (count > max_possible_samples) {
-                PyErr_Format(PyExc_ValueError,
-                    "Invalid RLE count %u exceeds maximum possible %zu for remaining data",
-                    count, max_possible_samples);
+            /* Reject a count larger than the remaining bytes can hold; each
+             * RLE sample needs at least 2 bytes (1-byte min varint + status). */
+            if (reader_validate_count("RLE", count,
+                                      reader->sample_data_size - offset, 2) < 0) {
                 return -1;
             }
             if ((uint64_t)count > (uint64_t)PY_SSIZE_T_MAX - (uint64_t)replayed) {