Skip to content

Commit 2969b0e

Browse files
gh-154848: Add check_frames option to pickletools.dis() and genops()
Framing (PEP 3154) is ignored by default, as an unpickler is free to do. When check_frames is true, an argument that straddles a frame boundary, or a frame that begins before the previous one ends, raises ValueError, as in the standard unpickler. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent ae18dae commit 2969b0e

4 files changed

Lines changed: 111 additions & 7 deletions

File tree

Doc/library/pickletools.rst

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,8 @@ Programmatic interface
8888
----------------------
8989

9090

91-
.. function:: dis(pickle, out=None, memo=None, indentlevel=4, annotate=0)
91+
.. function:: dis(pickle, out=None, memo=None, indentlevel=4, annotate=0, *, \
92+
check_frames=False)
9293

9394
Outputs a symbolic disassembly of the pickle to the file-like
9495
object *out*, defaulting to ``sys.stdout``. *pickle* can be a
@@ -101,17 +102,30 @@ Programmatic interface
101102
a short description. The value of *annotate* is used as a hint for
102103
the column where annotation should start.
103104

105+
Framing (:pep:`3154`) is ignored by default, as an unpickler is free to do.
106+
If *check_frames* is true, an argument that straddles a frame boundary, or a
107+
frame that begins before the previous one ends, raises a :exc:`ValueError`,
108+
as in the standard unpickler.
109+
104110
.. versionchanged:: 3.2
105111
Added the *annotate* parameter.
106112

107-
.. function:: genops(pickle)
113+
.. versionchanged:: next
114+
Added the *check_frames* parameter.
115+
116+
.. function:: genops(pickle, *, check_frames=False)
108117

109118
Provides an :term:`iterator` over all of the opcodes in a pickle, returning a
110119
sequence of ``(opcode, arg, pos)`` triples. *opcode* is an instance of an
111120
:class:`OpcodeInfo` class; *arg* is the decoded value, as a Python object, of
112121
the opcode's argument; *pos* is the position at which this opcode is located.
113122
*pickle* can be a string or a file-like object.
114123

124+
The *check_frames* argument has the same meaning as in :func:`dis`.
125+
126+
.. versionchanged:: next
127+
Added the *check_frames* parameter.
128+
115129
.. function:: optimize(picklestring)
116130

117131
Returns a new equivalent pickle string after eliminating unused ``PUT``

Lib/pickletools.py

Lines changed: 52 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2293,7 +2293,38 @@ def assure_pickle_consistency(verbose=False):
22932293
##############################################################################
22942294
# A pickle opcode generator.
22952295

2296-
def _genops(data, yield_end_pos=False):
2296+
class _FramedReader:
2297+
"""Wrap a file object to enforce pickle frame boundaries (PEP 3154).
2298+
2299+
Reads are limited to the current frame, so an argument that would straddle
2300+
a frame boundary is reported as truncated data by the argument readers,
2301+
just like a genuinely truncated pickle.
2302+
"""
2303+
2304+
def __init__(self, file):
2305+
self.file = file
2306+
self.frame_len = None # bytes left in the current frame, or None
2307+
2308+
def read(self, n):
2309+
if self.frame_len is None:
2310+
return self.file.read(n)
2311+
data = self.file.read(min(n, self.frame_len))
2312+
self.frame_len -= len(data)
2313+
if self.frame_len == 0:
2314+
self.frame_len = None
2315+
return data
2316+
2317+
def readline(self):
2318+
if self.frame_len is None:
2319+
return self.file.readline()
2320+
data = self.file.readline(self.frame_len)
2321+
self.frame_len -= len(data)
2322+
if self.frame_len == 0:
2323+
self.frame_len = None
2324+
return data
2325+
2326+
2327+
def _genops(data, yield_end_pos=False, check_frames=False):
22972328
if isinstance(data, bytes_types):
22982329
data = io.BytesIO(data)
22992330

@@ -2317,6 +2348,13 @@ def _genops(data, yield_end_pos=False):
23172348
arg = None
23182349
else:
23192350
arg = opcode.arg.reader(data)
2351+
if check_frames and opcode.name == 'FRAME':
2352+
if not isinstance(data, _FramedReader):
2353+
data = _FramedReader(data)
2354+
elif data.frame_len is not None:
2355+
raise ValueError("beginning of a new frame before end of "
2356+
"current frame")
2357+
data.frame_len = arg or None
23202358
if yield_end_pos:
23212359
yield opcode, arg, pos, getpos()
23222360
else:
@@ -2325,7 +2363,7 @@ def _genops(data, yield_end_pos=False):
23252363
assert opcode.name == 'STOP'
23262364
break
23272365

2328-
def genops(pickle):
2366+
def genops(pickle, *, check_frames=False):
23292367
"""Generate all the opcodes in a pickle.
23302368
23312369
'pickle' is a file-like object, or string, containing the pickle.
@@ -2347,8 +2385,13 @@ def genops(pickle):
23472385
it's wrapped in a BytesIO object, and the latter's tell() result is
23482386
used. Else (the pickle doesn't have a tell(), and it's not obvious how
23492387
to query its current position) pos is None.
2388+
2389+
Framing (PEP 3154) is ignored by default, as an unpickler is free to do.
2390+
If 'check_frames' is true, an argument that straddles a frame boundary, or
2391+
a frame that begins before the previous one ends, raises a ValueError, as
2392+
it does in the standard unpickler.
23502393
"""
2351-
return _genops(pickle)
2394+
return _genops(pickle, check_frames=check_frames)
23522395

23532396
##############################################################################
23542397
# A pickle optimizer.
@@ -2420,7 +2463,8 @@ def optimize(p):
24202463
##############################################################################
24212464
# A symbolic pickle disassembler.
24222465

2423-
def dis(pickle, out=None, memo=None, indentlevel=4, annotate=0):
2466+
def dis(pickle, out=None, memo=None, indentlevel=4, annotate=0, *,
2467+
check_frames=False):
24242468
"""Produce a symbolic disassembly of a pickle.
24252469
24262470
'pickle' is a file-like object, or string, containing a (at least one)
@@ -2457,6 +2501,9 @@ def dis(pickle, out=None, memo=None, indentlevel=4, annotate=0):
24572501
+ A memo entry isn't referenced before it's defined.
24582502
24592503
+ The markobject isn't stored in the memo.
2504+
2505+
Framing (PEP 3154) is ignored by default. If 'check_frames' is true,
2506+
frame boundaries are enforced as in the standard unpickler; see genops().
24602507
"""
24612508

24622509
# Most of the hair here is for sanity checks, but most of it is needed
@@ -2472,7 +2519,7 @@ def dis(pickle, out=None, memo=None, indentlevel=4, annotate=0):
24722519
errormsg = None
24732520
annocol = annotate # column hint for annotations
24742521
t = get_theme(tty_file=out).pickletools
2475-
for opcode, arg, pos in genops(pickle):
2522+
for opcode, arg, pos in genops(pickle, check_frames=check_frames):
24762523
if pos is not None:
24772524
print(f"{t.position}{pos:5d}:{t.reset}", end=' ', file=out)
24782525

Lib/test/test_pickletools.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,44 @@ def test_truncated_data(self):
142142
'not enough data in stream to read int4'):
143143
next(it)
144144

145+
def test_check_frames(self):
146+
# Valid framed pickles are disassembled the same way regardless of
147+
# whether frame boundaries are checked.
148+
valid = pickle.dumps([1, 2, 3], 4)
149+
for check in (False, True):
150+
with self.subTest(check_frames=check):
151+
ops = [op.name for op, arg, pos in
152+
pickletools.genops(valid, check_frames=check)]
153+
self.assertIn('FRAME', ops)
154+
self.assertEqual(ops[-1], 'STOP')
155+
156+
def test_check_frames_straddle(self):
157+
# An opcode argument that straddles a frame boundary is ignored by
158+
# default, but rejected when check_frames is true. See gh-154848.
159+
# FRAME 3; SHORT_BINBYTES argument straddles the frame.
160+
data = b'\x80\x05\x95\x03\x00\x00\x00\x00\x00\x00\x00C\x0ahelloworld.'
161+
self.assertEqual(list(pickletools.genops(data))[-1][0].name, 'STOP')
162+
with self.assertRaisesRegex(ValueError,
163+
'expected 10 bytes in a bytes1, but only 1 remain'):
164+
list(pickletools.genops(data, check_frames=True))
165+
166+
# FRAME 6; UNICODE argument (read by readline) straddles the frame.
167+
data = b'\x80\x05\x95\x06\x00\x00\x00\x00\x00\x00\x00Vhelloworld\n.'
168+
self.assertEqual(list(pickletools.genops(data))[-1][0].name, 'STOP')
169+
with self.assertRaisesRegex(ValueError,
170+
'no newline found when trying to read unicodestringnl'):
171+
list(pickletools.genops(data, check_frames=True))
172+
173+
def test_check_frames_nested(self):
174+
# A new frame beginning before the current one ends is rejected only
175+
# when check_frames is true.
176+
data = (b'\x80\x05\x95\x0c\x00\x00\x00\x00\x00\x00\x00'
177+
b'N\x95\x00\x00\x00\x00\x00\x00\x00\x00NN.')
178+
self.assertEqual(list(pickletools.genops(data))[-1][0].name, 'STOP')
179+
with self.assertRaisesRegex(ValueError,
180+
'beginning of a new frame before end of current frame'):
181+
list(pickletools.genops(data, check_frames=True))
182+
145183
def test_unknown_opcode(self):
146184
it = pickletools.genops(b'N\xff')
147185
item = next(it)
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
:func:`pickletools.dis` and :func:`pickletools.genops` now accept a
2+
*check_frames* keyword argument. When true, an argument that straddles a
3+
frame boundary, or a frame that begins before the previous one ends, raises
4+
:exc:`ValueError` instead of being read across the boundary (PEP 3154).
5+
Framing is still ignored by default, as an unpickler is free to do.

0 commit comments

Comments
 (0)