Skip to content

Commit 41a2abe

Browse files
[3.15] gh-152718: Reject oversized table counts in the profiling binary reader (GH-152719) (#153050)
gh-152718: Reject oversized table counts in the profiling binary reader (GH-152719) (cherry picked from commit 8e580b0) Co-authored-by: tonghuaroot (童话) <tonghuaroot@gmail.com>
1 parent 9f82c25 commit 41a2abe

4 files changed

Lines changed: 124 additions & 9 deletions

File tree

Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1003,6 +1003,9 @@ class TestBinaryFormatValidation(BinaryFormatTestBase):
10031003
HDR_OFF_STR_TABLE = 36
10041004
HDR_OFF_FRAME_TABLE = 44
10051005
FILE_HEADER_PLACEHOLDER_SIZE = 64
1006+
FILE_FOOTER_SIZE = 32
1007+
FTR_OFF_STRINGS = 0
1008+
FTR_OFF_FRAMES = 4
10061009

10071010
def test_replay_rejects_more_threads_than_declared(self):
10081011
"""Replay rejects files with more unique threads than the header declares."""
@@ -1064,6 +1067,88 @@ def test_replay_rejects_trailing_partial_sample_header(self):
10641067
reader.replay_samples(RawCollector())
10651068
self.assertEqual(str(cm.exception), "Truncated sample data: 1 trailing bytes")
10661069

1070+
# Minimum on-disk size of one table entry (see binary_io.h).
1071+
MIN_STRING_ENTRY_SIZE = 1
1072+
MIN_FRAME_ENTRY_SIZE = 7
1073+
1074+
def _read_offset(self, filename, hdr_off):
1075+
with open(filename, "rb") as raw:
1076+
raw.seek(hdr_off)
1077+
return struct.unpack("=Q", raw.read(8))[0]
1078+
1079+
def _patch_footer_count(self, filename, ftr_off, value):
1080+
size = os.path.getsize(filename)
1081+
with open(filename, "r+b") as raw:
1082+
raw.seek(size - self.FILE_FOOTER_SIZE + ftr_off)
1083+
raw.write(struct.pack("=I", value))
1084+
1085+
def test_open_rejects_string_count_larger_than_file(self):
1086+
"""Open rejects a footer string count larger than the file."""
1087+
samples = [[make_interpreter(0, [
1088+
make_thread(1, [make_frame("s.py", 10, "s")])
1089+
])]]
1090+
filename = self.create_binary_file(samples, compression="none")
1091+
size = os.path.getsize(filename)
1092+
str_off = self._read_offset(filename, self.HDR_OFF_STR_TABLE)
1093+
max_strings = (size - str_off) // self.MIN_STRING_ENTRY_SIZE
1094+
self._patch_footer_count(filename, self.FTR_OFF_STRINGS, 0xFFFFFFFF)
1095+
1096+
with self.assertRaises(ValueError) as cm:
1097+
with BinaryReader(filename):
1098+
pass
1099+
self.assertEqual(
1100+
str(cm.exception),
1101+
f"Invalid string count 4294967295 exceeds maximum "
1102+
f"possible {max_strings}",
1103+
)
1104+
1105+
def test_open_rejects_frame_count_larger_than_file(self):
1106+
"""Open rejects a footer frame count larger than the file."""
1107+
samples = [[make_interpreter(0, [
1108+
make_thread(1, [make_frame("f.py", 10, "f")])
1109+
])]]
1110+
filename = self.create_binary_file(samples, compression="none")
1111+
size = os.path.getsize(filename)
1112+
frame_off = self._read_offset(filename, self.HDR_OFF_FRAME_TABLE)
1113+
max_frames = (size - frame_off) // self.MIN_FRAME_ENTRY_SIZE
1114+
self._patch_footer_count(filename, self.FTR_OFF_FRAMES, 0xFFFFFFFF)
1115+
1116+
# On a 32-bit build this count overflows the allocation size, so the
1117+
# size_t overflow guard (OverflowError) fires before the count check.
1118+
with self.assertRaises((ValueError, OverflowError)) as cm:
1119+
with BinaryReader(filename):
1120+
pass
1121+
if isinstance(cm.exception, ValueError):
1122+
self.assertEqual(
1123+
str(cm.exception),
1124+
f"Invalid frame count 4294967295 exceeds maximum "
1125+
f"possible {max_frames}",
1126+
)
1127+
1128+
def test_open_accepts_frame_count_at_capacity_boundary(self):
1129+
"""A frame count at the file-size cap opens; one more is rejected."""
1130+
samples = [[make_interpreter(0, [
1131+
make_thread(1, [make_frame("f.py", 10, "f")])
1132+
])]]
1133+
filename = self.create_binary_file(samples, compression="none")
1134+
size = os.path.getsize(filename)
1135+
frame_off = self._read_offset(filename, self.HDR_OFF_FRAME_TABLE)
1136+
max_frames = (size - frame_off) // self.MIN_FRAME_ENTRY_SIZE
1137+
1138+
self._patch_footer_count(filename, self.FTR_OFF_FRAMES, max_frames)
1139+
with BinaryReader(filename):
1140+
pass
1141+
1142+
self._patch_footer_count(filename, self.FTR_OFF_FRAMES, max_frames + 1)
1143+
with self.assertRaises(ValueError) as cm:
1144+
with BinaryReader(filename):
1145+
pass
1146+
self.assertEqual(
1147+
str(cm.exception),
1148+
f"Invalid frame count {max_frames + 1} exceeds maximum "
1149+
f"possible {max_frames}",
1150+
)
1151+
10671152

10681153
class TestBinaryEncodings(BinaryFormatTestBase):
10691154
"""Tests specifically targeting different stack encodings."""
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix unbounded memory allocation in the :mod:`profiling.sampling` binary profile
2+
reader when a file declares more string or frame entries than it contains.

Modules/_remote_debugging/binary_io.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,10 @@ static_assert(SAMPLE_HEADER_FIXED_SIZE == 13,
9191
static_assert(FILE_FOOTER_SIZE == 32,
9292
"FILE_FOOTER_SIZE must remain 32");
9393

94+
/* Minimum on-disk bytes of a string (1) and frame (7) table entry. */
95+
#define MIN_STRING_ENTRY_SIZE 1
96+
#define MIN_FRAME_ENTRY_SIZE 7
97+
9498
/* Buffer sizes: 512KB balances syscall amortization against memory use,
9599
* and aligns well with filesystem block sizes and zstd dictionary windows */
96100
#define WRITE_BUFFER_SIZE (512 * 1024)

Modules/_remote_debugging/binary_io_reader.c

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -237,9 +237,32 @@ reader_decompress_samples(BinaryReader *reader, const uint8_t *data)
237237
}
238238
#endif
239239

240+
/* Reject a table/run count whose entries cannot fit in the bytes still
241+
* available; a malicious file could otherwise drive a huge allocation.
242+
* Each entry occupies at least min_entry_size bytes. */
243+
static int
244+
reader_validate_count(const char *what, uint32_t count,
245+
size_t available_bytes, size_t min_entry_size)
246+
{
247+
size_t max_possible = available_bytes / min_entry_size;
248+
if (count > max_possible) {
249+
PyErr_Format(PyExc_ValueError,
250+
"Invalid %s count %u exceeds maximum possible %zu",
251+
what, count, max_possible);
252+
return -1;
253+
}
254+
return 0;
255+
}
256+
240257
static inline int
241258
reader_parse_string_table(BinaryReader *reader, const uint8_t *data, size_t file_size)
242259
{
260+
if (reader_validate_count("string", reader->strings_count,
261+
file_size - reader->string_table_offset,
262+
MIN_STRING_ENTRY_SIZE) < 0) {
263+
return -1;
264+
}
265+
243266
reader->strings = PyMem_Calloc(reader->strings_count, sizeof(PyObject *));
244267
if (!reader->strings && reader->strings_count > 0) {
245268
PyErr_NoMemory();
@@ -280,6 +303,12 @@ reader_parse_frame_table(BinaryReader *reader, const uint8_t *data, size_t file_
280303
}
281304
#endif
282305

306+
if (reader_validate_count("frame", reader->frames_count,
307+
file_size - reader->frame_table_offset,
308+
MIN_FRAME_ENTRY_SIZE) < 0) {
309+
return -1;
310+
}
311+
283312
size_t alloc_size = (size_t)reader->frames_count * sizeof(FrameEntry);
284313
reader->frames = PyMem_Malloc(alloc_size);
285314
if (!reader->frames && reader->frames_count > 0) {
@@ -1053,15 +1082,10 @@ binary_reader_replay(BinaryReader *reader, PyObject *collector, PyObject *progre
10531082
return -1;
10541083
}
10551084

1056-
/* Validate RLE count to prevent DoS from malicious files.
1057-
* Each RLE sample needs at least 2 bytes (1 byte min varint + 1 status byte).
1058-
* Also reject absurdly large counts that would exhaust memory. */
1059-
size_t remaining_data = reader->sample_data_size - offset;
1060-
size_t max_possible_samples = remaining_data / 2;
1061-
if (count > max_possible_samples) {
1062-
PyErr_Format(PyExc_ValueError,
1063-
"Invalid RLE count %u exceeds maximum possible %zu for remaining data",
1064-
count, max_possible_samples);
1085+
/* Reject a count larger than the remaining bytes can hold; each
1086+
* RLE sample needs at least 2 bytes (1-byte min varint + status). */
1087+
if (reader_validate_count("RLE", count,
1088+
reader->sample_data_size - offset, 2) < 0) {
10651089
return -1;
10661090
}
10671091
if ((uint64_t)count > (uint64_t)PY_SSIZE_T_MAX - (uint64_t)replayed) {

0 commit comments

Comments
 (0)