diff --git a/Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py b/Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py index 5efc60a92111754..156662296ebfa95 100644 --- a/Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py +++ b/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.""" diff --git a/Misc/NEWS.d/next/Library/2026-07-01-12-00-00.gh-issue-152718.pr0f1l.rst b/Misc/NEWS.d/next/Library/2026-07-01-12-00-00.gh-issue-152718.pr0f1l.rst new file mode 100644 index 000000000000000..63d9d7f878cb07d --- /dev/null +++ b/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. diff --git a/Modules/_remote_debugging/binary_io.h b/Modules/_remote_debugging/binary_io.h index d4188335c0b6d0a..a4d6c52321e951c 100644 --- a/Modules/_remote_debugging/binary_io.h +++ b/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) diff --git a/Modules/_remote_debugging/binary_io_reader.c b/Modules/_remote_debugging/binary_io_reader.c index ce1c3d232c94e0f..73c4f5b6250d89b 100644 --- a/Modules/_remote_debugging/binary_io_reader.c +++ b/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) {