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

Author: tonghuaroot

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,85 @@ 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)
+
+        with self.assertRaises(ValueError) as cm:
+            with BinaryReader(filename):
+                pass
+        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,4 @@
+Fix unbounded memory allocation in the :mod:`profiling.sampling` binary profile
+reader, which trusted the string and frame table counts declared in a file's
+footer before allocating. A malformed file can no longer make the reader
+attempt a huge allocation.

Modules/_remote_debugging/binary_io.h

@@ -91,6 +91,11 @@ static_assert(SAMPLE_HEADER_FIXED_SIZE == 13,
 static_assert(FILE_FOOTER_SIZE == 32,
              "FILE_FOOTER_SIZE must remain 32");
 
+/* Minimum on-disk bytes per table entry: a string is a >=1-byte length
+ * varint; a frame is six >=1-byte varints plus the opcode byte. */
+#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

@@ -240,6 +240,16 @@ reader_decompress_samples(BinaryReader *reader, const uint8_t *data)
 static inline int
 reader_parse_string_table(BinaryReader *reader, const uint8_t *data, size_t file_size)
 {
+    /* Reject a count larger than the remaining bytes can hold. */
+    size_t max_strings =
+        (file_size - reader->string_table_offset) / MIN_STRING_ENTRY_SIZE;
+    if (reader->strings_count > max_strings) {
+        PyErr_Format(PyExc_ValueError,
+            "Invalid string count %u exceeds maximum possible %zu",
+            reader->strings_count, max_strings);
+        return -1;
+    }
+
     reader->strings = PyMem_Calloc(reader->strings_count, sizeof(PyObject *));
     if (!reader->strings && reader->strings_count > 0) {
         PyErr_NoMemory();
@@ -280,6 +290,16 @@ reader_parse_frame_table(BinaryReader *reader, const uint8_t *data, size_t file_
     }
 #endif
 
+    /* Reject a count larger than the remaining bytes can hold. */
+    size_t max_frames =
+        (file_size - reader->frame_table_offset) / MIN_FRAME_ENTRY_SIZE;
+    if (reader->frames_count > max_frames) {
+        PyErr_Format(PyExc_ValueError,
+            "Invalid frame count %u exceeds maximum possible %zu",
+            reader->frames_count, max_frames);
+        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) {